Expose packed: False, set log_peak_memory_stats: True, set compile: False #1872

krammnic · 2024-10-20T18:32:25Z

Context

What is the purpose of this PR? Is it to

add a new feature
fix a bug
update tests and/or documentation
other (Clean up)

Please link to any issues this PR addresses.

Changelog

What are the changes made in this PR?

update config defaults dataset.packed and log_peak_memory_stats #1732

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
add unit tests for any new functionality
update docstrings for any new or updated methods or classes
run unit tests via pytest tests
run recipe tests via pytest tests -m integration_test
manually run any new or modified recipes with sufficient proof of correctness
include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

I did not change any public API
I have added an example to docs or docstrings

pytorch-bot · 2024-10-20T18:32:28Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1872

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

joecummings

two nits: overall, big fan of this UX improvement

joecummings · 2024-10-21T13:28:01Z

recipes/configs/code_llama2/7B_full_low_memory.yaml

@@ -45,7 +45,9 @@ resume_from_checkpoint: False

 # Dataset
 dataset:
+  packed: False # Set to true for great speed ups


Huge nit: can we move this below the _component_ declaration?

That way it reads more as an option for the specific builder.

Yeah, though like this first time.. But I followed original declaration from issue:

dataset: packed=False # Set to true for great speed ups

Will be fixed

joecummings · 2024-10-21T13:29:05Z

recipes/configs/dev/8B_full_experimental.yaml

@@ -57,7 +58,7 @@ loss:
  _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
 max_steps_per_epoch: null
 gradient_accumulation_steps: 1
-
+compile: False


Should we be explicit that this will torch.compile the model and loss? compile seems like a vague name

agreed. I think we can add a comment, similar to packed

compile=False # pytorch compile, set to true for perf/memory improvement

wdyt?

Add some comment there?

Ok, will add

felipemello1 · 2024-10-21T15:11:33Z

@SalmanMohammadi , do we support compile and packed in PPO? If not, maybe we should not add those to these configs.

recipes/configs/mistral/7B_full_ppo_low_memory.yaml

SalmanMohammadi · 2024-10-21T15:15:24Z

@SalmanMohammadi , do we support compile and packed in PPO? If not, maybe we should not add those to these configs.

Compile... yes, kind of? It's going to be overhauled soon to properly support, it's very out-of-date, it doesn't even have hierarchical compilation.

felipemello1 · 2024-10-21T15:20:53Z

@SalmanMohammadi , do we support compile and packed in PPO? If not, maybe we should not add those to these configs.

Compile... yes, kind of? It's going to be overhauled soon to properly support, it's very out-of-date, it doesn't even have hierarchical compilation.

ok, so compile is NOT a no-op, and we SHOULD have it in the configs. However, packed does NOT work with PPO, and should NOT be added to ppo configs. Is that right?

SalmanMohammadi · 2024-10-21T15:21:59Z

@SalmanMohammadi , do we support compile and packed in PPO? If not, maybe we should not add those to these configs.

Compile... yes, kind of? It's going to be overhauled soon to properly support, it's very out-of-date, it doesn't even have hierarchical compilation.

ok, so compile is NOT a no-op, and we SHOULD have it in the configs. However, packed does NOT work with PPO, and should NOT be added to ppo configs. Is that right?

CORRECT

felipemello1 · 2024-10-21T15:22:00Z

@krammnic , we also have to check if compile/packed work for the knowledge distillation recipe. If not, we need to remove it from the configs. In short, i know for sure that these work for LORA and Full finetuning recipes/configs.

SalmanMohammadi · 2024-10-21T15:23:20Z

@krammnic packed should also be removed from all the DPO configs, please.

krammnic · 2024-10-21T15:39:03Z

@krammnic , we also have to check if compile/packed work for the knowledge distillation recipe. If not, we need to remove it from the configs. In short, i know for sure that these work for LORA and Full finetuning recipes/configs.

Sure, will test then

krammnic · 2024-10-21T15:44:34Z

@krammnic packed should also be removed from all the DPO configs, please.

Done

krammnic · 2024-10-21T16:00:27Z

Added, comment to compile: False

felipemello1 · 2024-10-21T20:34:17Z

We will need to update recipes to log memory. We are getting the error

FAILED tests/recipes/test_lora_dpo_single_device.py::TestLoRADPOSingleDeviceRecipe::test_save_and_load_merged_weights - ValueError: Logging memory stats is only supported on CUDA devices, got cpu

So where log_peak_memory_stats, we need to add "if device = 'cuda'" and add info "log.info("log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."

cc: @ebsmothers

krammnic · 2024-10-21T21:13:09Z

We will need to update recipes to log memory. We are getting the error
FAILED tests/recipes/test_lora_dpo_single_device.py::TestLoRADPOSingleDeviceRecipe::test_save_and_load_merged_weights - ValueError: Logging memory stats is only supported on CUDA devices, got cpu
So where log_peak_memory_stats, we need to add "if device = 'cuda'" and add info "log.info("log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."

