NeMo-RL DPO OOM issues with Nemotron-3-Nano-30B

[slic33]

Describe the bug

Ran NeMo-RL DPO on Nemotron-3-Nano-30B BF16 with custom {prompt, chosen, rejected} dataset and consistently hit OOM during DTensorPolicyWorker init. Set-up: Single node, 4×A100-80GB (Brev), TP=4, CPU offload, activation checkpointing, long context (~3.2-3.4k tokens). Open questions: Is 4×80GB expected to be insufficient for this recipe? Any known working DPO config for Nemotron-3-Nano-30B? Is DPO + LoRA supported?

Expected behavior

Provide working DPO configuration for Nemotron-3-Nano-30B. Document memory requirements and LoRA+DPO support.

[yfw]

@RayenTian would you be able to help with this?

[RayenTian]

Hi, @slic33. Thank you for opening this issue and reporting the problem.

1. Currently the dtensor backend does not support enabling TP for nano v3.
2. 4×80GB GPUs are not sufficient to run nano v3. We tested with 80GB H100 GPUs, and at least a 2n8g setup is required.
3. DPO LoRA is currently supported. You can refer to the tests/functional/dpo_automodel_lora.sh script.
4. We don’t currently have a public DPO recipe for nano v3. However, you can refer to the SFT recipe, as the two are similar.
   • sft-nanov3-30BA3B-2n8g-fsdp2.yaml
   • sft-nanov3-30BA3B-2n8g-fsdp2-lora.yaml
5. Or you can also directly launch the DPO task via the command, you only need to modify the specific configurations. To enable lora, just change policy.dtensor_cfg.lora_cfg.enabled to true. This command is a version I've tested and confirmed to work properly.

NRL_FORCE_REBUILD_VENVS=true \
uv run examples/run_dpo.py   \
  --config examples/configs/dpo.yaml \
  policy.model_name="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" \
  policy.tokenizer.name="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" \
  policy.dtensor_cfg.tensor_parallel_size=1 \
  policy.max_total_sequence_length=4096 \
  policy.train_global_batch_size=16 \
  policy.train_micro_batch_size=1 \
  policy.dtensor_cfg.enabled=true \
  ++policy.dtensor_cfg._v2=true \
  policy.dtensor_cfg.activation_checkpointing=true \
  policy.dtensor_cfg.lora_cfg.enabled=false \
  policy.dtensor_cfg.lora_cfg.dim=128 \
  policy.dtensor_cfg.lora_cfg.alpha=256 \
  policy.dtensor_cfg.lora_cfg.exclude_modules=["*out_proj*"] \
  policy.dtensor_cfg.lora_cfg.match_all_linear=false \
  policy.dtensor_cfg.lora_cfg.use_triton=false \
  dpo.max_num_steps=20 \
  dpo.val_period=50 \
  cluster.gpus_per_node=8 \
  cluster.num_nodes=2 \
  checkpointing.enabled=false

6. Currently, enabling both cpu_offload and val_at_start for DPO at the same time will encounter some issues. I have submitted a PR! 1930 to fix this. You can either:
   • Temporarily avoid enabling both cpu_offload and val_at_start for DPO until PR! 1930 merged.
   • Switch to the ruit/fix_dpo_cpu_offload branch to run.

For a max_total_sequence_length=4096 setting, without enabling CPU offload and LoRA, only 7 steps can be executed. Therefore, if the sequence length is relatively long, it is better to switch to branch ruit/fix_dpo_cpu_offload for running.

If anything is unclear, please let me know and we can discuss it further.