NVIDIA-NeMo / NVIDIA-NeMo/RL

OOM during refit in mcore path when using larger bucket sizes

Open
#630 6 comments 0 reactions 2 assignees Claimed by @guyueh1 View on GitHub
t-mcore vllm
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.