cc: @ebsmothers

Sure, will be done!

felipemello1 · 2024-10-22T14:01:57Z

We will need to update recipes to log memory. We are getting the error
FAILED tests/recipes/test_lora_dpo_single_device.py::TestLoRADPOSingleDeviceRecipe::test_save_and_load_merged_weights - ValueError: Logging memory stats is only supported on CUDA devices, got cpu
So where log_peak_memory_stats, we need to add "if device = 'cuda'" and add info "log.info("log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."

cc: @ebsmothers

edit: lets actually do this check in the init of the recipe. In the future, we can move all of these checks to some function like "config_parse". We already have multiple of these checks in the init

SalmanMohammadi · 2024-10-22T14:31:08Z

We will need to update recipes to log memory. We are getting the error
FAILED tests/recipes/test_lora_dpo_single_device.py::TestLoRADPOSingleDeviceRecipe::test_save_and_load_merged_weights - ValueError: Logging memory stats is only supported on CUDA devices, got cpu
So where log_peak_memory_stats, we need to add "if device = 'cuda'" and add info "log.info("log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."
cc: @ebsmothers
edit: lets actually do this check in the init of the recipe. In the future, we can move all of these checks to some function like "config_parse". We already have multiple of these checks in the init

I believe we're already doing this in most recipes when the stats are logged - the DPO recipe hasn't been updated.

SalmanMohammadi · 2024-10-22T14:55:07Z

We will need to update recipes to log memory. We are getting the error
FAILED tests/recipes/test_lora_dpo_single_device.py::TestLoRADPOSingleDeviceRecipe::test_save_and_load_merged_weights - ValueError: Logging memory stats is only supported on CUDA devices, got cpu
So where log_peak_memory_stats, we need to add "if device = 'cuda'" and add info "log.info("log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."
cc: @ebsmothers
edit: lets actually do this check in the init of the recipe. In the future, we can move all of these checks to some function like "config_parse". We already have multiple of these checks in the init
I believe we're already doing this in most recipes when the stats are logged - the DPO recipe hasn't been updated.

The DPO recipes uses:

                        if self._log_peak_memory_stats:
                            log_dict.update(
                                training.get_memory_stats(device=self._device)
                            )

it should be

                        if self._device.type == "cuda" and self._log_peak_memory_stats:
                            log_dict.update(
                                training.get_memory_stats(device=self._device)
                            )

krammnic · 2024-10-23T13:20:10Z

Add required check:

if self._log_peak_memory_stats and self._device.type == "cuda":
     log.info(
                "log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."
     )
     self._log_peak_memory_stats = False

SalmanMohammadi · 2024-10-23T14:05:45Z

recipes/qat_distributed.py

@@ -127,6 +127,12 @@ def __init__(self, cfg: DictConfig) -> None:
        self._log_every_n_steps = cfg.get("log_every_n_steps", 1)
        self._log_peak_memory_stats = cfg.get("log_peak_memory_stats", False)

+        if self._log_peak_memory_stats and self._device.type == "cuda":


Thoughts on whether we actually need this? I realise we kind of fail "silently" at the moment by just not logging if we aren't running on CUDA

torchtune/recipes/lora_finetune_single_device.py

Lines 716 to 719 in 17ba37d

if (

self._device.type == "cuda"

and self._log_peak_memory_stats

):

)

@joecummings @felipemello1

As-is we're now doing duplicating this check - once in in the init, and also every time we log the memory stats (in model setup, and during training) which isn't super clean. Personally I'd rather just make the check in the relevant logging util - but don't have to block on this.

let me get back to you on this @krammnic . Thanks for making all of these changes! :)

I'm not really see the problem in 2 partially duplicating checks. The point is, that if we have such logic in train:

log_dict = { "loss": loss_to_log, "lr": self._optimizer.param_groups[0]["lr"], "tokens_per_second_per_gpu": num_tokens / time_per_step, } if self._log_peak_memory_stats: log_dict.update( training.get_memory_stats(device=self._device) ) self._metric_logger.log_dict( log_dict, step=self.global_step, )

We can't do anything better, can we? Check in __init__ is about cuda and logging(once). Check in train probably should not be about "cuda"(there is no use case) and not about logging. I'm not sure if this should be in _metric_logger either

krammnic · 2024-10-24T13:38:51Z

Fixed some nits. Probably should be fine

felipemello1 · 2024-10-24T13:48:00Z

recipes/configs/llama3_1/8B_full.yaml


 # Training env
 device: cuda

 # Memory management
 enable_activation_checkpointing: True
 custom_sharded_layers: ['tok_embeddings', 'output']
-compile: False # set it to True for better memory and performance
+compile=False  # pytorch compile, set to true for perf/memory improvement# set it to True for better memory and performance


felipemello1 · 2024-10-24T13:48:50Z

recipes/configs/llama3_1/8B_full_single_device.yaml

@@ -61,7 +62,7 @@ loss:
 max_steps_per_epoch: null
 gradient_accumulation_steps: 1
 optimizer_in_bwd: True
