Skip to content

Commit dc2bcb7

Browse files
jepioCarlosGomes98filaretov
authored
[MLPerf 6.1][GRPO] Add RLVR GRPO reference (#890)
* Add reasoning benchmark scaffold * Remove setting seed * Update reasoning submodule pointer * Document GPUs per node for reasoning * add data processing instructions * update readme * update readme * update readme * update submodule * remove unnecessary parts of readme * update data instructions * Update SWE GRPO README Use placeholders for some parameters that still need to be fixed. * Update submodule to mlperf-training-qwen35 branch * Update SWE GRPO paths for Qwen3.5 branch * Add software versions * Rephrase first paragraph * explicitly mention RLVR * expand GRPO * remove commit * Address review comments and add note on reasoning * Rename benchmark to llm_moe_grpo * Add index placeholder for GRPO benchmark * llm_moe_grpo: Update reference and sync README with it * llm_moe_grpo: Update submodule reference to NVIDIA-NeMo org * update readme link to correct repo --------- Co-authored-by: Carlos Gomes <cgomes@nvidia.com> Co-authored-by: Hristo Filaretov <hfilaretov@nvidia.com> Co-authored-by: Carlos Gomes <carlosmiguel.gomes@live.com.pt>
1 parent 004cb4a commit dc2bcb7

4 files changed

Lines changed: 371 additions & 0 deletions

File tree

‎.gitmodules‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,7 @@
66
path = text_to_image/torchtitan
77
url = https://github.com/pytorch/torchtitan.git
88
branch = mlperf-training-flux.1
9+
[submodule "llm_moe_grpo/RL"]
10+
path = llm_moe_grpo/RL
11+
url = https://github.com/NVIDIA-NeMo/RL.git
12+
branch = v0.6.0-mlperf-training-qwen35

‎README.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,16 @@ Each benchmark will run until the target quality is reached and then stop, print
4242

4343
Some these benchmarks are rather slow or take a long time to run on the reference hardware. We expect to see significant performance improvements with more hardware and optimized implementations.
4444

45+
# MLPerf Training v6.1 (Submission Deadline XXX, 2026)
46+
47+
| Model | reference implementation | framework* | dataset | model parameter count**
48+
| ---- | ---- | ---- | ---- | ----
49+
| qwen3.5_397b_a17b_swe_grpo | [llm_moe_grpo](https://github.com/mlcommons/training/tree/master/llm_moe_grpo) | NeMo-RL / NeMo-Gym | SWE tasks | 397B
50+
51+
*Framework here is given for the reference implementation. Submitters are free to use their own frameworks to run the benchmark.
52+
53+
**Model parameter count is not the same as active parameter that are being trained in the benchmark.
54+
4555
# MLPerf Training v6.0 (Submission Deadline May 15, 2026)
4656

4757
| Model | reference implementation | framework* | dataset | model parameter count**

‎llm_moe_grpo/README.md‎

Lines changed: 356 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,356 @@
1+
# 1. Problem
2+
3+
## SWE Agent Reinforcement Learning - GRPO with NeMo-Gym SWE/OpenHands.
4+
5+
[NeMo-RL](https://github.com/NVIDIA-NeMo/RL/tree/v0.6.0-mlperf-training-qwen35) provides the implementation used for this benchmark from branch `mlperf-training-qwen35`. The benchmark uses Reinforcement Learning with Verifiable Rewards (RLVR) to train `Qwen/Qwen3.5-397B-A17B` with Group Relative Policy Optimization (GRPO) against a NeMo-Gym software-engineering environment driven by an OpenHands SWE agent.
6+
7+
The task is to improve the SWE agent's accuracy in solving held-out R2E-Gym software-engineering tasks. A rollout receives reward 1 when the generated patch passes the task evaluation and reward 0 otherwise.
8+
9+
The relevant config files are under `RL/examples/nemo_gym` and `RL/qwen_35`. The benchmark launch entrypoint is `RL/examples/nemo_gym/launch_qwen35_nemo_gym_multinode_training.sh`, using `RL/qwen_35/configs/grpo_qwen35_397b_swe_openhands_async_benchmark.yaml`.
10+
11+
# 2. Directions
12+
13+
## Steps to configure machine
14+
15+
To use this repository, please ensure your have access to a SLURM cluster with Enroot/Pyxis and at least 256 GB200 GPUs.
16+
17+
### Container setup
18+
19+
The Dockerfile to build for this benchmark is the NeMo-RL v0.6.0 Gym overlay Dockerfile.
20+
21+
```bash
22+
cd RL
23+
docker buildx build \
24+
--platform <linux/amd64 or linux/arm64> \
25+
-t <tag> \
26+
-f docker/Dockerfile.gym_v0.6.0 \
27+
.
28+
```
29+
30+
The Dockerfile overlays the SWE/NeMo-Gym pieces on top of `nvcr.io/nvidia/nemo-rl:v0.6.0` and prefetches Gym virtual environments for `qwen_35/configs/grpo_qwen35_397b_swe_openhands_async.yaml`.
31+
32+
## Steps to download and verify data
33+
34+
The run requires the following artifacts:
35+
36+
| Artifact | Description | Status |
37+
|---|---|---|
38+
| Policy model | Host directory containing `Qwen/Qwen3.5-397B-A17B`, passed as `HF_CKPT_PATH`, mounted into the container, and exposed to the recipe through `CONTAINER_HF_CKPT_PATH` | Download from Hugging Face |
39+
| Megatron-Core checkpoint cache | Host directory for the HF-to-Megatron converted checkpoint cache, passed as `NRL_MEGATRON_CHECKPOINT_DIR` and mounted into the container | Empty directory is allowed on first run |
40+
| Training JSONL | Host path to NeMo-Gym SWE training tasks, passed as `NEMO_GYM_SWE_TRAIN_DATA_PATH` and mounted into the container | Download `benchmark_r2e_gym_easy_train.jsonl` or rebuild with `RL/tools/create_r2e_gym_easy_subset_jsonl.py` |
41+
| Validation JSONL | Host path to NeMo-Gym SWE validation tasks, passed as `NEMO_GYM_SWE_VALIDATION_DATA_PATH` and mounted into the container | Download `benchmark_r2e_gym_easy_val.jsonl` or rebuild with `RL/tools/create_r2e_gym_easy_subset_jsonl.py` |
42+
| Task containers | Host directory containing Apptainer/Singularity SIF images in the layout expected by the recipe, passed as `NEMO_GYM_SWE_SIF_DIR` and mounted into the container | Build with `RL/docker/dataset-processing-container` |
43+
44+
To download the training and validation JSONL files using the HuggingFace CLI:
45+
46+
```bash
47+
hf download hfilaretov/Benchmark-R2E-Gym-Easy --repo-type dataset --local-dir hfilaretov__Benchmark-R2E-Gym-Easy
48+
...
49+
50+
tree hfilaretov__Benchmark-R2E-Gym-Easy
51+
hfilaretov__Benchmark-R2E-Gym-Easy
52+
├── benchmark_r2e_gym_easy_train.jsonl
53+
├── benchmark_r2e_gym_easy_val.jsonl
54+
└── README.md
55+
56+
1 directory, 3 files
57+
```
58+
59+
The environment also requires per-task SIF images. The recipe resolves task containers from `sif_dir` with this template:
60+
61+
```yaml
62+
- "${sif_dir}/r2egym/{instance_id}.sif"
63+
```
64+
65+
The task containers have to be built and converted to SIF format, please see [Section 3](#data-preprocessing) below.
66+
67+
### Model cache setup
68+
69+
From outside the container, download the Hugging Face model into a host directory. The launcher requires `HF_CKPT_PATH` to point at this directory, mounts that path into the container, and exports `CONTAINER_HF_CKPT_PATH` to the recipe as `policy.model_name`. By default, `CONTAINER_HF_CKPT_PATH` is the same path as `HF_CKPT_PATH`.
70+
71+
```bash
72+
python -m venv .venv
73+
source .venv/bin/activate
74+
pip install huggingface_hub
75+
76+
export HF_CKPT_PATH=$(pwd)/hf/Qwen/Qwen3.5-397B-A17B
77+
mkdir -p "$HF_CKPT_PATH"
78+
HF_TOKEN=<your hf token> hf download Qwen/Qwen3.5-397B-A17B --local-dir "$HF_CKPT_PATH"
79+
```
80+
81+
The launcher also creates and mounts a host Hugging Face cache. Set `HF_HOME` before launch if you want to use a cache outside `$(pwd)/.cache`.
82+
83+
## Steps to run and time
84+
85+
All steps below are assumed to be run from this `llm_moe_grpo` directory on the host; `cd RL` enters the NeMo-RL submodule checkout. The launcher submits `ray.sub` and runs training from the checkout baked into the container at `/opt/nemo-rl` by default.
86+
87+
```bash
88+
cd RL
89+
90+
export REPO_LOCATION=$(pwd)
91+
export EXP_NAME=<experiment name>
92+
export CONTAINER_IMAGE_PATH=<container image path or tag>
93+
export SLURM_ACCOUNT=<account>
94+
export SLURM_PARTITION=<partition>
95+
export GPUS_PER_NODE=4 # GB200 reference configuration
96+
export HF_CKPT_PATH=<host path to HF checkpoint directory>
97+
export NRL_MEGATRON_CHECKPOINT_DIR=<host path to Megatron-Core checkpoint cache directory> # may be empty on first run
98+
export NEMO_GYM_SWE_TRAIN_DATA_PATH=<host path to training JSONL>
99+
export NEMO_GYM_SWE_VALIDATION_DATA_PATH=<host path to validation JSONL>
100+
export NEMO_GYM_SWE_SIF_DIR=<host directory containing SWE task SIF images>
101+
102+
# Optional authentication/logging.
103+
export HF_TOKEN=<huggingface token>
104+
export WANDB_API_KEY=<wandb token>
105+
export MLPERF_TARGET_ACCURACY=<target reward mean> # default: 1.0 until the target is finalized
106+
export GRPO_SEED=<integer seed> # default: random per launch
107+
108+
# Defaults are defined by the launcher and may be overridden here.
109+
export TRAIN_NODES=<number of training nodes> # default: 32
110+
export GEN_NODES=<number of generation nodes> # default: 32
111+
export SLURM_TIME=<walltime> # default: 1:0:0
112+
export RECIPE=qwen_35/configs/grpo_qwen35_397b_swe_openhands_async_benchmark.yaml
113+
114+
# Optional extra mounts. The launcher automatically mounts the paths above.
115+
export EXTRA_MOUNTS=<host_path>:<container_path>[,<host_path>:<container_path>...]
116+
117+
bash examples/nemo_gym/launch_qwen35_nemo_gym_multinode_training.sh
118+
```
119+
120+
The launcher also accepts `NODES` to override `TRAIN_NODES + GEN_NODES`, `CONTAINER_REPO_LOCATION` to override the baked checkout path `/opt/nemo-rl`, `CONTAINER_INPUT_ROOT` and the `CONTAINER_*` path variables to override container-side paths.
121+
122+
# 3. Dataset/Environment
123+
124+
### Publication/Attribution
125+
126+
We use a subset of the [R2E-Gym/R2E-Gym-Subset](https://huggingface.co/datasets/R2E-Gym/R2E-Gym-Subset) dataset.
127+
128+
### Data preprocessing
129+
130+
The recipe consumes prebuilt JSONL files from [Benchmark-R2E-Gym-Easy](https://huggingface.co/datasets/hfilaretov/Benchmark-R2E-Gym-Easy).
131+
Each row represents a software-engineering task for the NeMo-Gym environment. We filtered the original `R2E-Gym/R2E-Gym-Subset` dataset based on these two conditions:
132+
* whether an environment container image successfully builds for both x86_64 and aarch64
133+
* complexity using the following condition:
134+
```
135+
where num_non_test_func_methods == 1 | where num_non_test_files == 1 | where num_non_test_lines <= 20
136+
```
137+
138+
To build the JSONL files yourself, run the converter from the RL checkout:
139+
140+
```bash
141+
cd RL
142+
143+
# Optional token
144+
export HF_TOKEN=<read-token>
145+
hf download R2E-Gym/R2E-Gym-Subset --repo-type dataset --local-dir tmp/R2E-Gym__R2E-Gym-Subset
146+
uv run --with pyarrow python tools/create_r2e_gym_easy_subset_jsonl.py \
147+
--dataset-dir tmp/R2E-Gym__R2E-Gym-Subset \
148+
--output-dir outputs/data/ \
149+
--cache-dir tmp/r2e_repo_cache \
150+
--train-ids tools/train-instance-ids.txt \
151+
--val-ids tools/val-instance-ids.txt
152+
```
153+
154+
You'll have the relevant output files in `outputs`:
155+
156+
```bash
157+
wc -l outputs/data/benchmark_r2e_gym_easy_train.jsonl \
158+
outputs/data/benchmark_r2e_gym_easy_val.jsonl \
159+
outputs/data/r2e_gym_subset_full.jsonl
160+
721 outputs/data/benchmark_r2e_gym_easy_train.jsonl
161+
256 outputs/data/benchmark_r2e_gym_easy_val.jsonl
162+
4578 outputs/data/r2e_gym_subset_full.jsonl
163+
5555 total
164+
```
165+
166+
The JSONL files refer to SIF container files that need to be generated.
167+
This is a two-step process:
168+
1. Images are built from the repository and git revision specified in the dataset.
169+
2. These images are converted to SIF file format.
170+
171+
You can build the container defined in `RL/docker/dataset-processing-container` that already pre-packages all necessary dependencies and can be used for both steps.
172+
173+
Prepare the builder image:
174+
175+
```bash
176+
cd RL/docker/dataset-processing-container
177+
export DOCKER_REGISTRY=<your-container-registry>
178+
docker build --push -t $DOCKER_REGISTRY/grpo-data-builder:latest .
179+
```
180+
181+
Note: to build the dataset images within the builder image, you need to mount the Docker daemon socket inside the container.
182+
If you do not want to do that, please set up an environment equivalent to the builder image, and then run the scripts outside a container.
183+
184+
To build the images and push them to a registry, run on a host that has your target architecture to build natively:
185+
186+
```bash
187+
export HF_TOKEN=<read-token-for-huggingface>
188+
export DOCKER_REGISTRY=<url-to-docker-registry>
189+
export DOCKER_TOKEN=<docker-registry-token>
190+
export DOCKER_USER=<docker-registry-username>
191+
export STATE_DIR=<path-to-persistent-storage>
192+
export MAX_WORKERS=<maximum-number-of-parallel-build-tasks>
193+
194+
# R2E-Gym Easy subset
195+
docker run -it --rm \
196+
-v /var/run/docker.sock:/var/run/docker.sock \
197+
-v $STATE_DIR:/workspace/state \
198+
-e DOCKER_REGISTRY -e DOCKER_TOKEN -e DOCKER_USER -e HF_TOKEN -e MAX_WORKERS \
199+
$DOCKER_REGISTRY/grpo-data-builder:latest \
200+
/workspace/run-r2e-gym-build-images.sh
201+
```
202+
203+
To convert the images from the registry to SIF files:
204+
205+
```bash
206+
export SIF_LOCAL_DIR=<local-directory-to-store-sif-containers>
207+
208+
# SIF images, final dataset
209+
docker run -it --rm \
210+
-v $SIF_LOCAL_DIR:/opt/data \
211+
-e DOCKER_REGISTRY -e DOCKER_TOKEN -e DOCKER_USER -e HF_TOKEN -e MAX_WORKERS \
212+
$DOCKER_REGISTRY/grpo-data-builder:latest \
213+
/workspace/run-build-sif-images.sh
214+
```
215+
216+
Please note that the above container will use its local storage to build the SIF files and then copy them over to your `SIF_LOCAL_DIR`.
217+
You therefore might be constrained in the number of `$MAX_WORKERS` by your available local storage.
218+
219+
### Training and test data separation
220+
221+
The config uses separate training and validation JSONL files:
222+
223+
```yaml
224+
policy:
225+
model_name: ${oc.env:CONTAINER_HF_CKPT_PATH}
226+
data:
227+
train:
228+
data_path: ${oc.env:NEMO_GYM_SWE_TRAIN_DATA_PATH}
229+
validation:
230+
data_path: ${oc.env:NEMO_GYM_SWE_VALIDATION_DATA_PATH}
231+
sif_dir: ${oc.env:NEMO_GYM_SWE_SIF_DIR}
232+
```
233+
234+
The official split is defined by the fixed instance-id files `RL/tools/train-instance-ids.txt` and `RL/tools/val-instance-ids.txt`. The conversion script validates that the lists do not overlap, writes matching rows to `benchmark_r2e_gym_easy_train.jsonl` and `benchmark_r2e_gym_easy_val.jsonl`, and leaves rows in neither list only in `r2e_gym_subset_full.jsonl`.
235+
236+
### Training data order
237+
238+
Training data order is preserved by the recipe with `data.shuffle: false`. The converter writes the training JSONL in the order encountered in the converted R2E-Gym subset after filtering by `RL/tools/train-instance-ids.txt`; the benchmark does not add runtime shuffling.
239+
240+
### Test data order
241+
242+
Validation data order is preserved by the recipe. The config uses `grpo.max_val_samples: null`, so validation thoroughness is inferred from the validation dataset size unless overridden. The benchmark does not add runtime shuffling of validation data.
243+
244+
### Simulation environment (RL models only)
245+
246+
The benchmark uses NeMo-Gym with the SWE/OpenHands agent configuration. Rollouts are collected through a vLLM-backed policy server, with OpenHands interacting with task containers via Apptainer/Singularity. Model reasoning/thinking is disabled through the vLLM chat template configuration. The async recipe uses non-colocated generation and training, with one-step-stale trajectories corrected by importance sampling.
247+
248+
# 4. Model
249+
250+
### Publication/Attribution
251+
252+
The policy starts from the [`Qwen/Qwen3.5-397B-A17B`](https://huggingface.co/Qwen/Qwen3.5-397B-A17B) checkpoint released by the Qwen team. The reference training implementation is NeMo-RL with Qwen 3.5 support from the `mlperf-training-qwen35` branch.
253+
254+
### Model details
255+
256+
Architecture values below are taken from the [Hugging Face model card](https://huggingface.co/Qwen/Qwen3.5-397B-A17B) and [`config.json`](https://huggingface.co/Qwen/Qwen3.5-397B-A17B/blob/main/config.json).
257+
258+
| Config | Value |
259+
| :-- | :-- |
260+
| # Total Parameters | 397B |
261+
| # Active Parameters | 17B |
262+
| # Layers | 60 |
263+
| Hidden Layout | 3 Gated DeltaNet + 1 Gated Attention layers per block (15 blocks) |
264+
| Attention Type | Hybrid Gated DeltaNet + Gated Attention |
265+
| Gated DeltaNet Heads (V / QK) | 64 / 16 |
266+
| Gated DeltaNet Head Dimension | 128 |
267+
| Gated Attention Heads (Q / KV) | 32 / 2 |
268+
| Gated Attention Head Dimension | 256 |
269+
| RoPE Dimension | 64 |
270+
| Model Dimension | 4,096 |
271+
| # Routed Experts | 512 |
272+
| # Active Routed Experts | 10 |
273+
| # Shared Experts | 1 |
274+
| Expert Intermediate Dimension | 1,024 |
275+
| Activation | SiLU (SwiGLU in MoE) |
276+
| Normalization | RMSNorm |
277+
| Vocab Size | 248,320 |
278+
| Native Context Length | 262,144 |
279+
| Benchmark Context Length | 65,536 |
280+
281+
### Benchmark runtime
282+
283+
| **Component** | **Architecture** | **Parameters** | **Technical Details** |
284+
|---------------|------------------|----------------|-----------------------|
285+
| **Training runtime** | Megatron-Bridge and Megatron-Core through NeMo-RL | Same policy weights | TP4 x PP2 x CP1, EP32, BF16 |
286+
| **Generation runtime** | vLLM | Same policy weights | TP8, EP8, 64k benchmark context, HTTP server exposed for NeMo-Gym |
287+
| **SWE environment** | NeMo-Gym + OpenHands | N/A | CodeActAgent, max 30 turns |
288+
289+
Source revisions identify the checked-out editable packages used by the reference implementation.
290+
291+
| **Runtime package** | **Package version** | **Source revision** |
292+
|---------------------|---------------------|---------------------|
293+
| NeMo-RL | `0.6.0` | `fbc91daf` |
294+
| Megatron-Bridge | `0.5.0` | `95e5f38f` |
295+
| Megatron-Core | `0.18.0` | `d30c3ae54` |
296+
| vLLM | `0.17.1` | PyPI package pin |
297+
| NeMo-Gym | `0.3.0rc0` | `1a4912e` |
298+
299+
### Weight and bias initialization
300+
301+
Training starts from the pretrained Hugging Face checkpoint converted to Megatron-Core format. Random initialization is not used for the policy model. The first run can populate `NRL_MEGATRON_CHECKPOINT_DIR` with the converted checkpoint cache.
302+
303+
MoE router weights are kept frozen.
304+
305+
### Loss function
306+
307+
The recipe uses token-level GRPO with reward normalization and a leave-one-out baseline. Reference-policy KL is disabled (`reference_policy_kl_penalty: 0`), and the async recipe uses importance-sampling correction for one-step-stale rollouts.
308+
309+
### Optimizer
310+
311+
AdamW with distributed optimizer state.
312+
313+
| Parameter | Value |
314+
| :-- | :-- |
315+
| Optimizer | AdamW |
316+
| Base learning rate | `2.0e-6` |
317+
| End learning rate | `2.0e-6` |
318+
| Learning-rate schedule | Constant |
319+
| Warmup steps | 2 |
320+
| Weight decay | `0.0` |
321+
| Adam beta1 | `0.9` |
322+
| Adam beta2 | `0.999` |
323+
| Adam epsilon | `1e-8` |
324+
| Gradient clipping | `1.0` |
325+
| Distributed optimizer | Enabled |
326+
| Optimizer parameters | FP32 |
327+
| Training precision | BF16 |
328+
329+
### Precision
330+
331+
The recipe uses BF16 policy precision by default.
332+
333+
# 5. Quality
334+
335+
### Quality metric
336+
337+
The quality metric is `val:accuracy`, computed from NeMo-Gym validation rollouts.
338+
339+
### Quality target
340+
341+
TODO: final
342+
343+
The quality target is pending MLCommons ratification. The current launcher reads `MLPERF_TARGET_ACCURACY` and defaults to `1.0`.
344+
345+
### Evaluation frequency
346+
347+
| Parameter | Value |
348+
| :-- | :-- |
349+
| Evaluate at start | Yes |
350+
| Evaluation period | Every 2 training steps |
351+
| Evaluate at end | Yes |
352+
| Maximum training steps | 20 |
353+
354+
### Evaluation thoroughness
355+
356+
The validation JSONL contains 256 R2E-Gym tasks and each evaluation uses the full validation set.

‎llm_moe_grpo/RL‎

Submodule RL added at fbc91da

0 commit comments

Comments
 (0)