Merge branch 'save' into support_vllm_dp

Author: guyueh1

docs/fp8.md

@@ -53,8 +53,8 @@ FP8 generations are recommended to be configured with the following settings:
                 num_first_layers_in_bf16: 0
                 # Use FP32 scaling factors. Rounding scaling factors to the nearest pow2 may improve quantization 
                 # fidelity however this feature is still under research.
-                use_weight_pow2_scale: False
-                use_activation_pow2_scale: False
+                pow2_weight_scaling_factors: False
+                pow2_activation_scaling_factors: False
 ```
 
 To train with FP8, you need to set the Megatron path and configure it using the following settings:

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/distillation_math.yaml

@@ -90,6 +90,17 @@ policy: &POLICY_BASE
         force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
         empty_unused_memory_level: 0
         activation_checkpointing: false
+        # recompute_granularity controls activation checkpointing depth.
+        # "full": recompute all activations (default, max memory savings).
+        # "selective": recompute only specific modules (see recompute_modules).
+        # Only takes effect when activation_checkpointing: true.
+        recompute_granularity: "full"
+        # Modules to selectively recompute when recompute_granularity="selective".
+        # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+        # Null defaults to ["core_attn"]. Full list and per-module constraints:
+        # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+        # Use ["moe"] for MoE models to recompute expert activations only.
+        recompute_modules: null
         converter_type: "Qwen3ForCausalLM"
         tensor_model_parallel_size: 2
         expert_tensor_parallel_size: 1

examples/configs/distillation_math_megatron.yaml

@@ -39,6 +39,17 @@ policy: &POLICY_BASE
         enabled: true
         empty_unused_memory_level: 0
         activation_checkpointing: false
+        # recompute_granularity controls activation checkpointing depth.
+        # "full": recompute all activations (default, max memory savings).
+        # "selective": recompute only specific modules (see recompute_modules).
+        # Only takes effect when activation_checkpointing: true.
+        recompute_granularity: "full"
+        # Modules to selectively recompute when recompute_granularity="selective".
+        # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+        # Null defaults to ["core_attn"]. Full list and per-module constraints:
+        # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+        # Use ["moe"] for MoE models to recompute expert activations only.
+        recompute_modules: null
         converter_type: "Qwen3ForCausalLM"
         tensor_model_parallel_size: 2
         expert_tensor_parallel_size: 1

examples/configs/dpo.yaml

@@ -129,6 +129,17 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     tensor_model_parallel_size: 2
     expert_tensor_parallel_size: 1
     expert_model_parallel_size: 1

examples/configs/grpo_math_1B.yaml

@@ -134,6 +134,17 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1

examples/configs/grpo_math_1B_megatron.yaml

@@ -86,6 +86,17 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1