fix: run dynamic sampling on unshaped rewards

docs/guides/dapo.md

@@ -13,10 +13,10 @@ This document focuses on DAPO-specific features: Dynamic Sampling and Overlong R
 
 ## Quickstart: Launch a DAPO Run
 
-To get started quickly, use the example configuration [examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml](../../examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml). You can launch this using the same script as GRPO:
+To get started quickly, use the example configuration [examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml](../../examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml). You can launch this using the same script as GRPO:
 
 ```bash
-uv run examples/run_grpo.py --config examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml {overrides}
+uv run examples/run_grpo.py --config examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml {overrides}
 ```
 
 **Reminder**: Don't forget to set your HF_HOME, WANDB_API_KEY, and HF_DATASETS_CACHE (if needed). You'll need to do a `huggingface-cli login` as well for LLaMA models.
@@ -85,10 +85,10 @@ grpo:
 > - Set `ratio_clip_max` to enable Clip-Higher (e.g., `ratio_clip_max: 0.28`)
 > - Set `token_level_loss: true` to enable Token-Level Policy Gradient Loss
 > 
-> See the full [DAPO example config](../../examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml) for reference.
+> See the full [DAPO example config](../../examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml) for reference.
 
 ## Example Training Results
-Using the [DAPO example config](../../examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml), you can expect to see intermediate plots such as the training reward curve and validation accuracy on AIME24 for Qwen/Qwen2.5-Math-7B. These plots serve as reference outputs to help verify reproducibility. They are not intended to reflect the best accuracy that can be achieved using DAPO for this model.
+Using the [DAPO example config](../../examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml), you can expect to see intermediate plots such as the training reward curve and validation accuracy on AIME24 for Qwen/Qwen2.5-Math-7B. These plots serve as reference outputs to help verify reproducibility. They are not intended to reflect the best accuracy that can be achieved using DAPO for this model.
 
 ![DAPO Qwen2.5-7B Training Reward](../assets/dapo_train_reward.png)
 ![DAPO Qwen2.5-7B Validation Accuracy](../assets/dapo_val_acc.png)

docs/guides/muon-optimizer.md

@@ -132,7 +132,7 @@ The command to generate the Muon results is:
 
 ```bash
 uv run examples/run_grpo_math.py \
-  --config examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml \
+  --config examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml \
   policy.megatron_cfg.enabled=true \
   policy.dtensor_cfg.enabled=false \
   ++policy.megatron_cfg.optimizer.optimizer=dist_muon \

docs/guides/prorlv2.md

@@ -1,6 +1,6 @@
 # An In-Depth Walkthrough of ProRLv2 in NeMo RL
 
-This guide covers the ProRLv2 configuration pattern in NeMo RL, based on the example config [`examples/configs/prorlv2.yaml`](../../examples/configs/prorlv2.yaml).
+This guide covers the ProRLv2 configuration pattern in NeMo RL, based on the example config [`examples/configs/prorlv2.v2.yaml`](../../examples/configs/prorlv2.v2.yaml).
 
 ProRLv2 is best thought of as **GRPO plus a bundle of stability/efficiency techniques** commonly used for long-horizon RL fine-tuning:
 
@@ -15,13 +15,13 @@ This document focuses on ProRLv2-specific knobs and gotchas. For foundational co
 
 ## Quickstart: Launch a ProRLv2 Run
 
-Use the example configuration [`examples/configs/prorlv2.yaml`](../../examples/configs/prorlv2.yaml):
+Use the example configuration [`examples/configs/prorlv2.v2.yaml`](../../examples/configs/prorlv2.v2.yaml):
 
 ```bash
