70B GRPO training is slower than reported
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
When training GRPO with Llama 3.1 70B, we observe moderately slower performance compared to what is reported in the [NVIDIA NeMo blog post](https://nvidia-nemo.github.io/blog/2025/07/21/nemo-rl-v0.3/).
This slowdown occurs with both the DTensor backend and the Megatron backend.
| Model| Backend| Nodes | GPUs per node | Total step time (s) | Policy training (s) | Refit (s) | Generation (s) | Get logprobs (s) | Avg. generated tokens per sample |
|---|---|---|---|---|---|---|---|---|---|
| Llama 3.1-8B Instruct | Megatron | 1 | 8 | 112 | 28 | 5 | 58 | 18 | 795 |
| | Megatron (Our rerun) | 1 | 8 | 116 | 27 | 8 | 56 | 18 | 811 |
| | PyT DTensor | 1 | 8 | 122 | 38 | 4 | 57 | 19 | 777 |
| | PyT DTensor (Our rerun) | 1 | 8 | 129 | 40 | 6 | 55 | 19 | 802 |
| Llama 3.1-70B Base | Megatron | 8 | 8 | 147 | 28 | 14 | 84 | 18 | 398 |
| | Megatron (Our rerun) | 8 | 8 | 178 | 31 | 26 |87 | 20 | 396 |
| | PyT DTensor | 8 | 8 | 230 | 97 | 15 | 82 | 28 | 395 |
| | PyT DTensor (Our rerun) | 8 | 8 | 266 | 117 | 20 | 88 | 29 | 411 |
- 8B GRPO: almost as fast as reported in the blog post
- 70B GRPO: Noticeably slower for both DTensor and Megatron backend. In particular,
- refit is ~2x slower on both DTensor and Megatron backends
- Policy training is also slower with DTensor
Do you have any suggestions or insights into why the refit step is much slower in our setup?
**Steps/Code to reproduce bug**
We are using the configurations described in the blog post, which corresponds to:
```
## dtensor
uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B.yaml \
policy.model_name=meta-llama/Llama-3.1-70B policy.tokenizer.name=meta-llama/Llama-3.1-70B-Instruct \
policy.generation.vllm_cfg.tensor_parallel_size=4 policy.max_total_sequence_length=4096 \
cluster.num_nodes=8 policy.dtensor_cfg.enabled=True policy.dtensor_cfg.tensor_parallel_size=8 \
policy.dtensor_cfg.sequence_parallel=True policy.dtensor_cfg.activation_checkpointing=False \
loss_fn.use_importance_sampling_correction=True \
policy.sequence_packing.enabled=False \
policy.dynamic_batching.enabled=True # as reported in the blog, we use dynamic batching instead of sequence packing
## megatron
uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_70B_megatron.yaml \
policy.model_name=meta-llama/Llama-3.1-70B policy.tokenizer.name=meta-llama/Llama-3.3-70B-Instruct \
policy.sequence_packing.enabled=True loss_fn.use_importance_sampling_correction=True
```
**Expected behavior**
GRPO training should be as fast.
**Environment overview (please complete the following information)**
- Environment location: Inhouse Slurm cluster with H100 cards
- Method of install: Internal Dockerfile to make NemoRL run on our cluster. NemoRL codebase is strictly based on v0.3.1 release.
**Environment details**
- OS version: Dockerfile is based on `nvcr.io/nvidia/cuda-dl-base:25.05-cuda12.9-devel-ubuntu24.04`
- PyTorch version: 2.7.0
- Python version: 3.12
**Additional context**
GPU is NVIDIA H100 80GB
Contributor guide
Assessment
This issue has not been assessed yet.