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.

✅ No Failures

As of commit b0b4b14 with merge base 3ca0d30 ():
💚 Looks good so far! There are no failures yet. 💚

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 · 2024-10-26T11:40:49Z

Fixed all typos

felipemello1

lgtm! thanks for the pr!

felipemello1 · 2024-10-25T22:46:52Z

recipes/configs/llama3/70B_full.yaml

@@ -99,7 +100,7 @@ device: cuda
 enable_activation_checkpointing: True
 custom_sharded_layers: ['tok_embeddings', 'output']
 fsdp_cpu_offload: 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


felipemello1 · 2024-10-25T22:46:59Z

recipes/configs/llama3/70B_lora.yaml

@@ -89,15 +90,15 @@ loss:
 epochs: 1
 max_steps_per_epoch: null
 gradient_accumulation_steps: 1
-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-25T22:48:04Z

recipes/configs/llama3_1/70B_lora.yaml

@@ -88,15 +89,15 @@ loss:
 epochs: 1
 max_steps_per_epoch: null
 gradient_accumulation_steps: 1
-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


i think that the easiest way is to ctrol+f and search for "compile=". There are >15 of such cases.

felipemello1 · 2024-10-25T22:48:55Z

recipes/ppo_full_finetune_single_device.py

@@ -119,6 +119,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":
