Does AutoTP support sharding the optimizer states?

[hijkzzz]

Previously, using ZeRO-3 (dp size = 4) to train llama3 8b with 4 A100 80G didn't cause OOM, but after switching to ZeRO-2 (dp size =2) with AutoTP(tp size = 2) or AutoTP(tp size = 4), it runs out of memory. I'm not sure why.

code: https://github.com/OpenRLHF/OpenRLHF/tree/auto_tp
script:

deepspeed --module openrlhf.cli.train_sft \
    --max_len 2048 \
    --dataset Open-Orca/OpenOrca \
    --input_key question \
    --output_key response \
    --train_batch_size 128 \
    --micro_train_batch_size 4 \
    --max_samples 1024 \
    --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
    --save_path ./checkpoint/llama3-8b-sft \
    --save_steps -1 \
    --logging_steps 1 \
    --eval_steps -1 \
    --zero_stage 1 \
    --max_epochs 1 \
    --bf16 \
    --flash_attn \
    --learning_rate 5e-6 \
    --gradient_checkpointing \
    --packing_samples \
    --ds_tensor_parallel_size 4 \

logs

[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.96 GiB. GPU 1 has a total capacity of 79.14 GiB of which 1.01 GiB is free. Including non-PyTorch memory, this process has 0 bytes memory in use. Of the allocated memory 76.49 GiB is allocated by PyTorch, and 76.30 MiB 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)
[rank2]: Traceback (most recent call last):
[rank2]:   File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
[rank2]:     return _run_code(code, main_globals, None,
[rank2]:   File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
[rank2]:     exec(code, run_globals)
[rank2]:   File "/openrlhf/openrlhf/cli/train_sft.py", line 284, in <module>
[rank2]:     train(args)
[rank2]:   File "/openrlhf/openrlhf/cli/train_sft.py", line 141, in train
[rank2]:     trainer.fit(args, consumed_samples, num_update_steps_per_epoch)
[rank2]:   File "/openrlhf/openrlhf/trainer/sft_trainer.py", line 154, in fit
[rank2]:     self.strategy.backward(loss, self.model, self.optimizer)
[rank2]:   File "/openrlhf/openrlhf/utils/deepspeed/deepspeed.py", line 142, in backward
[rank2]:     model.backward(loss)
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank2]:     ret_val = func(*args, **kwargs)
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2216, in backward
[rank2]:     self._do_optimizer_backward(loss, retain_graph)
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2175, in _do_optimizer_backward
[rank2]:     self.optimizer.backward(loss, retain_graph=retain_graph)
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py", line 327, in backward
[rank2]:     loss.backward(retain_graph=retain_graph, **bwd_kwargs)
[rank2]:   File "/usr/local/lib/python3.10/dist-packages/torch/_tensor.py", line 626, in backward
[rank2]:     torch.autograd.backward(
[rank2]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/__init__.py", line 347, in backward
[rank2]:     _engine_run_backward(
[rank2]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
[rank2]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py", line 551, in accumulate_hp_grads_and_remove_lp
[rank2]:     self.accumulate_hp_grads_and_remove_lp(param, i, j)
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py", line 541, in accumulate_hp_grads_and_remove_lp
[rank2]:     self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=False)
[rank2]:   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py", line 341, in _update_hp_grad
[rank2]:     hp_grad.data.add_(lp.grad.data.to(hp_grad.dtype).view(hp_grad.shape))
[rank2]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.96 GiB. GPU 2 has a total capacity of 79.14 GiB of which 1.01 GiB is free. Including non-PyTorch memory, this process has 0 bytes memory in use. Of the allocated memory 76.49 GiB is allocated by PyTorch, and 76.30 MiB 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.

[inkcherry]

hi, @hijkzzz , It supports shardingoptimizer_states across dp_groups for ZeRO-1 and ZeRO-2,For safety, you can use print(model) to check whether it has been correctly partitioned.

However, by default, TP does not partition the embedding and lm_head layers( This is a deliberate trade-off between memory and performance, as the machines used in experiments have memory capacities greater than 80GB(96GB), Perhaps they need you to lower the mbs) For models with a large vocab_size, memory usage in these layers can be more, we're planning to make this behavior configurable—some work is still in progress. #7181
Thank you for your feedback! You're very welcome to share more suggestions or discussion about this feature to help us improve it further, and feel free to ping me if you need any help.

[hijkzzz]

hi, @hijkzzz , It supports shardingoptimizer_states across dp_groups for ZeRO-1 and ZeRO-2,For safety, you can use print(model) to check whether it has been correctly partitioned.

However, by default, TP does not partition the embedding and lm_head layers( This is a deliberate trade-off between memory and performance, as the machines used in experiments have memory capacities greater than 80GB, Perhaps they need you to lower the max mbs) For models with a large vocab_size, memory usage in these layers can be more, we're planning to make this behavior configurable—some work is still in progress. #7181 Thank you for your feedback! You're very welcome to share more suggestions or discussion about this feature to help us improve it further, and feel free to ping me if you need any help.

the logs show that stage_1_and_2.py", line 1891, in step (PPOTrainer pid=555946) fp32_partition.to(get_accelerator().current_device_name()).data) (PPOTrainer pid=555946) torch.OutOfMemoryError: CUDA out of memory

(PPOTrainer pid=555946)   File "/tmp/ray/session_2025-04-25_08-49-20_792240_524513/runtime_resources/working_dir_files/_ray_pkg_5477d70341eb83b5/openrlhf/trainer/ray/ppo_actor.py", line 447, in fit
(PPOTrainer pid=555946)     status = self.trainer.ppo_train(kl_ctl)
(PPOTrainer pid=555946)   File "/tmp/ray/session_2025-04-25_08-49-20_792240_524513/runtime_resources/working_dir_files/_ray_pkg_5477d70341eb83b5/openrlhf/trainer/ray/ppo_actor.py", line 165, in ppo_train
(PPOTrainer pid=555946)     status = self.training_step(experience, kl_ctl)
(PPOTrainer pid=555946)   File "/tmp/ray/session_2025-04-25_08-49-20_792240_524513/runtime_resources/working_dir_files/_ray_pkg_5477d70341eb83b5/openrlhf/trainer/ray/ppo_actor.py", line 249, in training_step
(PPOTrainer pid=555946)     self.strategy.optimizer_step(self.actor_optim, self.actor, self.actor_scheduler, name="actor")
(PPOTrainer pid=555946)   File "/tmp/ray/session_2025-04-25_08-49-20_792240_524513/runtime_resources/working_dir_files/_ray_pkg_5477d70341eb83b5/openrlhf/utils/deepspeed/deepspeed.py", line 154, in optimizer_step
(PPOTrainer pid=555946)     model.step()
(PPOTrainer pid=555946)   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2378, in step
(PPOTrainer pid=555946)     self._take_model_step(lr_kwargs)
(PPOTrainer pid=555946)   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2281, in _take_model_step
(PPOTrainer pid=555946)     self.optimizer.step()
(PPOTrainer pid=555946)   File "/home/jianh/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1891, in step
(PPOTrainer pid=555946)     fp32_partition.to(get_accelerator().current_device_name()).data)
(PPOTrainer pid=555946) torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 29.91 GiB. GPU 0 has a total capacity of 79.10 GiB of which 15.15 GiB is free. Process 1647313 has 4.70 GiB memory in use. Process 1674734 has 23.29 GiB memory in use. Process 1676063 has 18.37 GiB memory in use. Process 1678007 has 17.39 GiB memory in use. Of the allocated memory 17.16 GiB is allocated by PyTorch, and 513.41 MiB 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)