-compile: False # set it to True for better memory and performance
+compile=False  # pytorch compile, set to true for perf/memory improvement# set it to True for better memory and performance


do you mind quickly double checking? Also, if you are using a script, maybe for a sanity check make just that compile/packed dont appear twice?

this should be a colon? compile: False

Yes, obviously. Fixed

felipemello1 · 2024-10-24T13:50:54Z

recipes/full_finetune_distributed.py

@@ -121,6 +121,12 @@ def __init__(self, cfg: DictConfig) -> None:
        self._log_every_n_steps = cfg.get("log_every_n_steps", 1)
        self._log_peak_memory_stats = cfg.get("log_peak_memory_stats", False)

+        if self._log_peak_memory_stats and self._device.type == "cuda":
+            log.info(
+                "log_peak_memory_stats was se to True, however, training does not use cuda. Setting log_peak_memory_stats=False."


i think that there are typos. Thats my fault, i guess you just copied/pasted what i wrote earlier.

Yes))) I didn't double check because it pretty minor

Fixed typos

also, it should be self._device.type != "cuda" not "=="

@krammnic sorry, just realized that the condition is wrong

Good catch! Yeah need to be fixed

felipemello1 · 2024-10-24T15:34:22Z

recipes/configs/llama3_1/405B_qlora.yaml

@@ -71,7 +72,7 @@ fsdp:
 epochs: 1
 max_steps_per_epoch: null
 gradient_accumulation_steps: 16
-compile: False # set it to True for better memory and performance
+compile=False  # pytorch compile, set to true for perf/memory improvement# set it to True for better memory and performance


other configs still need fixing :(

krammnic added 2 commits October 20, 2024 14:23

add required arguments in configs

3978d15

fix incorrect replacement

00d09d0

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 20, 2024

joecummings reviewed Oct 21, 2024

View reviewed changes

SalmanMohammadi reviewed Oct 21, 2024

View reviewed changes

recipes/configs/mistral/7B_full_ppo_low_memory.yaml Outdated Show resolved Hide resolved

krammnic added 2 commits October 21, 2024 11:42

remove packed from dpo

bb7648f

remove packed from dpo

22c84ab

krammnic force-pushed the configs branch from a7bf9bc to 22c84ab Compare October 21, 2024 15:44

add comment

8acfa75

add cuda check

9b47501

fix lint

e60eebd

SalmanMohammadi reviewed Oct 23, 2024

View reviewed changes

fixes

dacc3d2

felipemello1 reviewed Oct 24, 2024

View reviewed changes

krammnic added 8 commits October 24, 2024 17:29

Update 8B_full_single_device.yaml

b756391

Update full_finetune_distributed.py

06df423

Update full_finetune_single_device.py

6e8328f

Update knowledge_distillation_single_device.py

89eaa67

Update lora_dpo_distributed.py

641d081

Update lora_dpo_single_device.py

42b51fd

Update lora_finetune_single_device.py

442b26b

Update qat_distributed.py

28a9406

felipemello1 reviewed Oct 24, 2024

View reviewed changes

krammnic added 5 commits October 25, 2024 23:18

Update 405B_qlora.yaml

77bf771

Update 70B_full.yaml

b4a3e25

Update 1B_qlora_single_device.yaml

41b0012

Update 3B_qlora_single_device.yaml

c0e9645

Update 3B_lora_single_device.yaml

763b23d

	if (
	self._device.type == "cuda"
	and self._log_peak_memory_stats
	):

Expose packed: False, set log_peak_memory_stats: True, set compile: False #1872

Are you sure you want to change the base?

Expose packed: False, set log_peak_memory_stats: True, set compile: False #1872

Conversation

krammnic commented Oct 20, 2024

Context

Changelog

Test plan

UX

pytorch-bot bot commented Oct 20, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1872

joecummings left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 commented Oct 21, 2024

SalmanMohammadi commented Oct 21, 2024

felipemello1 commented Oct 21, 2024

SalmanMohammadi commented Oct 21, 2024

felipemello1 commented Oct 21, 2024

SalmanMohammadi commented Oct 21, 2024

krammnic commented Oct 21, 2024 • edited Loading

krammnic commented Oct 21, 2024

krammnic commented Oct 21, 2024

felipemello1 commented Oct 21, 2024

krammnic commented Oct 21, 2024

felipemello1 commented Oct 22, 2024

SalmanMohammadi commented Oct 22, 2024

SalmanMohammadi commented Oct 22, 2024

krammnic commented Oct 23, 2024

SalmanMohammadi Oct 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

krammnic commented Oct 24, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 Oct 24, 2024 • edited Loading

Choose a reason for hiding this comment

pytorch-bot bot commented Oct 20, 2024 •

edited

Loading

felipemello1 Oct 21, 2024 •

edited

Loading

krammnic commented Oct 21, 2024 •

edited

Loading

SalmanMohammadi Oct 23, 2024 •

edited

Loading

felipemello1 Oct 24, 2024 •

edited

Loading