OOM handling and recovery

[vadimkantorov]

We just hit OOM, revealing that by default torchtune does not use torch.compile and that it does not use fused linear cross entropy yet...

I found the following report from 2024:

• https://www.reddit.com/r/LocalLLaMA/comments/1di0fhv/torchtune_vs_axolotl_vs_unsloth_trainer/
• https://wandb.ai/augmxnt/train-bench/reports/Trainer-performance-comparison-torchtune-vs-axolotl-vs-Unsloth---Vmlldzo4MzU3NTAx

Are there any plans to make torchtune excellent for peak GPU memory usage and practical OOM handling?

For fine-tuning on long chain-of-thought's, OOM coming from some long not-filtered-out examples is an issue.

Could torchtune somehow make it OOM-crash-safe out-of-the-box? E.g. robustly skip batches in runtime if OOM occurs once in a while? Or aggressively cpu-offload some huge activation tensors, when OOM hits and re-try after off-loading...

File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/recipes/full_finetune_distributed.py", line 919, in train                
[rank6]:     current_loss = self._loss_step(batch) * current_num_tokens                                                                       
[rank6]:                    ^^^^^^^^^^^^^^^^^^^^^^                                                                                            
[rank6]:   File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/recipes/full_finetune_distributed.py", line 822, in _loss_step           
[rank6]:     loss = self._loss_fn(outputs, labels)                                                                                            
[rank6]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                            
[rank6]:   File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl            
[rank6]:     return self._call_impl(*args, **kwargs)                                                                                          
[rank6]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                          
[rank6]:   File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl                    
[rank6]:     return forward_call(*args, **kwargs)                                                                                             
[rank6]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                             
[rank6]:   File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/torchtune/modules/loss/cross_entropy_loss.py", line 136, in forward      
[rank6]:     total_loss += self.compute_cross_entropy(                                                                                        
[rank6]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                        
[rank6]:   File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/torchtune/modules/loss/cross_entropy_loss.py", line 105, in compute_cross
_entropy                                                                                                                                      
[rank6]:     return F.cross_entropy(                                                                                                          
[rank6]:            ^^^^^^^^^^^^^^^^                                                                                                          
[rank6]:   File "/mnt/fs/venv_torchtune/lib/python3.12/site-packages/torch/nn/functional.py", line 3494, in cross_entropy                     
[rank6]:     return torch._C._nn.cross_entropy_loss(                                                                                          
[rank6]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                          
[rank6]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 5.15 GiB. GPU 6 has a total capacity of 79.10 GiB of which 3.79 GiB is 
free. Including non-PyTorch memory, this process has 75.30 GiB memory in use. Of the allocated memory 58.28 GiB is allocated by PyTorch, and 1
5.06 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_se
gments:True to avoid fragmentation.  See documentation for Memory Management

[felipemello1]

hey @vadimkantorov , you can see all the knobs we have here to enable low memory: https://github.com/pytorch/torchtune#optimization-flags

You can also enable context parallelism for very long sequence length: #2668

In short, you can get to pretty long context if you enable our flags + context parallel. Which model and context length do you have in mind?

Could torchtune somehow make it OOM-crash-safe out-of-the-box? E.g. robustly skip batches in runtime if OOM occurs once in a while?

It should error in the very first iteration if you have dataset.packed=True, because it forces all sequences to have max_seq_len by packing them together. We are moving more towards packed=True being the default, so at the time we are not prioritizing skipping batch if OOM. Let me know if that makes sense.

---

extra:
There is a PR to enable liger loss. It is close to merging, but you wont see benefits vs our chunked loss unless you use bwd_in_opt flag: #2741

We also have a PR to do distributed parallel loss, but it doesnt work with compile yet: #2782

[vadimkantorov]

Let me know if that makes sense.

Mostly, it does make sense. But likely due to memory fragmentation, even with packed / bucketed dataset, OOMs can sometimes occur. For actual robust usage, it's important to have some (even if crude) mechanisms for occasional OOM recovery/skipping happening on some batches in some ranks.

[felipemello1]

I see. I dont think that skipping batches would solve memory fragmentation. Maybe some sort of collect?

But please, give this a flag a try. It should help. We will make it default in our recipes:

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

[vadimkantorov]

I dont think that skipping batches would solve memory fragmentation.

It would not solve it, but at least OOM handling via skipping batches (or another way could be some on-OOM CPU off-loading and op retry) prevent the fine-tuning run from failing if OOM occurs only once in a while - as a last resort to prevent unfortunate crashes.

Fragmentation is also an issue if the batches/packed batches are not padded to multiples of some size...

[felipemello1]

Sorry, what i am trying to say is that with packed=True, the batchsize is constant. In this scenario, skipping batch X to get batch Y of the same size wouldnt solve it if the issue is some sort of fragmentation and reserved memory is growing over time.

I do agree that this would be helpful for batches of variable size.

[vadimkantorov]

batchsize is constant

Are the tensors padded on the right for making allocations completely constant and never falling into some other allocator bin? Because even if inputs are padded on the right, kernels may like to do intermediate smaller allocations to profit from discarded padding...

Because otherwise I don't see how processing varlen prompts and completions will result in allocations of perfectly same size and always making perfectly same allocations...

[felipemello1]

yes, its packed on the right. What we see in practice is constant peak memory allocated and peak memory reserved.

[vadimkantorov]

I also found some special package https://github.com/pytorch/torchft, but don't know if they provide any recipes to guard/recover against OOMs...

[vadimkantorov]

the fsdp + OOM recovery problem is relevant not only for torchtune, e.g. in GRPO context, where one does not know the completion seqlens a-priori and might not be able to easily do packing:

• OOM handling with multi-node training on long-response reasoning tasks volcengine/verl#2159

So, I would say that a generic solution for OOM recovery under FSDP 1/2 (and especially multi-node, so maybe HSDP) is needed...