CUDA OOM when running QAT for gpt-oss-20b

[eduardzl]

I am trying to run QAT on gpt-oss-20b model after FFT on instance with 8 A100 GPUs (8x80GB).
I am getting CUDA OOM error for both configurations : sft_full.yaml or sft_lora.yaml, if the max_length is > 4096, in my case 32k.

The command is (flash_attn is installed) :

accelerate launch --config_file configs/zero3.yaml sft.py  --config configs/sft_lora.yaml \
        --quant_cfg MXFP4_MLP_WEIGHT_ONLY_CFG --attn_implementation kernels-community/vllm-flash-attn3

Doesn't matter what, after couple of training steps, the training OOM.

12%|███████████████                                                                                                         | 1/8 [01:38<11:28, 98.31s/it][Rank 0] Memory usage at training step 1: memory (MB) | allocated:  6.44e+01 | max_allocated:  8.34e+03 | reserved:  5.54e+04 | max_reserved:  5.54e+04
{'loss': 7.1055, 'grad_norm': 7.644806500757786, 'learning_rate': 0.0, 'entropy': 2.1449486911296844, 'num_tokens': 138061.0, 'mean_token_accuracy': 0.1823525046929717, 'epoch': 0.13}
{'loss': 7.1342, 'grad_norm': 8.164548440793329, 'learning_rate': 0.0002, 'entropy': 2.1739209294319153, 'num_tokens': 278805.0, 'mean_token_accuracy': 0.1771835247054696, 'epoch': 0.26}
{'loss': 6.7401, 'grad_norm': 6.377030560710944, 'learning_rate': 0.00019108719811121772, 'entropy': 2.2128825187683105, 'num_tokens': 416105.0, 'mean_token_accuracy': 0.19561117608100176, 'epoch': 0.38}
{'loss': 6.1746, 'grad_norm': 5.935622939274496, 'learning_rate': 0.00016611408216728603, 'entropy': 2.240649074316025, 'num_tokens': 554550.0, 'mean_token_accuracy': 0.21982254460453987, 'epoch': 0.51}
 50%|████████████████████████████████████████████████████████████                                                            | 4/8 [06:28<06:27, 96.90s/it][rank1]: Traceback (most recent call last):