-uv run examples/run_grpo_math.py --config examples/configs/prorlv2.yaml {overrides}
+uv run examples/run_grpo_math.py --config examples/configs/prorlv2.v2.yaml {overrides}
 ```
 
-`prorlv2.yaml` inherits from [`examples/configs/grpo_math_1B.yaml`](../../examples/configs/grpo_math_1B.yaml) and only overrides a small set of fields under `grpo` and `loss_fn`, plus output directories.
+`prorlv2.v2.yaml` inherits from [`examples/configs/grpo_math_1B.yaml`](../../examples/configs/grpo_math_1B.yaml) and only overrides a small set of fields under `grpo` and `loss_fn`, plus output directories.
 
 **Reminder**: Don't forget to set your `HF_HOME`, `WANDB_API_KEY`, and `HF_DATASETS_CACHE` (if needed). You'll need to do a `huggingface-cli login` as well for gated models.
 
@@ -47,7 +47,7 @@ Quick intuition:
 Computation (as implemented in this repo, with the ProRLv2 example defaults):
 
 ```text
-Defaults in examples/configs/prorlv2.yaml:
+Defaults in examples/configs/prorlv2.v2.yaml:
   grpo.adv_estimator.minus_baseline = true
   loss_fn.use_kl_in_reward          = false
 
@@ -204,26 +204,26 @@ Under Pure Online training (single update per rollout), the PPO policy ratio is
 - **Seq-mask-tis** — sequence-level filtering instead of token-level ICE-POP.
 - **No DAPO dynamic sampling** (`use_dynamic_sampling: false`).
 
