NVIDIA-NeMo / NVIDIA-NeMo/RL

70B GRPO training is slower than reported

Open
#959 5 comments 0 reactions 1 assignee Claimed by @guyueh1 View on GitHub
bug community-request external Performance waiting-on-customer x-SBIntuitions
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

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.