[rank1]:   File "/workspace/TensorRT-Model-Optimizer/examples/gpt-oss/sft.py", line 123, in <module>
[rank1]:     main(script_args, training_args, model_args, quant_args)
[rank1]:   File "/workspace/TensorRT-Model-Optimizer/examples/gpt-oss/sft.py", line 112, in main
[rank1]:     trainer.train()
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/modelopt/torch/quantization/plugins/transformers_trainer.py", line 254, in train
[rank1]:     outputs = super().train(*args, **kwargs)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/trainer.py", line 2328, in train
[rank1]:     return inner_training_loop(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/trainer.py", line 2672, in _inner_training_loop
[rank1]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank1]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/modelopt/torch/quantization/plugins/transformers_trainer.py", line 235, in training_step
[rank1]:     return super().training_step(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 1189, in training_step
[rank1]:     return super().training_step(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/trainer.py", line 4060, in training_step
[rank1]:     self.accelerator.backward(loss, **kwargs)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/accelerate/accelerator.py", line 2726, in backward
[rank1]:     self.deepspeed_engine_wrapped.backward(loss, sync_gradients=self.sync_gradients, **kwargs)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/accelerate/utils/deepspeed.py", line 270, in backward
[rank1]:     self.engine.backward(loss, **kwargs)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]:     ret_val = func(*args, **kwargs)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2298, in backward
[rank1]:     self._do_optimizer_backward(loss, retain_graph)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2244, in _do_optimizer_backward
[rank1]:     self.optimizer.backward(loss, retain_graph=retain_graph)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]:     ret_val = func(*args, **kwargs)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/deepspeed/runtime/zero/stage3.py", line 2305, in backward
[rank1]:     self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 65, in backward
[rank1]:     scaled_loss.backward(retain_graph=retain_graph)
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/torch/_tensor.py", line 647, in backward
[rank1]:     torch.autograd.backward(
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/torch/autograd/__init__.py", line 354, in backward
[rank1]:     _engine_run_backward(
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/torch/autograd/graph.py", line 829, in _engine_run_backward
[rank1]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 24.55 GiB. GPU 1 has a total capacity of 79.25 GiB of which 18.90 GiB is free. Including non-PyTorch memory, this process has 60.33 GiB memory in use. Of the allocated memory 56.34 GiB is allocated by PyTorch, and 2.42 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
W0917 10:24:13.135000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17457 closing signal SIGTERM
W0917 10:24:13.138000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17459 closing signal SIGTERM
W0917 10:24:13.140000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17460 closing signal SIGTERM
W0917 10:24:13.156000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17461 closing signal SIGTERM
W0917 10:24:13.159000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17462 closing signal SIGTERM
W0917 10:24:13.165000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17463 closing signal SIGTERM
W0917 10:24:13.167000 17232 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 17464 closing signal SIGTERM

I was able to run FFT on the same instance for 32k context length with Axolotl. No CUDA OOMs..
When trying QAT, I am not even able to run LoRA training.

Env:
Latest code from "main" was used.
torch==2.8.0
CUDA 12.6
transformers==4.56.1

Any help would be greatly appreciated.

[realAsma]

@eduardzl
Have you tracked the GPU memory consumption over time with tensorboard or trackio? If you have them, could you please share it?

This looks like a memory leak to me.

In the past, I had seen such a memory leak with FSDP backend. We solved that by calling gc.collect() during training. Please see a utitility function, monkey_patch_training_step_to_fix_memory_leak which does this here -

TensorRT-Model-Optimizer/examples/llm_qat/main.py

Line 259 in 682bf6d

monkey_patch_training_step_to_fix_memory_leak(trainer)

Could you please call monkey_patch_training_step_to_fix_memory_leak before trainer.train() and see if that works?

[eduardzl]

@realAsma
Thanks for the reply. I have applied the suggested fix, run with fsdp and deepspeed, same result : CUDA OOM.
All GPUs allocated memory goes from 12GB on first step, to 70GB and then OOM.

rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1879, in _call_impl
[rank1]:     return inner()
[rank1]:            ^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1827, in inner
[rank1]:     result = forward_call(*args, **kwargs)
[rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/peft/tuners/tuners_utils.py", line 222, in forward
[rank1]:     return self.model.forward(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/utils/generic.py", line 940, in wrapper
[rank1]:     output = func(self, *args, **kwargs)
[rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 682, in forward
[rank1]:     loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/loss/loss_utils.py", line 67, in ForCausalLMLoss
[rank1]:     loss = fixed_cross_entropy(logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/transformers/loss/loss_utils.py", line 36, in fixed_cross_entropy
[rank1]:     loss = nn.functional.cross_entropy(source, target, ignore_index=ignore_index, reduction=reduction)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/workspace/quant-env/lib/python3.12/site-packages/torch/nn/functional.py", line 3462, in cross_entropy
[rank1]:     return torch._C._nn.cross_entropy_loss(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 24.55 GiB. GPU 1 has a total capacity of 79.25 GiB of which 23.00 GiB is free. Including non-PyTorch memory, this process has 56.24 GiB memory in use. Of the allocated memory 51.58 GiB is allocated by PyTorch, and 3.09 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
W0918 13:30:47.725000 6381 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6606 closing signal SIGTERM

[realAsma]

@eduardzl Could you try running the training without quantization? i.e, simply skip the argument --quant_cfg MXFP4_MLP_WEIGHT_ONLY_CFG ?

I initially suggested this would be a memory leak because the CUDA OOM is happening after many training steps in this.

However if you are running into issues during the first step itself, this is not likely a memory leak. I would guess the OOM is happening even for non-quantized training.