-Use [`examples/configs/prorlv2_1_moe.yaml`](../../examples/configs/prorlv2_1_moe.yaml) directly. See the [ProRL v2.1 blog post](https://developer.nvidia.com/blog/scaling-llm-reinforcement-learning-with-prolonged-training-using-prorl-v2/) for the full motivation.
+Use [`examples/configs/prorlv2_1_moe.v2.yaml`](../../examples/configs/prorlv2_1_moe.v2.yaml) directly. See the [ProRL v2.1 blog post](https://developer.nvidia.com/blog/scaling-llm-reinforcement-learning-with-prolonged-training-using-prorl-v2/) for the full motivation.
 :::
 
 ```bash
 # Launch ProRL v2.1 for MoE models
-uv run examples/run_grpo_math.py --config examples/configs/prorlv2_1_moe.yaml {overrides}
+uv run examples/run_grpo_math.py --config examples/configs/prorlv2_1_moe.v2.yaml {overrides}
 ```
 
 ## Full Example Configs
 
-- **ProRLv2** (ICE-POP, DAPO, clipping): [`examples/configs/prorlv2.yaml`](../../examples/configs/prorlv2.yaml) — inherits from [`grpo_math_1B.yaml`](../../examples/configs/grpo_math_1B.yaml)
-- **ProRL v2.1** (seq-mask-tis, pure online, MoE): [`examples/configs/prorlv2_1_moe.yaml`](../../examples/configs/prorlv2_1_moe.yaml) — inherits from `prorlv2.yaml`, adds `force_on_policy_ratio`, switches to seq-mask-tis, disables dynamic sampling
+- **ProRLv2** (ICE-POP, DAPO, clipping): [`examples/configs/prorlv2.v2.yaml`](../../examples/configs/prorlv2.v2.yaml) — inherits from [`grpo_math_1B.yaml`](../../examples/configs/grpo_math_1B.yaml)
+- **ProRL v2.1** (seq-mask-tis, pure online, MoE): [`examples/configs/prorlv2_1_moe.v2.yaml`](../../examples/configs/prorlv2_1_moe.v2.yaml) — inherits from `prorlv2.v2.yaml`, adds `force_on_policy_ratio`, switches to seq-mask-tis, disables dynamic sampling
 
 ## Practical Overrides
 
 A few common overrides when launching:
 
 ```bash
 uv run examples/run_grpo_math.py \
-  --config examples/configs/prorlv2.yaml \
+  --config examples/configs/prorlv2.v2.yaml \
   policy.model_name="Qwen/Qwen2.5-1.5B" \
   logger.wandb_enabled=true \
   logger.wandb.project="prorlv2-dev" \
@@ -235,7 +235,7 @@ If you want to enable DAPO overlong reward shaping instead of stop-properly:
 
 ```bash
 uv run examples/run_grpo_math.py \
-  --config examples/configs/prorlv2.yaml \
+  --config examples/configs/prorlv2.v2.yaml \
   grpo.reward_shaping.stop_properly_penalty_coef=null \
   grpo.reward_shaping.overlong_buffer_length=4096 \
   grpo.reward_shaping.overlong_buffer_penalty=1.0 \

docs/guides/quantization-aware-rl.md

@@ -15,7 +15,7 @@ The following workflow + quantization recipe combinations have been validated en
 | Workflow | Quantization | Recipe | Status | Example Config |
 |---|---|---|---|---|
 | QA-Distillation | W4A4 | `NVFP4_DEFAULT_CFG` (NVFP4 weights + NVFP4 activations) | ✅ Converges | `examples/modelopt/qa_distillation_math_megatron.yaml` |
-| QA-GRPO | W4A16 | `examples/modelopt/quant_configs/nvfp4_a16.yaml` (NVFP4 weights, native-dtype activations) | ✅ Converges | `examples/modelopt/qa_grpo_llama8b_megatron.yaml` |
+| QA-GRPO | W4A16 | `examples/modelopt/quant_configs/nvfp4_a16.yaml` (NVFP4 weights, native-dtype activations) | ✅ Converges | `examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml` |
 | QA-GRPO | W4A4 | `NVFP4_DEFAULT_CFG` | ⚠️ Known convergence issue | `examples/modelopt/qa_grpo_math_megatron.yaml` |
 
 The `nvfp4_a16.yaml` custom YAML enables NVFP4 e2m1 weight quantization (with dynamic e4m3 micro-block scales) and leaves activations unquantized; weights are still exercised through both Megatron training and vLLM generation.
@@ -27,7 +27,7 @@ The `nvfp4_a16.yaml` custom YAML enables NVFP4 e2m1 weight quantization (with dy
 The QA-GRPO config extends the standard Megatron GRPO config by adding quantization parameters. See [Verified Configurations](#verified-configurations) for the status of W4A4 vs W4A16 on GRPO.
 
 ```yaml
-# examples/modelopt/qa_grpo_llama8b_megatron.yaml
+# examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml
 defaults: "../configs/grpo_math_8B_megatron.yaml"
 
 policy:
@@ -47,15 +47,15 @@ policy:
 
 ```bash
 uv run examples/run_grpo.py \
-  --config examples/modelopt/qa_grpo_llama8b_megatron.yaml \
+  --config examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml \
   policy.model_name=meta-llama/Llama-3.1-8B-Instruct
 ```
 
 **Via Slurm:**
 
 ```bash
 COMMAND="uv run examples/run_grpo.py \
-  --config examples/modelopt/qa_grpo_llama8b_megatron.yaml \
+  --config examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml \
   policy.model_name=meta-llama/Llama-3.1-8B-Instruct \
   checkpointing.checkpoint_dir=results/qa_grpo" \
 CONTAINER=YOUR_CONTAINER \

examples/configs/prorlv2_1_moe.yaml → examples/configs/prorlv2_1_moe.v2.yaml

@@ -17,7 +17,7 @@
 # Seq-Mask-TIS: https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda
 # Online IcePop: https://hijkzzz.notion.site/online-ice-pop
 
-defaults: "prorlv2.yaml"
+defaults: "prorlv2.v2.yaml"
 
 grpo:
   # ============================================================================

examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.yaml → examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.v2.yaml

@@ -1,4 +1,4 @@
-defaults: ./dapo-qwen2.5-7b.yaml
+defaults: ./dapo-qwen2.5-7b.v2.yaml
 policy:
   dtensor_cfg:
     context_parallel_size: 2

examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.yaml → examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.yaml

@@ -1,4 +1,4 @@
-defaults: ../../prorlv2.yaml
+defaults: ../../prorlv2.v2.yaml
 grpo:
   max_num_steps: 450
 checkpointing:

nemo_rl/algorithms/grpo.py

@@ -1616,6 +1616,17 @@ def grpo_train(
                     rewards = repeated_batch["total_reward"]
 
                     print("▶ Computing advantages...", flush=True)
+                    # For DAPO with reward shaping, compute std on the raw
+                    # pre-shaping reward so dynamic sampling filters prompt
+                    # groups on the raw task metric (e.g. acc) instead of on
+                    # length-dependent shaped reward variance. Baseline
+                    # (which drives advantages) stays on the shaped reward.
+                    std_rewards = (
+                        repeated_batch["unshaped_total_reward"]
+                        if master_config.grpo["use_dynamic_sampling"]
+                        and "unshaped_total_reward" in repeated_batch
+                        else None
+                    )
                     if master_config.grpo.get("calculate_advantages_on_gpu"):
                         print("Computing advantages on GPU!")
                         # Just fix the device id for now
@@ -1627,6 +1638,11 @@ def grpo_train(
                             leave_one_out_baseline=master_config.grpo[
                                 "use_leave_one_out_baseline"
                             ],
+                            std_rewards=(
+                                std_rewards.cuda(device_id)
+                                if std_rewards is not None
+                                else None
+                            ),
                         )
                         baseline = baseline.cpu()
                         std = std.cpu()
@@ -1638,6 +1654,7 @@ def grpo_train(
                             leave_one_out_baseline=master_config.grpo[
                                 "use_leave_one_out_baseline"
                             ],
+                            std_rewards=std_rewards,
                         )
 
                     # Apply dynamic sampling to filter prompts with non-zero std (DAPO algorithm)

nemo_rl/algorithms/reward_functions.py

@@ -61,6 +61,11 @@ def apply_reward_shaping(
     if not cfg["enabled"]:
         return batch
 
+    # Preserve the pre-shaping reward so downstream consumers (e.g. DAPO
+    # dynamic sampling) can filter prompt groups on the raw task metric
+    # rather than on length-dependent shaped rewards.
+    batch["unshaped_total_reward"] = rewards.clone()
+
     # Apply stop properly penalty if configured
     stop_properly_penalty_coef = cfg.get("stop_properly_penalty_coef", None)
     if stop_properly_penalty_coef is not None:

nemo_rl/algorithms/utils.py

@@ -88,6 +88,7 @@ def calculate_baseline_and_std_per_prompt(
     rewards: torch.Tensor,
     valid_mask: torch.Tensor,
     leave_one_out_baseline: bool = True,
+    std_rewards: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Function to compute a baseline for each (prompt, response) pair in the batch.
 
@@ -99,10 +100,17 @@ def calculate_baseline_and_std_per_prompt(
     valid_mask: tensor (b,)       Vector of 0/1, where 0 is to ignore and 1 is to keep
     leave_one_out_baseline: bool  Compute an unbiased baseline by leaving out the sample that
                                   the baseline is for (from RLOO https://arxiv.org/abs/2402.14740)
+    std_rewards: tensor (b,)      Optional separate reward tensor used only for the std
+                                  calculation. Defaults to `rewards`. Useful for DAPO,
+                                  which needs std on the raw task metric for dynamic
+                                  sampling filtering while keeping baseline on the
+                                  shaped reward.
 
     Returns:
     tensor (b,), tensor (b,) of baselines and std on the same device as 'rewards'
     """
+    if std_rewards is None:
+        std_rewards = rewards
     unique_prompts = torch.unique(prompts, dim=0)
 
     baseline = torch.zeros_like(rewards)
@@ -141,19 +149,28 @@ def calculate_baseline_and_std_per_prompt(
                 )
                 / num_valid
             )
-            prompt_baseline_square = (
+            std_prompt_baseline = (
+                prompt_baseline
+                if std_rewards is rewards
+                else torch.matmul(
+                    baseline_mask_matrix,
+                    std_rewards[prompt_idx] * valid_mask[prompt_idx],
+                )
+                / num_valid
+            )
+            std_prompt_baseline_square = (
                 torch.matmul(
                     baseline_mask_matrix,
-                    torch.pow(rewards[prompt_idx], 2) * valid_mask[prompt_idx],
+                    torch.pow(std_rewards[prompt_idx], 2) * valid_mask[prompt_idx],
                 )
                 / num_valid
             )
 
             baseline[prompt_idx] = prompt_baseline
-            sq_baseline[prompt_idx] = prompt_baseline_square
+            sq_baseline[prompt_idx] = std_prompt_baseline_square
             std[prompt_idx] = (
                 (
-                    (prompt_baseline_square - prompt_baseline.square())
+                    (std_prompt_baseline_square - std_prompt_baseline.square())
                     * (num_valid / (num_valid - 1))
                 )
                 .sqrt()

tests/functional/prorlv2.sh

@@ -20,7 +20,7 @@ mkdir -p $EXP_DIR $LOG_DIR
 cd $PROJECT_ROOT
 uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
     $PROJECT_ROOT/examples/run_grpo.py \
-    --config $PROJECT_ROOT/examples/configs/prorlv2.yaml \
+    --config $PROJECT_ROOT/examples/configs/prorlv2.v2.yaml \
     policy.model_name=Qwen/Qwen3-0.6B \
     policy.tokenizer.name=Qwen/Qwen3-0.6B \
     grpo.num_prompts_per_step=2 \

tests/test_suites/disabled.txt

@@ -4,4 +4,4 @@ tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.sh
 tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.sh
 
 # Disable this due to know vLLM bug with Qwen3.5: https://github.com/vllm-project/vllm/issues/37856
-tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.sh
+tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.sh

tests/test_suites/llm/dapo-qwen2.5-7b.sh → tests/test_suites/llm/dapo-qwen2.5-7b.v2.sh

@@ -36,8 +36,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma
     uv run tests/check_metrics.py $JSON_METRICS \
         'median(data["train/token_mult_prob_error"]) < 1.1' \
         'data["train/token_mult_prob_error"]["20"] < 1.05' \
-        'data["train/reward"]["20"] > -0.45' \
-        'data["train/filtered_reward"]["20"] > -0.2'
+        'data["train/reward"]["20"] > -0.40' \
+        'data["train/filtered_reward"]["20"] > -0.15'
 
     # Clean up checkpoint directory after successful run to save space.
     rm -rf "$CKPT_DIR"

tests/test_suites/nightly.txt

@@ -8,7 +8,7 @@ tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v3.sh
 tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1.sh
 
 # ProRLv2 convergence test
-tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.sh
+tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.sh
 
 # SGLang backend
 tests/test_suites/llm/grpo-qwen3-0.6b-1n8g-sglang.sh

tests/test_suites/performance.txt

@@ -13,7 +13,7 @@ tests/test_suites/llm/performance/grpo-deepseek-v3-64n8g.sh
 tests/test_suites/llm/performance/grpo-qwen3-32b-4n8g.sh
 tests/test_suites/llm/performance/grpo-qwen3-235b-16n8g.sh
 tests/test_suites/llm/performance/grpo-qwen3-235b-32n8g.sh
-tests/test_suites/llm/performance/dapo-deepseek-v3-64n8g.sh
+tests/test_suites/llm/performance/dapo-deepseek-v3-64n8g.v2.sh
 
 ## ASYNC 1-off
 tests/test_suites/llm/performance/grpo-deepseek-v3-64n8g-async-1off.sh

tests/test_suites/release.txt

@@ -18,7 +18,7 @@ tests/test_suites/llm/grpo-gemma3-27b-it-8n8g-fsdp2tp8-actckpt-long.sh
 tests/test_suites/llm/grpo-qwen3-30ba3b-8n8g-megatron.sh
 
 # DAPO 4h run
-tests/test_suites/llm/dapo-qwen2.5-7b.sh
+tests/test_suites/llm/dapo-qwen2.5-7b.v2.sh
 
 # Qwen3.5-35B DAPO GRPO run
 tests/test_suites/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.sh

tests/test_suites/release_gb200.txt

@@ -18,7 +18,7 @@ tests/test_suites/llm/grpo-gemma3-27b-it-8n4g-fsdp2tp4-actckpt-long.sh
 tests/test_suites/llm/grpo-qwen3-30ba3b-8n4g-megatron.sh
 
 # DAPO 4h run
-tests/test_suites/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.sh
+tests/test_suites/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.v2.sh
 
 # Deepseek-V3 on DAPO dataset
 tests/test_suites/llm/grpo-dapomath17k-dsv3-32n4g-megatron.sh