OOM during refit in mcore path when using larger bucket sizes
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
Using a larger bucket size for refit by setting a higher percentage of available bytes [here](https://github.com/NVIDIA-NeMo/RL/blob/c5ee4a5460e4844ce2aac919aec0ac4eba733b70/nemo_rl/models/policy/megatron_policy_worker.py#L1449) (e.g. 0.4) results in OOM for DSv2-Lite model.
Sample command to reproduce (after increasing bucket size to 40%):
```
uv run python examples/run_grpo_math.py --config=examples/configs/grpo_math_1B_megatron.yaml grpo.val_batch_size=2 policy.model_name=deepseek-ai/DeepSeek-V2-Lite-Chat cluster.gpus_per_node=8 policy.megatron_cfg.pipeline_model_parallel_size=4 policy.megatron_cfg.num_layers_in_first_pipeline_stage=7 policy.megatron_cfg.num_layers_in_last_pipeline_stage=6 policy.refit_buffer_size_gb=8 policy.max_total_sequence_length=1024 checkpointing.enabled=False checkpointing.save_period=5 grpo.val_period=-1 grpo.max_val_samples=16 grpo.val_batch_size=4 checkpointing.keep_top_k=100 checkpointing.checkpoint_dir=results/dsv2-lite-testsave grpo.val_at_start=False grpo.max_val_samples=16 policy.megatron_cfg.expert_model_parallel_size=2 policy.megatron_cfg.apply_rope_fusion=False
```
This seems to happen because we use double the memory when [packing the params](https://github.com/NVIDIA-NeMo/RL/blob/c5ee4a5460e4844ce2aac919aec0ac4eba733b70/nemo_rl/models/policy/megatron_policy_worker.py#L1506-L1522). We may need to multiply the size_in_bytes of each param by 2 in `prepare_weights_for_ipc`.
Contributor guide
Assessment
This issue has not been assessed yet.