diff --git a/docs/guides/dapo.md b/docs/guides/dapo.md index 6d65b33ac63..4b50a6db205 100644 --- a/docs/guides/dapo.md +++ b/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) diff --git a/docs/guides/muon-optimizer.md b/docs/guides/muon-optimizer.md index e3098f9d4be..0dd6864d415 100644 --- a/docs/guides/muon-optimizer.md +++ b/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 \ diff --git a/docs/guides/prorlv2.md b/docs/guides/prorlv2.md index 3a9df63c511..f14fd818d37 100644 --- a/docs/guides/prorlv2.md +++ b/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,18 +204,18 @@ 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 @@ -223,7 +223,7 @@ 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 \ diff --git a/docs/guides/quantization-aware-rl.md b/docs/guides/quantization-aware-rl.md index 541f64cee5f..8b9799eb940 100644 --- a/docs/guides/quantization-aware-rl.md +++ b/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,7 +47,7 @@ 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 ``` @@ -55,7 +55,7 @@ uv run examples/run_grpo.py \ ```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 \ diff --git a/examples/configs/prorlv2.yaml b/examples/configs/prorlv2.v2.yaml similarity index 100% rename from examples/configs/prorlv2.yaml rename to examples/configs/prorlv2.v2.yaml diff --git a/examples/configs/prorlv2_1_moe.yaml b/examples/configs/prorlv2_1_moe.v2.yaml similarity index 98% rename from examples/configs/prorlv2_1_moe.yaml rename to examples/configs/prorlv2_1_moe.v2.yaml index b3db0a1df5b..b88f8d658f0 100644 --- a/examples/configs/prorlv2_1_moe.yaml +++ b/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: # ============================================================================ diff --git a/examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.yaml b/examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.v2.yaml similarity index 87% rename from examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.yaml rename to examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.v2.yaml index 6135f42da79..11539a6dd2f 100644 --- a/examples/configs/recipes/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.yaml +++ b/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 diff --git a/examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml b/examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml similarity index 100% rename from examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml rename to examples/configs/recipes/llm/dapo-qwen2.5-7b.v2.yaml diff --git a/examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.yaml b/examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.yaml similarity index 100% rename from examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.yaml rename to examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.yaml diff --git a/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml b/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.v2.yaml similarity index 100% rename from examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml rename to examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.v2.yaml diff --git a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.yaml b/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.yaml similarity index 95% rename from examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.yaml rename to examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.yaml index 89d364c521a..90fc26dd7a9 100644 --- a/examples/configs/recipes/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.yaml +++ b/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: diff --git a/examples/modelopt/qa_grpo_llama8b_megatron.yaml b/examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml similarity index 100% rename from examples/modelopt/qa_grpo_llama8b_megatron.yaml rename to examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d26fb0bcae4..8eff9478c13 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/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) diff --git a/nemo_rl/algorithms/reward_functions.py b/nemo_rl/algorithms/reward_functions.py index 87c826db26f..974bebc6392 100644 --- a/nemo_rl/algorithms/reward_functions.py +++ b/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: diff --git a/nemo_rl/algorithms/utils.py b/nemo_rl/algorithms/utils.py index 8ca463bb2c1..f145081aa99 100644 --- a/nemo_rl/algorithms/utils.py +++ b/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() diff --git a/tests/functional/prorlv2.sh b/tests/functional/prorlv2.sh index c39cff1da7f..9a2fe114201 100755 --- a/tests/functional/prorlv2.sh +++ b/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 \ diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 1730e4ed54d..d5ee913e6d9 100644 --- a/tests/test_suites/disabled.txt +++ b/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 \ No newline at end of file +tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.sh diff --git a/tests/test_suites/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.sh b/tests/test_suites/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.v2.sh similarity index 100% rename from tests/test_suites/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.sh rename to tests/test_suites/llm/dapo-qwen2.5-7b-16n4g-fsdp2cp2.v2.sh diff --git a/tests/test_suites/llm/dapo-qwen2.5-7b.sh b/tests/test_suites/llm/dapo-qwen2.5-7b.v2.sh similarity index 93% rename from tests/test_suites/llm/dapo-qwen2.5-7b.sh rename to tests/test_suites/llm/dapo-qwen2.5-7b.v2.sh index 721a2aabfff..2bbe7baf59e 100755 --- a/tests/test_suites/llm/dapo-qwen2.5-7b.sh +++ b/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" diff --git a/tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.sh b/tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.sh similarity index 100% rename from tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.sh rename to tests/test_suites/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.sh diff --git a/tests/test_suites/llm/performance/dapo-deepseek-v3-64n8g.sh b/tests/test_suites/llm/performance/dapo-deepseek-v3-64n8g.v2.sh similarity index 100% rename from tests/test_suites/llm/performance/dapo-deepseek-v3-64n8g.sh rename to tests/test_suites/llm/performance/dapo-deepseek-v3-64n8g.v2.sh diff --git a/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.sh b/tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.sh similarity index 100% rename from tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.sh rename to tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2.sh diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 0ceb4b13fdd..3615ab14a13 100644 --- a/tests/test_suites/nightly.txt +++ b/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 diff --git a/tests/test_suites/performance.txt b/tests/test_suites/performance.txt index c7ac5102771..cba2c73d491 100644 --- a/tests/test_suites/performance.txt +++ b/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 diff --git a/tests/test_suites/release.txt b/tests/test_suites/release.txt index 96837012bf2..aa47eab881f 100644 --- a/tests/test_suites/release.txt +++ b/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 diff --git a/tests/test_suites/release_gb200.txt b/tests/test_suites/release_gb200.txt index 59cee6f24c2..f0c346509db 100644 --- a/tests/test_suites/release_gb200.txt +++ b/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 diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index a6193f9f854..2510f0cd55a 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -35,6 +35,11 @@ validate, ) from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn +from nemo_rl.algorithms.reward_functions import ( + RewardShapingConfig, + apply_reward_shaping, +) +from nemo_rl.algorithms.utils import calculate_baseline_and_std_per_prompt from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( @@ -936,6 +941,99 @@ def test_dapo_dynamic_sampling_disabled(): assert batch_cache is None # No caching when disabled +def test_dapo_dynamic_sampling_filters_on_raw_metric_after_overlong_shaping(): + """Regression test for the issue where DAPO dynamic sampling filtered on + shaped reward std instead of the raw task metric. + + The first prompt group: all responses are raw-wrong (acc=0) but lengths + differ, so the overlong penalty produces non-zero shaped std. The fix + must recompute std on the raw (pre-shaping) reward so this group is + filtered out. The second prompt group has genuinely varied raw rewards + and must be kept. + """ + batch_size = 6 + # Vary the assistant response length for the first group so the overlong + # penalty creates spurious shaped-reward variance. + response_lengths = [10, 22, 30, 10, 20, 10] + message_logs = [] + for i, length in enumerate(response_lengths): + message_logs.append( + [ + { + "role": "user", + "content": f"prompt_{i // 3}", + "token_ids": torch.tensor([100 + (i // 3), 101, 102]), + }, + { + "role": "assistant", + "content": f"response_{i}", + "token_ids": torch.arange(length, dtype=torch.long), + }, + ] + ) + task_names = ["math"] * batch_size + repeated_batch = create_mock_batch(batch_size, task_names, message_logs) + # Group 0: all raw-wrong (acc=0). Group 1: mixed. + repeated_batch["total_reward"] = torch.tensor([0.0, 0.0, 0.0, 1.0, 0.0, 1.0]) + + shaping_cfg = RewardShapingConfig( + enabled=True, + overlong_buffer_length=5, + overlong_buffer_penalty=0.5, + max_response_length=25, + ) + repeated_batch = apply_reward_shaping(repeated_batch, shaping_cfg) + + # Shaped reward of group 0 is now [0.0, -0.2, -1.0] -> non-zero std. + shaped_std = repeated_batch["total_reward"][:3].std(unbiased=False) + assert shaped_std.item() > 0.0 + + # Use prompt-only token_ids to identify groups, mirroring the grpo.py + # call site. + input_ids = torch.stack([m[0]["token_ids"] for m in repeated_batch["message_log"]]) + rewards = repeated_batch["total_reward"] + baseline, raw_std = calculate_baseline_and_std_per_prompt( + input_ids, + rewards, + torch.ones_like(rewards), + leave_one_out_baseline=False, + std_rewards=repeated_batch["unshaped_total_reward"], + ) + + # Raw std is 0 for the homogeneous group, non-zero for the mixed group. + assert torch.allclose(raw_std[:3], torch.zeros(3)) + assert (raw_std[3:] > 0).all() + + master_config = MasterConfig.model_construct( + **{ + "grpo": { + "use_dynamic_sampling": True, + "num_prompts_per_step": 1, + "num_generations_per_prompt": 3, + "dynamic_sampling_max_gen_batches": 5, + } + } + ) + + result_batch, is_batch_complete, _, _ = dynamic_sampling( + repeated_batch, + raw_std, + baseline, + dynamic_sampling_num_gen_batches=1, + master_config=master_config, + timer=Timer(), + ) + + # Only the second group should survive — the first group's raw rewards are + # identical, so filtering on the raw metric drops it. + assert is_batch_complete is True + assert result_batch.size == 3 + surviving_prompts = [ + result_batch["message_log"][i][0]["content"] for i in range(result_batch.size) + ] + assert surviving_prompts == ["prompt_1", "prompt_1", "prompt_1"] + + def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node(): """Test that non-colocated inference requires explicit gpus_per_node when policy_nodes=1.""" from unittest.mock import MagicMock, patch diff --git a/tests/unit/algorithms/test_reward_functions.py b/tests/unit/algorithms/test_reward_functions.py index 2f72b300cb4..01334a56b58 100755 --- a/tests/unit/algorithms/test_reward_functions.py +++ b/tests/unit/algorithms/test_reward_functions.py @@ -186,6 +186,80 @@ def test_reward_shaping_with_penalties(): assert torch.allclose(result_batch["total_reward"], expected_rewards, atol=1e-6) +def test_reward_shaping_preserves_unshaped_reward_overlong(): + """Reward shaping must save the pre-shaping reward so dynamic sampling can + filter prompt groups on the raw task metric, not on shaped reward whose + std is corrupted by length-dependent overlong penalties. + """ + raw_rewards = [0.0, 0.0, 0.0, 0.0] + batch = create_mock_batch_with_responses( + num_samples=4, + response_lengths=[10, 22, 25, 30], + initial_rewards=raw_rewards, + ) + + config = RewardShapingConfig( + enabled=True, + overlong_buffer_length=5, + overlong_buffer_penalty=0.5, + max_response_length=25, + ) + + result_batch = apply_reward_shaping(batch, config) + + # Shaped rewards differ across the group due to length-based penalty even + # though all raw rewards are 0 (i.e. all responses are wrong). + expected_shaped = torch.tensor([0.0, -0.2, -0.5, -1.0]) + assert torch.allclose(result_batch["total_reward"], expected_shaped, atol=1e-6) + + # The unshaped reward must be retained verbatim. + assert "unshaped_total_reward" in result_batch + assert torch.allclose( + result_batch["unshaped_total_reward"], torch.tensor(raw_rewards) + ) + # The two tensors must be independent — mutating one must not affect the other. + assert ( + result_batch["unshaped_total_reward"].data_ptr() + != result_batch["total_reward"].data_ptr() + ) + + +def test_reward_shaping_preserves_unshaped_reward_stop_properly(): + """The stop-properly penalty path must also preserve the raw reward.""" + raw_rewards = [1.0, 0.8, 0.6, 0.4] + batch = create_mock_batch_with_responses( + num_samples=4, + response_lengths=[10, 20, 30, 40], + initial_rewards=raw_rewards, + ) + batch["truncated"] = torch.tensor([False, True, False, True]) + + config = RewardShapingConfig(enabled=True, stop_properly_penalty_coef=0.5) + result_batch = apply_reward_shaping(batch, config) + + assert "unshaped_total_reward" in result_batch + assert torch.allclose( + result_batch["unshaped_total_reward"], torch.tensor(raw_rewards) + ) + + +def test_reward_shaping_disabled_does_not_save_unshaped_reward(): + """When shaping is disabled, the unshaped_total_reward field should not be added.""" + batch = create_mock_batch_with_responses( + num_samples=3, response_lengths=[10, 20, 30], initial_rewards=[1.0, 0.5, 0.8] + ) + + config = RewardShapingConfig( + enabled=False, + overlong_buffer_length=5, + overlong_buffer_penalty=0.1, + max_response_length=25, + ) + + result_batch = apply_reward_shaping(batch, config) + assert "unshaped_total_reward" not in result_batch + + def test_reward_shaping_missing_config_values(): """Test that missing required config values raise ValueError.""" batch = create_mock_batch_with_responses( diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index c90e6d3b115..760ce027b51 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -45,7 +45,7 @@ "distillation": "examples/configs/distillation_math.yaml", "rm": "examples/configs/rm.yaml", "dapo": "examples/configs/grpo_math_1B.yaml", - "prorlv2": "examples/configs/prorlv2.yaml", + "prorlv2": "examples/configs/prorlv2.v2.yaml", } # Configuration keys that are allowed to be added to base configs during testing