From de1095a90f455445c009d3fe5d2cc214b36727b0 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Mon, 19 Oct 2026 10:36:05 -0700 Subject: [PATCH 01/26] pertoken nvfp4 zero kl recipe Signed-off-by: YigongQin --- ...noncolocated-zero-kl-mxfp8-flashinfer.yaml | 26 +++++ ...noncolocated-zero-kl-nvfp4-flashinfer.yaml | 24 +++++ nemo_rl/models/generation/megatron/config.py | 4 + nemo_rl/models/megatron/setup.py | 41 +++++++ .../megatron/zero_train_gen_mismatch.py | 87 +++++++++++++-- nemo_rl/models/policy/__init__.py | 15 +++ tests/test_suites/disabled.txt | 6 ++ ...n-noncolocated-zero-kl-mxfp8-flashinfer.sh | 63 +++++++++++ ...n-noncolocated-zero-kl-nvfp4-flashinfer.sh | 60 +++++++++++ .../models/megatron/test_megatron_setup.py | 76 +++++++++++++ .../megatron/test_zero_train_gen_mismatch.py | 101 ++++++++++++++++++ 11 files changed, 495 insertions(+), 8 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml new file mode 100644 index 00000000000..1a6d79f8afd --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml @@ -0,0 +1,26 @@ +# Nemotron Nano v3 zero train/generation KL with MXFP8 routed experts on generation's +# FlashInfer routed MXFP8 kernel. The FlashInfer kernel does not match TE training, so +# each training MoE forward runs that kernel (moe_inference_training_forward) and the +# TE recompute pass of the MoE layer supplies the backward. +# +# Needs Megatron-LM PR #7798 on top of the zero-kl-mxfp8 parent's requirements. +defaults: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer +policy: + megatron_cfg: + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls + moe_inference_training_forward: true + activation_checkpointing: true + recompute_granularity: selective + recompute_modules: + - moe + generation: + mcore_generation_config: + inference_grouped_gemm_backend: flashinfer +logger: + log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer + wandb: + name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml new file mode 100644 index 00000000000..886350aab86 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -0,0 +1,24 @@ +# Nemotron Nano v3 zero train/generation KL with per-token NVFP4 routed experts on +# FlashInfer. Training and refit stay BF16; generation quantizes the experts to NVFP4, +# and each training MoE forward runs the same kernel on the same quantization +# (moe_inference_training_forward) with the backward from the TE recompute pass. +# +# Needs Megatron-LM PR #7798 on top of the zero-kl-mxfp8 parent's requirements, and +# flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 MoE kernels). +defaults: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer +policy: + megatron_cfg: + inference_moe_nvfp4_recipe: per_token + te_precision_config_file: null + fp8_cfg: + enabled: false + generation: + mcore_generation_config: + fp8_cfg: + enabled: false +logger: + log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + wandb: + name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index e2706688a51..5d6ef696362 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -214,6 +214,10 @@ def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any] "activation_checkpointing": False, "context_parallel_size": 1, } + # Training-only: generation already runs the inference kernel, and MCore + # rejects the mode without the MoE recompute that generation disables. + if merged.get("moe_inference_training_forward"): + merged["moe_inference_training_forward"] = False # inference_optimized layers hard-require SP with TP>1. Raise with the # config key: the colocated build bypasses validate_and_set_config, so this # merge is the only spot the inference cfg gets a named error instead of a diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 765ce858768..ed31cb6facb 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1345,6 +1345,35 @@ def _apply_multimodal_config(model_cfg: Any, config: PolicyConfig) -> None: setattr(model_cfg, provider_attr, megatron_cfg[config_key]) +def moe_inference_training_max_tokens_per_rank(config: PolicyConfig) -> int: + """Tokens one rank feeds the largest moe_inference_training_forward pass. + + The NVLS symmetric heap is allocated once and cannot grow, so it is sized for + the largest train or log-prob microbatch at the padded maximum sequence length. + """ + configured = config["megatron_cfg"].get( + "moe_inference_training_max_tokens_per_rank" + ) + if configured is not None: + return configured + divisor = config["make_sequence_length_divisible_by"] + padded_seq_len = -(-config["max_total_sequence_length"] // divisor) * divisor + micro_batch = max( + config["train_micro_batch_size"], + config.get("logprob_batch_size", config["train_micro_batch_size"]), + ) + tokens = micro_batch * padded_seq_len + for batching_key in ("sequence_packing", "dynamic_batching"): + batching = config.get(batching_key) or {} + if batching.get("enabled"): + tokens = max( + tokens, + batching["train_mb_tokens"], + batching.get("logprob_mb_tokens", batching["train_mb_tokens"]), + ) + return tokens + + def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: """Apply Mixture of Experts configuration.""" model_cfg.expert_tensor_parallel_size = config["megatron_cfg"][ @@ -1389,6 +1418,18 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.inference_flashinfer_mxfp8_token_capacity = config["megatron_cfg"][ "inference_flashinfer_mxfp8_token_capacity" ] + if "inference_moe_nvfp4_recipe" in config["megatron_cfg"]: + model_cfg.inference_moe_nvfp4_recipe = config["megatron_cfg"][ + "inference_moe_nvfp4_recipe" + ] + if "moe_inference_training_forward" in config["megatron_cfg"]: + model_cfg.moe_inference_training_forward = config["megatron_cfg"][ + "moe_inference_training_forward" + ] + if model_cfg.moe_inference_training_forward: + model_cfg.moe_inference_training_max_tokens_per_rank = ( + moe_inference_training_max_tokens_per_rank(config) + ) if "moe_router_num_groups" in config["megatron_cfg"]: model_cfg.moe_router_num_groups = config["megatron_cfg"][ "moe_router_num_groups" diff --git a/nemo_rl/models/megatron/zero_train_gen_mismatch.py b/nemo_rl/models/megatron/zero_train_gen_mismatch.py index 2e99e4cfdf7..ea862f06f00 100644 --- a/nemo_rl/models/megatron/zero_train_gen_mismatch.py +++ b/nemo_rl/models/megatron/zero_train_gen_mismatch.py @@ -15,7 +15,8 @@ """Zero train/generation KL (batch-invariant Megatron): validate the recipe, then enable. The recipe sets every required value; this module only checks them. Supported -precisions are BF16 and MXFP8 (the same on train and generation). +precisions are BF16 and MXFP8 (the same on train and generation), and FlashInfer +MXFP8 or per-token NVFP4 experts with moe_inference_training_forward. """ from __future__ import annotations @@ -487,19 +488,89 @@ def _validate_backend(config: PolicyConfig, out: ZeroTrainGenValidation) -> None ) fp8_cfg = inference_cfg.get("fp8_cfg") or {} - backend = inference_cfg.get("inference_grouped_gemm_backend") - backend = getattr(backend, "value", backend) + backend = _grouped_gemm_backend(inference_cfg) + quantized_experts = ( + fp8_cfg.get("enabled") and fp8_cfg.get("fp8_recipe") == "mxfp8" + ) or inference_cfg.get("inference_moe_nvfp4_recipe") is not None if ( impl == "inference_optimized" and backend == "flashinfer" - and fp8_cfg.get("enabled") - and fp8_cfg.get("fp8_recipe") == "mxfp8" + and quantized_experts + and not config["megatron_cfg"].get("moe_inference_training_forward") ): out.violations.append( - "zero_train_gen_mismatch does not support " - "inference_grouped_gemm_backend='flashinfer' with MXFP8: the FlashInfer " + "zero_train_gen_mismatch with inference_grouped_gemm_backend='flashinfer' " + "and MXFP8 or NVFP4 experts requires " + "policy.megatron_cfg.moe_inference_training_forward=true: the FlashInfer " "kernel is batch-invariant but not bitwise identical to TE training. " - "Use 'torch' or 'vllm'." + "Otherwise use 'torch' or 'vllm' with MXFP8." + ) + if config["megatron_cfg"].get("moe_inference_training_forward"): + _validate_inference_training_forward(config, inference_cfg, out) + + +def _grouped_gemm_backend(cfg: dict[str, Any]) -> Any: + backend = cfg.get("inference_grouped_gemm_backend") + return getattr(backend, "value", backend) + + +def _validate_inference_training_forward( + config: PolicyConfig, inference_cfg: dict[str, Any], out: ZeroTrainGenValidation +) -> None: + """The training MoE forward runs generation's FlashInfer kernel; TE recompute backs it.""" + megatron_cfg = config["megatron_cfg"] + prefix = "moe_inference_training_forward=true requires" + if megatron_cfg.get("transformer_impl") != "inference_optimized": + out.violations.append( + f"{prefix} policy.megatron_cfg.transformer_impl='inference_optimized' " + f"(got {megatron_cfg.get('transformer_impl')!r})." + ) + for label, cfg in ( + ("policy.megatron_cfg", megatron_cfg), + ("policy.generation.mcore_generation_config", inference_cfg), + ): + if _grouped_gemm_backend(cfg) != "flashinfer": + out.violations.append( + f"{prefix} {label}.inference_grouped_gemm_backend='flashinfer' " + f"(got {_grouped_gemm_backend(cfg)!r})." + ) + fp8_cfg = megatron_cfg.get("fp8_cfg") or {} + mxfp8 = fp8_cfg.get("enabled") and fp8_cfg.get("fp8_recipe") == "mxfp8" + if not mxfp8 and megatron_cfg.get("inference_moe_nvfp4_recipe") is None: + out.violations.append( + f"{prefix} MXFP8 experts (policy.megatron_cfg.fp8_cfg with " + "fp8_recipe='mxfp8') or policy.megatron_cfg.inference_moe_nvfp4_recipe: " + "the FlashInfer BF16 kernel is not batch invariant." + ) + recompute_modules = megatron_cfg.get("recompute_modules") or [] + if not ( + megatron_cfg.get("activation_checkpointing") + and megatron_cfg.get("recompute_granularity") == "selective" + and "moe" in recompute_modules + ): + out.violations.append( + f"{prefix} activation_checkpointing=true, " + "recompute_granularity='selective' and 'moe' in recompute_modules: the " + "backward comes from the TE recompute pass." + ) + if megatron_cfg.get("cuda_graph_impl") == "local": + out.violations.append( + f"{prefix} a training cuda_graph_impl other than 'local', which disables " + "MoE-layer recompute." + ) + mismatched = [ + key + for key in ( + "inference_moe_nvfp4_recipe", + "inference_moe_token_dispatcher_type", + "expert_model_parallel_size", + ) + if megatron_cfg.get(key) != inference_cfg.get(key) + ] + if mismatched: + out.violations.append( + f"{prefix} the same training and generation values for: " + f"{', '.join(mismatched)}." ) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 45846aab66c..fbaf1f85d1e 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -492,6 +492,21 @@ class MegatronConfig(TypedDict): # rows for the local expert count (Nano-30B: >=256 rows at 16 local experts, # >=1024 at 64); smaller capacities fail CUDA-graph capture. inference_flashinfer_mxfp8_token_capacity: NotRequired[int | None] + # Routed experts as NVFP4 weights x per-token NVFP4 activations on FlashInfer + # (MCore inference_moe_nvfp4_recipe). Only 'per_token'; None keeps BF16 experts. + # Requires inference_grouped_gemm_backend='flashinfer' and BF16 (no fp8_cfg). + inference_moe_nvfp4_recipe: NotRequired[Literal["per_token"] | None] + # Training forward of each MoE layer runs generation's FlashInfer MXFP8 or + # per-token NVFP4 expert kernel; the backward comes from the TE recompute pass + # (straight-through estimator). Requires transformer_impl='inference_optimized', + # inference_grouped_gemm_backend='flashinfer', batch_invariant_mode, and + # activation_checkpointing with recompute_granularity='selective' and 'moe' in + # recompute_modules. Generation never sets it. + moe_inference_training_forward: NotRequired[bool] + # Token bound per rank for the NVLS symmetric heap of that training forward. + # Unset: max(train_micro_batch_size, logprob_batch_size) * padded + # max_total_sequence_length, or the packing / dynamic batching mb_tokens if larger. + moe_inference_training_max_tokens_per_rank: NotRequired[int | None] # InferenceTopKRouter requires moe_router_num_groups=None # (used when transformer_impl='inference_optimized') moe_router_num_groups: NotRequired[int | None] diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 2d43c913cc5..dc09a4c6989 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -102,3 +102,9 @@ research/template_project/tests/test_suites/llm/single_update_1n8g.sh # Full Flow-GRPO recipe requires an externally served GenRM OCR judge. research/flow_grpo/tests/test_suites/flow_grpo_qwen_image_ocr.sh + +# FlashInfer MXFP8 / per-token NVFP4 zero-KL via moe_inference_training_forward: needs +# Megatron-LM PR #7798 in the submodule pin (NVFP4 also flashinfer-python>=0.7.0) and +# has not run end to end yet. +tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh new file mode 100755 index 00000000000..ac0bbe943fd --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with MXFP8 +# routed experts on generation's FlashInfer routed MXFP8 kernel. Each training MoE forward +# runs that kernel (moe_inference_training_forward); the TE recompute pass supplies the +# backward. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=4 +STEPS_PER_RUN=5 +MAX_STEPS=5 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run --no-sync examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# The recipe must have switched the batch-invariant kernels on. +grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG + +# The per-module MXFP8 recipe must have been loaded. +grep -F -q "TE per-module precision recipe loaded from" $RUN_LOG + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +# Zero train/generation KL means bitwise agreement: gen_kl_error is exactly 0 and +# token_mult_prob_error (max multiplicative probability error) is exactly 1 on every +# step. Both min and max are checked so a negative KL estimate cannot slip through. +uv run --no-sync tests/check_metrics.py $JSON_METRICS \ + 'min(data["train/num_valid_samples"]) > 0' \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/gen_kl_error"]) == 0' \ + 'max(data["train/gen_kl_error"]) == 0' \ + 'min(data["train/token_mult_prob_error"]) == 1' \ + 'max(data["train/token_mult_prob_error"]) == 1' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh new file mode 100755 index 00000000000..947547652cc --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with +# per-token NVFP4 routed experts on FlashInfer. Training and refit stay BF16; each +# training MoE forward runs generation's NVFP4 kernel (moe_inference_training_forward) +# and the TE recompute pass supplies the backward. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=4 +STEPS_PER_RUN=5 +MAX_STEPS=5 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run --no-sync examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# The recipe must have switched the batch-invariant kernels on. +grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +# Zero train/generation KL means bitwise agreement: gen_kl_error is exactly 0 and +# token_mult_prob_error (max multiplicative probability error) is exactly 1 on every +# step. Both min and max are checked so a negative KL estimate cannot slip through. +uv run --no-sync tests/check_metrics.py $JSON_METRICS \ + 'min(data["train/num_valid_samples"]) > 0' \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/gen_kl_error"]) == 0' \ + 'max(data["train/gen_kl_error"]) == 0' \ + 'min(data["train/token_mult_prob_error"]) == 1' \ + 'max(data["train/token_mult_prob_error"]) == 1' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 066065f6cbe..16aeab5c0cc 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -994,6 +994,8 @@ def test_moe_grouped_gemm_absent_keeps_default(self): assert not hasattr(model_cfg, "moe_grouped_gemm") assert not hasattr(model_cfg, "inference_flashinfer_mxfp8_token_capacity") + assert not hasattr(model_cfg, "inference_moe_nvfp4_recipe") + assert not hasattr(model_cfg, "moe_inference_training_forward") def test_flashinfer_mxfp8_token_capacity_forwarded(self): """Without this key MCore runs FlashInfer MXFP8 MoE over the full dispatcher buffer.""" @@ -1040,6 +1042,80 @@ def test_inference_optimized_moe_requires_fp32_router( _apply_moe_config(model_cfg, config) assert model_cfg.moe_router_dtype == moe_router_dtype + def test_inference_training_forward_forwarded_with_token_bound(self): + """The NVLS heap bound covers the larger of the train and log-prob microbatches.""" + from nemo_rl.models.megatron.setup import _apply_moe_config + + model_cfg = MagicMock() + megatron_cfg = self._base_moe_megatron_cfg() + megatron_cfg["inference_moe_nvfp4_recipe"] = "per_token" + megatron_cfg["moe_inference_training_forward"] = True + config = { + "megatron_cfg": megatron_cfg, + "train_micro_batch_size": 1, + "logprob_batch_size": 4, + "max_total_sequence_length": 1000, + "make_sequence_length_divisible_by": 64, + } + + _apply_moe_config(model_cfg, config) + + assert model_cfg.inference_moe_nvfp4_recipe == "per_token" + assert model_cfg.moe_inference_training_forward is True + assert model_cfg.moe_inference_training_max_tokens_per_rank == 4 * 1024 + + @pytest.mark.parametrize("batching_key", ["sequence_packing", "dynamic_batching"]) + def test_inference_training_forward_token_bound_covers_packed_microbatches( + self, batching_key + ): + from nemo_rl.models.megatron.setup import ( + moe_inference_training_max_tokens_per_rank, + ) + + config = { + "megatron_cfg": {}, + "train_micro_batch_size": 1, + "logprob_batch_size": 1, + "max_total_sequence_length": 1000, + "make_sequence_length_divisible_by": 64, + batching_key: { + "enabled": True, + "train_mb_tokens": 4096, + "logprob_mb_tokens": 8192, + }, + } + + assert moe_inference_training_max_tokens_per_rank(config) == 8192 + + def test_inference_training_forward_keeps_an_explicit_token_bound(self): + from nemo_rl.models.megatron.setup import _apply_moe_config + + model_cfg = MagicMock() + megatron_cfg = self._base_moe_megatron_cfg() + megatron_cfg["moe_inference_training_forward"] = True + megatron_cfg["moe_inference_training_max_tokens_per_rank"] = 300 + config = {"megatron_cfg": megatron_cfg} + + _apply_moe_config(model_cfg, config) + + assert model_cfg.moe_inference_training_max_tokens_per_rank == 300 + + def test_inference_training_forward_false_resets_a_copied_provider(self): + """Generation can start from the training provider and must clear the flag.""" + from nemo_rl.models.megatron.setup import _apply_moe_config + + model_cfg = MagicMock() + model_cfg.moe_inference_training_forward = True + model_cfg.moe_inference_training_max_tokens_per_rank = None + megatron_cfg = self._base_moe_megatron_cfg() + megatron_cfg["moe_inference_training_forward"] = False + config = {"megatron_cfg": megatron_cfg} + + _apply_moe_config(model_cfg, config) + + assert model_cfg.moe_inference_training_forward is False + assert model_cfg.moe_inference_training_max_tokens_per_rank is None + def test_hybridep_input_prepadding_wins_after_bridge_validation(self): from nemo_rl.models.megatron import setup diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 549e7aa08b3..0a7048b2ee8 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -40,6 +40,8 @@ ZERO_KL_EXEMPLARS = [ "recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.yaml", "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml", + "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml", + "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml", ] # Installed-package versions that satisfy every gate in `_validate_packages`. @@ -501,6 +503,105 @@ def test_validate_rejects_flashinfer_grouped_gemm_with_mxfp8(): assert any("flashinfer" in v for v in _validate(config).violations) +def _inference_training_forward_config(precision: str) -> dict[str, Any]: + """FlashInfer experts on both sides, with the training forward running that kernel.""" + config = _policy_config() + config["megatron_cfg"].update( + transformer_impl="inference_optimized", + inference_grouped_gemm_backend="flashinfer", + inference_moe_token_dispatcher_type="nvls", + activation_checkpointing=True, + recompute_granularity="selective", + recompute_modules=["moe"], + moe_inference_training_forward=True, + ) + if precision == "mxfp8": + config["megatron_cfg"]["fp8_cfg"] = dict(_MXFP8) + else: + config["megatron_cfg"]["inference_moe_nvfp4_recipe"] = "per_token" + return config + + +@pytest.mark.parametrize("precision", ["mxfp8", "nvfp4"]) +def test_validate_accepts_flashinfer_with_inference_training_forward(precision): + assert _validate(_inference_training_forward_config(precision)).violations == [] + + +def test_validate_rejects_flashinfer_nvfp4_without_inference_training_forward(): + config = _inference_training_forward_config("nvfp4") + config["megatron_cfg"]["moe_inference_training_forward"] = False + assert any( + "requires policy.megatron_cfg.moe_inference_training_forward=true" in v + for v in _validate(config).violations + ) + + +def _set_generation(config: dict[str, Any], key: str, value: Any) -> None: + config["generation"]["mcore_generation_config"][key] = value + + +@pytest.mark.parametrize( + ("mutate", "expected"), + [ + ( + lambda c: c["megatron_cfg"].update(transformer_impl="transformer_engine"), + "transformer_impl='inference_optimized'", + ), + ( + lambda c: _set_generation(c, "inference_grouped_gemm_backend", "torch"), + "mcore_generation_config.inference_grouped_gemm_backend='flashinfer'", + ), + ( + lambda c: c["megatron_cfg"].update(activation_checkpointing=False), + "recompute_granularity='selective'", + ), + ( + lambda c: c["megatron_cfg"].update(recompute_granularity="full"), + "recompute_granularity='selective'", + ), + ( + lambda c: c["megatron_cfg"].update(recompute_modules=["core_attn"]), + "'moe' in recompute_modules", + ), + ( + lambda c: c["megatron_cfg"].update(cuda_graph_impl="local"), + "cuda_graph_impl other than 'local'", + ), + ( + lambda c: _set_generation(c, "inference_moe_token_dispatcher_type", "nccl"), + "inference_moe_token_dispatcher_type", + ), + ( + lambda c: _set_generation(c, "inference_moe_nvfp4_recipe", None), + "inference_moe_nvfp4_recipe", + ), + ( + lambda c: _set_generation(c, "expert_model_parallel_size", 2), + "expert_model_parallel_size", + ), + ( + lambda c: c["megatron_cfg"].pop("inference_moe_nvfp4_recipe"), + "FlashInfer BF16 kernel is not batch invariant", + ), + ], +) +def test_validate_rejects_unsupported_inference_training_forward(mutate, expected): + config = _inference_training_forward_config("nvfp4") + mutate(config) + assert any(expected in v for v in _validate(config).violations) + + +def test_generation_config_never_runs_the_inference_training_forward(): + from nemo_rl.models.generation.megatron.config import ( + merged_inference_megatron_cfg, + ) + + config = _inference_training_forward_config("mxfp8") + assert ( + merged_inference_megatron_cfg(config)["moe_inference_training_forward"] is False + ) + + # --------------------------------------------------------------------------- # # Package and platform gates # --------------------------------------------------------------------------- # From 140a1784287e0956c5ba9a3e5fdcf62f0601a765 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Mon, 19 Oct 2026 11:21:39 -0700 Subject: [PATCH 02/26] bump flashinfer to 0.7.0.post1 for per-token NVFP4 Signed-off-by: YigongQin --- pyproject.toml | 5 ++--- uv.lock | 60 ++++++++++++++++++++++++++------------------------ 2 files changed, 33 insertions(+), 32 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4e2770faaf7..a568ba70c0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,9 +236,8 @@ mcore = [ "mamba-ssm", "causal-conv1d", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@17cfb817bccec3a9c247013360cc550c2bac441e", - "flashinfer-python==0.6.18.post1", - "flashinfer-cubin==0.6.18.post1", - "flashinfer-jit-cache==0.6.18.post1", + "flashinfer-python==0.7.0.post1", + "flashinfer-cubin==0.7.0.post1", "flash_mla ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", "nvidia-cutlass-dsl[cu13]==4.6.2 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", "torch-memory-saver>=0.0.9.post1", diff --git a/uv.lock b/uv.lock index 80477156eaa..fd3a923f623 100644 --- a/uv.lock +++ b/uv.lock @@ -1819,14 +1819,14 @@ wheels = [ [[package]] name = "flashinfer-cubin" -version = "0.6.18.post1" +version = "0.7.0.post1" source = { registry = "https://flashinfer.ai/whl/" } resolution-markers = [ "platform_machine == 'x86_64' and sys_platform == 'linux'", "platform_machine == 'aarch64' and sys_platform == 'linux'", ] wheels = [ - { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.18.post1/flashinfer_cubin-0.6.18.post1-py3-none-any.whl", hash = "sha256:bbacb5b8bbf429e43bf2740bb45551d981f41d0462842e36ee6fb3e3452763ab" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_cubin-0.7.0.post1-py3-none-any.whl", hash = "sha256:f5b5fbf6e914b358219b845ed3ee586ffc1ae3c9768c8b135115ac14afbceed3" }, ] [[package]] @@ -1855,19 +1855,6 @@ wheels = [ { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.18/flashinfer_jit_cache-0.6.18+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:428a47a554ade93c30a818e142b781df58582bf056bab94611fb4c906cc366bf" }, ] -[[package]] -name = "flashinfer-jit-cache" -version = "0.6.18.post1+cu130" -source = { registry = "https://flashinfer.ai/whl/cu130" } -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", - "platform_machine == 'aarch64' and sys_platform == 'linux'", -] -wheels = [ - { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.18.post1/flashinfer_jit_cache-0.6.18.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:860ec5ce11a1d686bde04dbef4bedee15c39e3ec5e10adccb186dfee2e4113f2" }, - { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.18.post1/flashinfer_jit_cache-0.6.18.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:d1729490636a98f22f158518e10c257eec5a188d3dadf30b4c9de298675d8d4b" }, -] - [[package]] name = "flashinfer-python" version = "0.6.14" @@ -1935,7 +1922,7 @@ wheels = [ [[package]] name = "flashinfer-python" -version = "0.6.18.post1" +version = "0.7.0.post1" source = { registry = "https://flashinfer.ai/whl/" } resolution-markers = [ "platform_machine == 'x86_64' and sys_platform == 'linux'", @@ -1947,6 +1934,7 @@ dependencies = [ { name = "cuda-python" }, { name = "cuda-tile" }, { name = "einops" }, + { name = "nccl-extensions" }, { name = "nccl4py", version = "0.5.0", source = { registry = "https://pypi.org/simple" } }, { name = "ninja" }, { name = "numpy" }, @@ -1960,7 +1948,7 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.18.post1/flashinfer_python-0.6.18.post1-py3-none-any.whl", hash = "sha256:adb6d2952b47131138e108bd2999b1fe18b71f1a4c0b2684f168753db9563991" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_python-0.7.0.post1-py3-none-any.whl", hash = "sha256:c7adf826568d61fc1b7d3aadd4cae387a35a138bfc08e2deca2f34ecfa280716" }, ] [[package]] @@ -3237,8 +3225,8 @@ dependencies = [ { name = "einops" }, { name = "fast-hadamard-transform" }, { name = "flash-linear-attention" }, - { name = "flashinfer-cubin", version = "0.6.18.post1", source = { registry = "https://flashinfer.ai/whl/" } }, - { name = "flashinfer-python", version = "0.6.18.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-cubin", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-python", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, { name = "hydra-core" }, { name = "megatron-core", extra = ["dev", "inference", "mlm", "te"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "mistral-common" }, @@ -3281,8 +3269,8 @@ requires-dist = [ { name = "einops" }, { name = "fast-hadamard-transform", specifier = ">=1.1.0" }, { name = "flash-linear-attention" }, - { name = "flashinfer-cubin", specifier = "==0.6.18.post1", index = "https://flashinfer.ai/whl/" }, - { name = "flashinfer-python", specifier = "==0.6.18.post1", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-cubin", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-python", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl/" }, { name = "hydra-core", specifier = ">=1.3.4,<1.4" }, { name = "librosa", marker = "extra == 'audio'" }, { name = "mamba-ssm", marker = "extra == 'ssm'", specifier = "==2.3.1" }, @@ -3375,7 +3363,7 @@ dev = [ { name = "emerging-optimizers" }, { name = "fast-hadamard-transform" }, { name = "fastapi", extra = ["standard"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "flashinfer-python", version = "0.6.18.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-python", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, { name = "hypercorn" }, { name = "megatron-energon", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "multi-storage-client" }, @@ -3423,7 +3411,7 @@ requires-dist = [ { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "flash-linear-attention", marker = "extra == 'ssm'", specifier = "==0.5.1" }, - { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, + { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.8.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "gigatoken", marker = "extra == 'mlm'" }, @@ -3889,6 +3877,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/b5/1b84b2c784db76d69442334bc8b8748c840f13ca53be086f4f250ad4a0bc/narwhals-2.26.0-py3-none-any.whl", hash = "sha256:29326d74f107c347fd1009bd58e38d9f7c7c5b51e6de97bc93dbc325d9038b54", size = 474034, upload-time = "2026-09-08T13:32:07.159Z" }, ] +[[package]] +name = "nccl-extensions" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-core" }, + { name = "cuda-pathfinder" }, + { name = "nccl4py", version = "0.5.0", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy" }, + { name = "packaging" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/8b/0d98ca96c8625af88d834230b34115fac032ed69eb84c8bde7c38b956d16/nccl_extensions-0.1.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3678bce16864064a9f8774ec69f5c023cb2cbcd91db60a7f98e7c87fd4f18a50", size = 22265222, upload-time = "2026-09-02T15:39:12.533Z" }, + { url = "https://files.pythonhosted.org/packages/5b/66/0f0a3fe84ffdf6ddfb13c01a036ea370519204820cc8c53303f479c43a09/nccl_extensions-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3328dc1cfeed0ddf37caac00be8c49b3e4f727dce0e7980f8831696165619a4f", size = 22564580, upload-time = "2026-09-02T15:39:32.862Z" }, +] + [[package]] name = "nccl4py" version = "0.3.1.post1" @@ -4303,9 +4307,8 @@ mcore = [ { name = "flash-attn" }, { name = "flash-attn-4" }, { name = "flash-mla" }, - { name = "flashinfer-cubin", version = "0.6.18.post1", source = { registry = "https://flashinfer.ai/whl/" } }, - { name = "flashinfer-jit-cache", version = "0.6.18.post1+cu130", source = { registry = "https://flashinfer.ai/whl/cu130" } }, - { name = "flashinfer-python", version = "0.6.18.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-cubin", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-python", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, { name = "mamba-ssm" }, { name = "megatron-bridge", extra = ["ssm", "te"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "megatron-energon", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, @@ -4429,13 +4432,12 @@ requires-dist = [ { name = "flash-attn", marker = "extra == 'mcore'", specifier = "==2.8.1" }, { name = "flash-attn-4", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'mcore')", specifier = "==4.0.0b19" }, { name = "flash-mla", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'mcore')", git = "https://github.com/deepseek-ai/FlashMLA?rev=b7643bd54521f563b839b98289b5cd048c062ba2" }, - { name = "flashinfer-cubin", marker = "extra == 'mcore'", specifier = "==0.6.18.post1", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-cubin", marker = "extra == 'mcore'", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl/" }, { name = "flashinfer-cubin", marker = "extra == 'sglang'", specifier = "==0.6.14", index = "https://flashinfer.ai/whl/" }, { name = "flashinfer-cubin", marker = "extra == 'vllm'", specifier = "==0.6.18", index = "https://flashinfer.ai/whl/" }, - { name = "flashinfer-jit-cache", marker = "extra == 'mcore'", specifier = "==0.6.18.post1", index = "https://flashinfer.ai/whl/cu130" }, { name = "flashinfer-jit-cache", marker = "extra == 'sglang'", specifier = "==0.6.14", index = "https://flashinfer.ai/whl/cu130" }, { name = "flashinfer-jit-cache", marker = "extra == 'vllm'", specifier = "==0.6.18", index = "https://flashinfer.ai/whl/cu130" }, - { name = "flashinfer-python", marker = "extra == 'mcore'", specifier = "==0.6.18.post1" }, + { name = "flashinfer-python", marker = "extra == 'mcore'", specifier = "==0.7.0.post1" }, { name = "flashinfer-python", marker = "extra == 'sglang'", specifier = "==0.6.14" }, { name = "flashinfer-python", marker = "extra == 'vllm'", specifier = "==0.6.18" }, { name = "hydra-core" }, From 96bc3109670b6d21425482b802a360b8de4e3770 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 20 Oct 2026 09:20:26 -0700 Subject: [PATCH 03/26] pipeclean pertoken nvfp4 changes Signed-off-by: YigongQin --- nemo_rl/models/megatron/setup.py | 8 +++++-- .../megatron/zero_train_gen_mismatch.py | 24 ++++++++++++++----- nemo_rl/models/policy/__init__.py | 4 ++-- .../policy/workers/megatron_policy_worker.py | 2 +- .../megatron/test_zero_train_gen_mismatch.py | 23 ++++++++++++++++++ 5 files changed, 50 insertions(+), 11 deletions(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index ed31cb6facb..0581ece74fb 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -447,9 +447,13 @@ def _resolve_optimizer_dtype_kwargs(optimizer_cfg: dict[str, Any]) -> dict[str, return resolved -def enable_batch_invariant_mode(config: PolicyConfig) -> None: +def enable_batch_invariant_mode( + config: PolicyConfig, *, generation_only: bool = False +) -> None: """Validate, then enable Megatron-Core batch-invariant kernels before CUDA init.""" - validate_zero_train_gen_kl(config, check_environment=True) + validate_zero_train_gen_kl( + config, check_environment=True, generation_only=generation_only + ) if config["megatron_cfg"].get("batch_invariant_mode"): enable_batch_invariant_kernels(config) diff --git a/nemo_rl/models/megatron/zero_train_gen_mismatch.py b/nemo_rl/models/megatron/zero_train_gen_mismatch.py index ea862f06f00..b72c9e67025 100644 --- a/nemo_rl/models/megatron/zero_train_gen_mismatch.py +++ b/nemo_rl/models/megatron/zero_train_gen_mismatch.py @@ -100,14 +100,21 @@ def validate_zero_train_gen_mismatch( *, check_packages: bool = True, check_platform: bool = True, + generation_only: bool = False, ) -> ZeroTrainGenValidation: - """Run all zero-KL gates; return violations and non-fatal warnings.""" + """Run all zero-KL gates; return violations and non-fatal warnings. + + ``generation_only`` is for a dedicated generation policy, whose + ``megatron_cfg`` is the merged inference config: the training-only + ``moe_inference_training_forward`` is cleared there, so the checks on the + training side's value are left to the driver, which sees the full config. + """ out = ZeroTrainGenValidation() if not config["megatron_cfg"].get("zero_train_gen_mismatch"): return out _validate_required_values(config, out) - _validate_backend(config, out) + _validate_backend(config, out, generation_only=generation_only) _validate_platform(out, check_device=check_platform) if check_packages: _validate_packages(out) @@ -427,13 +434,14 @@ def enable_batch_invariant_kernels(config: PolicyConfig) -> None: def validate_zero_train_gen_kl( - config: PolicyConfig, *, check_environment: bool + config: PolicyConfig, *, check_environment: bool, generation_only: bool = False ) -> None: """Raise if the zero-KL / batch-invariant config is invalid. ``check_environment=False`` runs only the config checks, so the driver can call it before Ray allocates GPUs; workers also check the installed - packages and the GPU. + packages and the GPU. ``generation_only`` is as in + :func:`validate_zero_train_gen_mismatch`. """ if "megatron_cfg" not in config: # Automodel policies return @@ -441,6 +449,7 @@ def validate_zero_train_gen_kl( config, check_packages=check_environment, check_platform=check_environment, + generation_only=generation_only, ) result.emit_warnings() result.raise_if_invalid( @@ -465,7 +474,9 @@ def _validate_required_values( ) -def _validate_backend(config: PolicyConfig, out: ZeroTrainGenValidation) -> None: +def _validate_backend( + config: PolicyConfig, out: ZeroTrainGenValidation, *, generation_only: bool = False +) -> None: generation = config.get("generation") if generation is None or generation.get("backend") != "megatron": out.violations.append( @@ -493,7 +504,8 @@ def _validate_backend(config: PolicyConfig, out: ZeroTrainGenValidation) -> None fp8_cfg.get("enabled") and fp8_cfg.get("fp8_recipe") == "mxfp8" ) or inference_cfg.get("inference_moe_nvfp4_recipe") is not None if ( - impl == "inference_optimized" + not generation_only + and impl == "inference_optimized" and backend == "flashinfer" and quantized_experts and not config["megatron_cfg"].get("moe_inference_training_forward") diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index fbaf1f85d1e..b395459ea64 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -622,8 +622,8 @@ class MegatronConfig(TypedDict): num_layers_at_start_in_bf16: NotRequired[int] num_layers_at_end_in_bf16: NotRequired[int] # Path to a per-module Transformer Engine precision recipe loaded into - # Megatron quant_recipe. - te_precision_config_file: NotRequired[str] + # Megatron quant_recipe. null clears one inherited from a parent recipe. + te_precision_config_file: NotRequired[str | None] # Passed through to the Megatron model's freeze() method. # Supported keys are model-specific, such as freeze_vision_model, # freeze_vision_projection, and freeze_language_model. diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index dd530f22e6f..be2ae650589 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -851,7 +851,7 @@ def __init__( local_rank = int(gpu_ids[0]) os.environ["LOCAL_RANK"] = str(local_rank) # Batch-invariant / zero-KL kernels must be enabled before CUDA init. - enable_batch_invariant_mode(config) + enable_batch_invariant_mode(config, generation_only=is_refit_destination) torch.cuda.set_device(local_rank) # Apply patch from https://github.com/NVIDIA/TransformerEngine/pull/2286/files diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 0a7048b2ee8..6764b9dbfc7 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -602,6 +602,25 @@ def test_generation_config_never_runs_the_inference_training_forward(): ) +@pytest.mark.parametrize("precision", ["mxfp8", "nvfp4"]) +def test_dedicated_generation_policy_passes_validation(precision): + from nemo_rl.models.generation.megatron.config import ( + merged_inference_megatron_cfg, + ) + + config = _inference_training_forward_config(precision) + generation_policy = { + **config, + "megatron_cfg": merged_inference_megatron_cfg(config), + } + + assert _validate(generation_policy, generation_only=True).violations == [] + assert any( + "requires policy.megatron_cfg.moe_inference_training_forward=true" in v + for v in _validate(generation_policy).violations + ) + + # --------------------------------------------------------------------------- # # Package and platform gates # --------------------------------------------------------------------------- # @@ -821,6 +840,10 @@ def test_enable_batch_invariant_kernels_passes_the_recipe_settings(monkeypatch): monkeypatch.setattr( zgm, "_enable_mcore_deterministic_tuning", lambda: tuned.append(True) ) + tuned: list[bool] = [] + monkeypatch.setattr( + zgm, "_enable_mcore_deterministic_tuning", lambda: tuned.append(True) + ) zgm.enable_batch_invariant_kernels(_policy_config()) From 152726eb3bc96aa1c3dc08a5588a875d82ede980 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 20 Oct 2026 10:48:45 -0700 Subject: [PATCH 04/26] pin Bridge/Megatron-LM forks for per-token NVFP4 zero-KL Bridge 7c984ce = Bridge main f73818d + Megatron-LM #7798 (d9b0ba0) + FlashInfer 0.7.0.post1. Signed-off-by: YigongQin --- .gitmodules | 4 ++-- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.gitmodules b/.gitmodules index 8efea66a2a4..9daf89a4a49 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge - url = https://github.com/NVIDIA-NeMo/Megatron-Bridge.git - branch = main + url = https://github.com/YigongQin/Megatron-Bridge.git + branch = yigongq/nano-minf-training-forward shallow = true [submodule "3rdparty/Automodel-workspace/Automodel"] path = 3rdparty/Automodel-workspace/Automodel diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 7e9dcd609ed..7c984ce023e 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 7e9dcd609edf2bb9478d3fa4d9b9f8ee3875c16e +Subproject commit 7c984ce023e0bdee48848bdceff8cdeae97a2bfa From fd140c398e16a61d2da74e9749dde0063fbc1c74 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Mon, 26 Oct 2026 16:17:19 -0700 Subject: [PATCH 05/26] uv lock for Bridge 7c984ce (Megatron-LM #7798, FlashInfer 0.7.0.post1) Signed-off-by: YigongQin --- uv.lock | 94 +++++++++++++++------------------------------------------ 1 file changed, 24 insertions(+), 70 deletions(-) diff --git a/uv.lock b/uv.lock index fd3a923f623..98d21c1ec66 100644 --- a/uv.lock +++ b/uv.lock @@ -223,20 +223,20 @@ wheels = [ [[package]] name = "aiobotocore" -version = "3.9.1" +version = "3.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, { name = "aioitertools" }, - { name = "botocore", version = "1.43.75", source = { registry = "https://pypi.org/simple" } }, + { name = "botocore" }, { name = "jmespath" }, { name = "multidict" }, { name = "python-dateutil" }, { name = "wrapt" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a4/0d/98230a1320e78f2c35277ca194c494690dce3161ad05200556417ab9e816/aiobotocore-3.9.1.tar.gz", hash = "sha256:617d4a78da5adb6ae7d532835727ad388f73555f2ef129b9de1b1add73a43ee6", size = 516440, upload-time = "2026-09-07T17:14:07.099Z" } +sdist = { url = "https://files.pythonhosted.org/packages/16/40/659f2ac3915d329e3e8e5bf674d35ce1774f3117c6f60339d5e05e505c36/aiobotocore-3.10.0.tar.gz", hash = "sha256:35480ef11c51e23087273b85d97f80da2f356a46ce5f3fb88e24483cd343ffdb", size = 530781, upload-time = "2026-10-23T10:44:50.843Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/13/f5/20b5c946fdcfafa8e0538d4d278abf10dda4b4012f29aa98a32fe09a46c5/aiobotocore-3.9.1-py3-none-any.whl", hash = "sha256:23aa448b4daff3fdce0714c962a2050b63c0222b5f2b5c17d0c874eb76a50ac0", size = 102044, upload-time = "2026-09-07T17:14:05.224Z" }, + { url = "https://files.pythonhosted.org/packages/7f/bc/3bb8965204ba51bbf4d021bc89d1357cd2179a17ef71b5fc429e21b511a2/aiobotocore-3.10.0-py3-none-any.whl", hash = "sha256:de0ccbcb82b6b0009f1b6667f48442e9d0432f9b8bfe6172fa51777cbea89d03", size = 102759, upload-time = "2026-10-23T10:44:49.109Z" }, ] [[package]] @@ -646,72 +646,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/4d/1392562369b1139e741b30d624f09fe7091d17dd5579fae5732f044b12bb/blobfile-3.0.0-py3-none-any.whl", hash = "sha256:48ecc3307e622804bd8fe13bf6f40e6463c4439eba7a1f9ad49fd78aa63cc658", size = 75413, upload-time = "2024-08-27T00:02:51.518Z" }, ] -[[package]] -name = "boto3" -version = "1.43.75" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", - "platform_machine == 'aarch64' and sys_platform == 'linux'", -] -dependencies = [ - { name = "botocore", version = "1.43.75", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "jmespath", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "s3transfer", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b1/dd/b93be617dd5d1c6a3a248066030a450e62d211c0af0f49ed3a0d761383c9/boto3-1.43.75.tar.gz", hash = "sha256:86da93d3d5b46a58b03fa51b598ffa4aeaff485a3000affacf78491c105e44f0", size = 112673, upload-time = "2026-08-19T19:54:17.236Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/52/89/d0baf1269e57ac548387c11e5d9ba78d1e118ceb43f1e5e5114bdd2e8391/boto3-1.43.75-py3-none-any.whl", hash = "sha256:08a79edb68e6a0d65d305eb90188639313650e155aed015e86eef59d8475dde0", size = 140026, upload-time = "2026-08-19T19:54:15.509Z" }, -] - [[package]] name = "boto3" version = "1.43.115" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", - "platform_machine == 'aarch64' and sys_platform == 'linux'", -] dependencies = [ - { name = "botocore", version = "1.43.115", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "jmespath", marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "s3transfer", marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/b0ce191492a1724be73d8aacaf70d508bc8cbb9674c046d2ff7181dc51fb/boto3-1.43.115.tar.gz", hash = "sha256:c0f35e5f9cdf0edb365195e306d5586d3970b224660a11a87a4ec361ec3d7168", size = 112677, upload-time = "2026-10-19T19:19:53.093Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/be/bc/40dffae3238ef8d3415393132751fe2c88e8663ef84a31f0950e439ed2d5/boto3-1.43.115-py3-none-any.whl", hash = "sha256:6667a0c21d48dc125dbba88444bc8a601ce881a99341a0e6169f9333898e620d", size = 140044, upload-time = "2026-10-19T19:19:51.444Z" }, ] -[[package]] -name = "botocore" -version = "1.43.75" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", - "platform_machine == 'aarch64' and sys_platform == 'linux'", -] -dependencies = [ - { name = "jmespath", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "python-dateutil", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "urllib3", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a3/80/5e39eda5dfb66df691d71212799a9f9d35aa9e447511604030447a00e9d6/botocore-1.43.75.tar.gz", hash = "sha256:e8ed6b0f3cd398dfb9e08d7ca3a0b964152166a317a04e89c45ec91003327ffe", size = 15973868, upload-time = "2026-08-19T19:54:12.318Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/46/19d3178e70f37fda251f3d9560f29654f09c6eedbe1c8b790adeee84da49/botocore-1.43.75-py3-none-any.whl", hash = "sha256:121abc8b0b529bc4e29a28d1e8b096b20f26708cabe3d5ae4b3446747712aece", size = 15667237, upload-time = "2026-08-19T19:54:08.802Z" }, -] - [[package]] name = "botocore" version = "1.43.115" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", - "platform_machine == 'aarch64' and sys_platform == 'linux'", -] dependencies = [ - { name = "jmespath", marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "python-dateutil", marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "urllib3", marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/10/ca525e88db50f5153ddad2a3d55908b5ace915d10f443d5e5135ba808d8d/botocore-1.43.115.tar.gz", hash = "sha256:aa4d9c4e7ef99988b7347e0cd276ba84e1eb1625510557b0e36cd57ecb19bbf2", size = 16383796, upload-time = "2026-10-19T19:19:45.382Z" } wheels = [ @@ -3297,7 +3253,7 @@ requires-dist = [ { name = "timm" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", specifier = ">=4.67.1" }, - { name = "transformers", specifier = ">=5.15,<5.20" }, + { name = "transformers", specifier = ">=5.15,<5.21" }, { name = "typing-extensions" }, { name = "wandb", specifier = ">=0.29.0" }, ] @@ -3423,7 +3379,7 @@ requires-dist = [ { name = "nemo-lens", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'otel'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", specifier = "==1.29.0" }, + { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", specifier = "==1.29.1" }, { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'", specifier = ">=0.47" }, { name = "nvidia-resiliency-ext", marker = "extra == 'dev'", specifier = "==0.7.0" }, { name = "omegaconf", marker = "extra == 'mlm'" }, @@ -3433,7 +3389,7 @@ requires-dist = [ { name = "opentelemetry-api", marker = "extra == 'dev'", specifier = ">=1.33.1,<2" }, { name = "orjson", marker = "extra == 'dev'" }, { name = "packaging", specifier = ">=24.2" }, - { name = "pillow", marker = "extra == 'dev'" }, + { name = "pillow", marker = "extra == 'dev'", specifier = ">=13.0.0" }, { name = "quart", marker = "extra == 'dev'" }, { name = "sentencepiece", marker = "extra == 'mlm'" }, { name = "sentencepiece", marker = "extra == 'training'" }, @@ -3816,8 +3772,7 @@ wheels = [ [package.optional-dependencies] boto3 = [ - { name = "boto3", version = "1.43.75", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "boto3", version = "1.43.115", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "boto3" }, ] [[package]] @@ -6090,17 +6045,17 @@ wheels = [ [[package]] name = "pillow" -version = "12.3.0" +version = "13.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ce/a4/9cc30690a394a8305e8be4c9036aa16d0dfddadd09079860b582f562496d/pillow-13.0.0.tar.gz", hash = "sha256:cdcd1ba258178f7356d74170c9a81ba4eab6a5d55748bc402a5a2a6fe6be3510", size = 47058631, upload-time = "2026-10-16T22:18:09.718Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/ac/31fb64e1e7efb5a4b50cd3d92049ba89ac6e4d8d3bb6a74e15048ca3353e/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:21900ce7ba264168cd50defae43cd75d25c833ad4ad6e73ffc5596d12e25ac89", size = 4161684, upload-time = "2026-07-01T11:54:25.934Z" }, - { url = "https://files.pythonhosted.org/packages/87/b4/9805e23d2b4d77842b468513841fda254ee42f0289d25088340e4ff46e2d/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:4e8c2a84d977f50b9daed6eeaf3baef67d00d5d74d932288f02cb94518ee3ace", size = 4255487, upload-time = "2026-07-01T11:54:27.935Z" }, - { url = "https://files.pythonhosted.org/packages/df/39/ecf519435a200c693fe053a6ee4d835b41cf963a4dfc2551c4e637cb2a71/pillow-12.3.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:ae26d61dfa7a47befdc7572b521024e8745f3d809bd95ca9505a7bba9ef849ec", size = 3696433, upload-time = "2026-07-01T11:54:29.813Z" }, - { url = "https://files.pythonhosted.org/packages/1f/01/e18aff37cb0b4aac47ac90f016d347a49aca667ef97f190b06ac2aabc928/pillow-12.3.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f7401aebd7f581d7f83a439d87d474999317ee099218e5ad25d125290990ba65", size = 6263736, upload-time = "2026-07-01T11:54:36.131Z" }, - { url = "https://files.pythonhosted.org/packages/f7/62/de5bdd77d935331f4f802edc11e4d82950f642caad6cb2f949837b8560e2/pillow-12.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0847a763afefb695bc912d7c131e7e0632d4edc1d8698f58ddabec8e46b8b6d3", size = 6937129, upload-time = "2026-07-01T11:54:38.216Z" }, - { url = "https://files.pythonhosted.org/packages/70/4d/105627a13300c5e0df1d174230b32fd1273062c96f7745fd552b945d1e1d/pillow-12.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:571b9fcb07b97ef3a492028fb3d2dc0993ca23a06138b0315286566d29ef718a", size = 6339562, upload-time = "2026-07-01T11:54:40.354Z" }, - { url = "https://files.pythonhosted.org/packages/6b/1d/f13de01a553988ab895ba1c722e06cf3144d4f57656fd5b81b6d881f1179/pillow-12.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:756c768d0c9c2955feb7a56c37ea24aea2e369f8d36a88da270b6a9f19e62b5e", size = 7049439, upload-time = "2026-07-01T11:54:42.489Z" }, + { url = "https://files.pythonhosted.org/packages/82/e2/9132d38df320e8099e2642d9fa538f18a7dfebf68a875a31725dc2a115aa/pillow-13.0.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:e98b4f91f99b24b1da72e10621ba181f7281922384a7403094debb4199de6c00", size = 4190620, upload-time = "2026-10-16T22:16:02.527Z" }, + { url = "https://files.pythonhosted.org/packages/cb/be/436dac0996876f65ef1916ffb711c5ef4af2ebc57414b6f16642e07f6b47/pillow-13.0.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:34fb3db1149c879343861831728c243c432d4a5afc8f1878ca7612ef8de2a033", size = 4290286, upload-time = "2026-10-16T22:16:04.861Z" }, + { url = "https://files.pythonhosted.org/packages/60/16/7c1007162e4169ea0d5bec47e2e0b2174f56c50cf5b8c02a1dc23e41037e/pillow-13.0.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:2d50500a4a3bc6469adcdb319cf05d688b08ca3bf0f22a05339f3b25b433a599", size = 3723169, upload-time = "2026-10-16T22:16:06.956Z" }, + { url = "https://files.pythonhosted.org/packages/fb/3c/cd0770e96d20bfd990db024bcbd43897ff3e88da1f4782cf90ea231b1d4f/pillow-13.0.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:19b66ca0a4d25cc0200b381ce824c7ea090db61427315e73cfb05dbf1b35c369", size = 6515908, upload-time = "2026-10-16T22:16:12.896Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fe/f5cdaa2d1c4086e757430f368c271a58ee29cfe7378e7ae7b9d964015c25/pillow-13.0.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f56499c1a74d2b862beb0c33b8016f99d095858455df9d21583bbbc1c15b582", size = 7142544, upload-time = "2026-10-16T22:16:14.844Z" }, + { url = "https://files.pythonhosted.org/packages/fb/b5/a54f6d44adbcd23737c8c12953f42be107e4cca64e247c0ca3a5b563061c/pillow-13.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a2a14cd73d5e6c927f644b0110ee81f59c4ef506259ec9183904747f2caacc81", size = 6613097, upload-time = "2026-10-16T22:16:16.893Z" }, + { url = "https://files.pythonhosted.org/packages/e6/b4/9e6c1207c21b2e91db0d15c16d30f3701ce5add1bd4da0e787761e66cb6e/pillow-13.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:08b826a9bf1265c76de04a42061d61ea6b109418b7826f28ff1ecdfaa26c5502", size = 7224846, upload-time = "2026-10-16T22:16:19.195Z" }, ] [[package]] @@ -7145,8 +7100,7 @@ name = "s3transfer" version = "0.19.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "botocore", version = "1.43.75", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "botocore", version = "1.43.115", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "botocore" }, ] sdist = { url = "https://files.pythonhosted.org/packages/49/c5/0a5ce655109cb96b8dfad36908a4e108d15f5dc2eff8fbc3146c4312e704/s3transfer-0.19.3.tar.gz", hash = "sha256:8326acc3d056d7f7b8a37110c7d2aea41f67e83f6ae79cd3b5fbcca34c726406", size = 166087, upload-time = "2026-10-19T19:19:49.751Z" } wheels = [ From 3c33b9f67d57a0c8a52e942f4bb1cb95a4804522 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Mon, 26 Oct 2026 18:19:35 -0700 Subject: [PATCH 06/26] recipe: base NVFP4 FlashInfer zero-KL recipe on grpo_math_1B Signed-off-by: YigongQin --- ...noncolocated-zero-kl-nvfp4-flashinfer.yaml | 84 +++++++++++++++++-- 1 file changed, 76 insertions(+), 8 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index 886350aab86..a72dd60ea16 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -3,22 +3,90 @@ # and each training MoE forward runs the same kernel on the same quantization # (moe_inference_training_forward) with the backward from the TE recompute pass. # -# Needs Megatron-LM PR #7798 on top of the zero-kl-mxfp8 parent's requirements, and -# flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 MoE kernels). -defaults: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml +# Needs Megatron-LM PR #7798 and flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 +# MoE kernels). +defaults: ../../grpo_math_1B.yaml +grpo: + max_num_steps: 5 + num_prompts_per_step: 8 + num_generations_per_prompt: 4 checkpointing: + enabled: false checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer policy: + model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 + tokenizer: + name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + hf_config_overrides: + router_aux_loss_coef: 0 + train_global_batch_size: 32 + train_micro_batch_size: 1 + logprob_batch_size: 4 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 128 + automodel_cfg: + enabled: false megatron_cfg: + enabled: true + zero_train_gen_mismatch: true + batch_invariant_mode: true + batch_invariant_backend: te_native + batch_invariant_collective: ordered + attention_backend: flash + flash_attention_version: 4 + moe_permute_fusion: false + env_vars: + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False + CUBLASLT_WORKSPACE_SIZE: '0' + CUBLAS_WORKSPACE_CONFIG: :0:0 + MAMBA_DETERMINISTIC: '1' + TRITON_CACHE_AUTOTUNING: '1' + bias_activation_fusion: false + tensor_model_parallel_size: 2 + expert_model_parallel_size: 4 + sequence_parallel: true + moe_router_dtype: fp32 + defer_fp32_logits: true + use_fused_weighted_squared_relu: true + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls inference_moe_nvfp4_recipe: per_token - te_precision_config_file: null - fp8_cfg: - enabled: false + moe_inference_training_forward: true + activation_checkpointing: true + recompute_granularity: selective + recompute_modules: + - moe generation: + backend: megatron + refit_transport: mcore mcore_generation_config: - fp8_cfg: - enabled: false + refit_backend: nccl + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls + tensor_model_parallel_size: 2 + expert_model_parallel_size: 4 + sequence_parallel: true + moe_router_dtype: fp32 + logprobs_mode: raw_logprobs + mamba_inference_ssm_states_dtype: float32 + max_tokens: 256 + max_requests: 32 + enable_prefix_caching: false + moe_pad_experts_for_cuda_graph_inference: false + colocated: + enabled: false + resources: + gpus_per_node: ${cluster.gpus_per_node} + num_nodes: 1 logger: log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + wandb_enabled: true + tensorboard_enabled: true wandb: + project: nemo-rl name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer +cluster: + gpus_per_node: 4 + num_nodes: 2 From c9754345e7e68d84e7ab8a6964af86387461b049 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Fri, 30 Oct 2026 15:46:57 -0700 Subject: [PATCH 07/26] match_generation_forward rename, release value-pass weights, MXFP8 experts-only recipe, FlashInfer 0.7.2.post1 + jit-cache, enable FlashInfer nightlies Signed-off-by: YigongQin --- .../Megatron-Bridge-workspace/Megatron-Bridge | 2 +- ...noncolocated-zero-kl-mxfp8-flashinfer.yaml | 5 +- ...noncolocated-zero-kl-nvfp4-flashinfer.yaml | 4 +- .../te_precision/nanov3_experts_mxfp8.yaml | 51 ++++++ nemo_rl/models/generation/megatron/config.py | 4 +- nemo_rl/models/megatron/setup.py | 39 +---- nemo_rl/models/megatron/train.py | 5 + .../megatron/zero_train_gen_mismatch.py | 16 +- nemo_rl/models/policy/__init__.py | 6 +- pyproject.toml | 44 +++++- tests/test_suites/disabled.txt | 6 - ...n-noncolocated-zero-kl-mxfp8-flashinfer.sh | 2 +- ...n-noncolocated-zero-kl-nvfp4-flashinfer.sh | 2 +- tests/test_suites/nightly_gb200.txt | 2 + .../models/megatron/test_megatron_setup.py | 64 ++------ tests/unit/models/megatron/test_train.py | 33 ++++ .../megatron/test_zero_train_gen_mismatch.py | 30 ++-- uv.lock | 146 ++++++++++++++++-- 18 files changed, 310 insertions(+), 151 deletions(-) create mode 100644 examples/te_precision/nanov3_experts_mxfp8.yaml diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 7c984ce023e..9a9b39a52f9 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 7c984ce023e0bdee48848bdceff8cdeae97a2bfa +Subproject commit 9a9b39a52f9bf94d8773f017f9c666ad4c6f0a39 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml index 1a6d79f8afd..2a36cd698bb 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml @@ -1,6 +1,6 @@ # Nemotron Nano v3 zero train/generation KL with MXFP8 routed experts on generation's # FlashInfer routed MXFP8 kernel. The FlashInfer kernel does not match TE training, so -# each training MoE forward runs that kernel (moe_inference_training_forward) and the +# each training MoE forward runs that kernel (moe_match_generation_forward) and the # TE recompute pass of the MoE layer supplies the backward. # # Needs Megatron-LM PR #7798 on top of the zero-kl-mxfp8 parent's requirements. @@ -12,7 +12,8 @@ policy: transformer_impl: inference_optimized inference_grouped_gemm_backend: flashinfer inference_moe_token_dispatcher_type: nvls - moe_inference_training_forward: true + moe_match_generation_forward: true + te_precision_config_file: examples/te_precision/nanov3_experts_mxfp8.yaml activation_checkpointing: true recompute_granularity: selective recompute_modules: diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index a72dd60ea16..2748ab2797a 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -1,7 +1,7 @@ # Nemotron Nano v3 zero train/generation KL with per-token NVFP4 routed experts on # FlashInfer. Training and refit stay BF16; generation quantizes the experts to NVFP4, # and each training MoE forward runs the same kernel on the same quantization -# (moe_inference_training_forward) with the backward from the TE recompute pass. +# (moe_match_generation_forward) with the backward from the TE recompute pass. # # Needs Megatron-LM PR #7798 and flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 # MoE kernels). @@ -52,7 +52,7 @@ policy: inference_grouped_gemm_backend: flashinfer inference_moe_token_dispatcher_type: nvls inference_moe_nvfp4_recipe: per_token - moe_inference_training_forward: true + moe_match_generation_forward: true activation_checkpointing: true recompute_granularity: selective recompute_modules: diff --git a/examples/te_precision/nanov3_experts_mxfp8.yaml b/examples/te_precision/nanov3_experts_mxfp8.yaml new file mode 100644 index 00000000000..340cf778a5c --- /dev/null +++ b/examples/te_precision/nanov3_experts_mxfp8.yaml @@ -0,0 +1,51 @@ +# MXFP8 on the routed experts (GroupedLinear linear_fc1 / linear_fc2) of Nemotron Nano v3 +# only; every other linear stays BF16, including linear_qkv and linear_proj. For +# transformer_impl='inference_optimized' on the training side (moe_match_generation_forward): +# its log-prob forward runs under eval(), where the inference linears skip TE, so a TE MXFP8 +# qkv/proj weight would run a BF16 GEMM while generation runs MXFP8. +# +# Use the same file on train and generation (generation inherits megatron_cfg). +configs: + bf16: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: {} + mxfp8: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: + fp8_quantization_recipe: "mxfp8" + fp8_format: "hybrid" + fp8_param: true + override_quantized_autocast: true + override_nonquantized_autocast: true + # Grad-disabled model construction (the generation model) reads evaluation_recipe, so it + # repeats the training precision; an empty block would mean BF16. + evaluation_recipe: + fp8_quantization_recipe: "mxfp8" + fp8_format: "hybrid" + fp8_param: true + override_quantized_autocast: true + override_nonquantized_autocast: true + +matchers: + # First match wins. + # Before the routed experts so that no broader experts pattern can pull them into MXFP8. + shared_experts_bf16: + config: "bf16" + type: "glob" + pattern: "*.mlp.shared_experts.*" + enabled: true + experts_fc1_mxfp8: + config: "mxfp8" + type: "glob" + pattern: "*.mlp.experts.linear_fc1" + enabled: true + experts_fc2_mxfp8: + config: "mxfp8" + type: "glob" + pattern: "*.mlp.experts.linear_fc2" + enabled: true + fallthrough_bf16: + config: "bf16" + type: "glob" + pattern: "*" + enabled: true diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 5d6ef696362..fbb48055b96 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -216,8 +216,8 @@ def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any] } # Training-only: generation already runs the inference kernel, and MCore # rejects the mode without the MoE recompute that generation disables. - if merged.get("moe_inference_training_forward"): - merged["moe_inference_training_forward"] = False + if merged.get("moe_match_generation_forward"): + merged["moe_match_generation_forward"] = False # inference_optimized layers hard-require SP with TP>1. Raise with the # config key: the colocated build bypasses validate_and_set_config, so this # merge is the only spot the inference cfg gets a named error instead of a diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 0581ece74fb..cfdec14cb26 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1349,35 +1349,6 @@ def _apply_multimodal_config(model_cfg: Any, config: PolicyConfig) -> None: setattr(model_cfg, provider_attr, megatron_cfg[config_key]) -def moe_inference_training_max_tokens_per_rank(config: PolicyConfig) -> int: - """Tokens one rank feeds the largest moe_inference_training_forward pass. - - The NVLS symmetric heap is allocated once and cannot grow, so it is sized for - the largest train or log-prob microbatch at the padded maximum sequence length. - """ - configured = config["megatron_cfg"].get( - "moe_inference_training_max_tokens_per_rank" - ) - if configured is not None: - return configured - divisor = config["make_sequence_length_divisible_by"] - padded_seq_len = -(-config["max_total_sequence_length"] // divisor) * divisor - micro_batch = max( - config["train_micro_batch_size"], - config.get("logprob_batch_size", config["train_micro_batch_size"]), - ) - tokens = micro_batch * padded_seq_len - for batching_key in ("sequence_packing", "dynamic_batching"): - batching = config.get(batching_key) or {} - if batching.get("enabled"): - tokens = max( - tokens, - batching["train_mb_tokens"], - batching.get("logprob_mb_tokens", batching["train_mb_tokens"]), - ) - return tokens - - def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: """Apply Mixture of Experts configuration.""" model_cfg.expert_tensor_parallel_size = config["megatron_cfg"][ @@ -1426,14 +1397,10 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.inference_moe_nvfp4_recipe = config["megatron_cfg"][ "inference_moe_nvfp4_recipe" ] - if "moe_inference_training_forward" in config["megatron_cfg"]: - model_cfg.moe_inference_training_forward = config["megatron_cfg"][ - "moe_inference_training_forward" + if "moe_match_generation_forward" in config["megatron_cfg"]: + model_cfg.moe_match_generation_forward = config["megatron_cfg"][ + "moe_match_generation_forward" ] - if model_cfg.moe_inference_training_forward: - model_cfg.moe_inference_training_max_tokens_per_rank = ( - moe_inference_training_max_tokens_per_rank(config) - ) if "moe_router_num_groups" in config["megatron_cfg"]: model_cfg.moe_router_num_groups = config["megatron_cfg"][ "moe_router_num_groups" diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 45d514f2772..6f41c2e0f18 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -587,6 +587,11 @@ def megatron_forward_backward( finally: if use_router_replay: clear_router_replay(model) + # The next call rebuilds the moe_match_generation_forward kernel weights + # on its first microbatch, so keeping them between calls only holds memory. + for model_chunk in model if isinstance(model, list) else [model]: + if hasattr(model_chunk, "release_match_generation_weights"): + model_chunk.release_match_generation_weights() class LossPostProcessor: diff --git a/nemo_rl/models/megatron/zero_train_gen_mismatch.py b/nemo_rl/models/megatron/zero_train_gen_mismatch.py index b72c9e67025..7a901d9143b 100644 --- a/nemo_rl/models/megatron/zero_train_gen_mismatch.py +++ b/nemo_rl/models/megatron/zero_train_gen_mismatch.py @@ -16,7 +16,7 @@ The recipe sets every required value; this module only checks them. Supported precisions are BF16 and MXFP8 (the same on train and generation), and FlashInfer -MXFP8 or per-token NVFP4 experts with moe_inference_training_forward. +MXFP8 or per-token NVFP4 experts with moe_match_generation_forward. """ from __future__ import annotations @@ -106,7 +106,7 @@ def validate_zero_train_gen_mismatch( ``generation_only`` is for a dedicated generation policy, whose ``megatron_cfg`` is the merged inference config: the training-only - ``moe_inference_training_forward`` is cleared there, so the checks on the + ``moe_match_generation_forward`` is cleared there, so the checks on the training side's value are left to the driver, which sees the full config. """ out = ZeroTrainGenValidation() @@ -508,17 +508,17 @@ def _validate_backend( and impl == "inference_optimized" and backend == "flashinfer" and quantized_experts - and not config["megatron_cfg"].get("moe_inference_training_forward") + and not config["megatron_cfg"].get("moe_match_generation_forward") ): out.violations.append( "zero_train_gen_mismatch with inference_grouped_gemm_backend='flashinfer' " "and MXFP8 or NVFP4 experts requires " - "policy.megatron_cfg.moe_inference_training_forward=true: the FlashInfer " + "policy.megatron_cfg.moe_match_generation_forward=true: the FlashInfer " "kernel is batch-invariant but not bitwise identical to TE training. " "Otherwise use 'torch' or 'vllm' with MXFP8." ) - if config["megatron_cfg"].get("moe_inference_training_forward"): - _validate_inference_training_forward(config, inference_cfg, out) + if config["megatron_cfg"].get("moe_match_generation_forward"): + _validate_match_generation_forward(config, inference_cfg, out) def _grouped_gemm_backend(cfg: dict[str, Any]) -> Any: @@ -526,12 +526,12 @@ def _grouped_gemm_backend(cfg: dict[str, Any]) -> Any: return getattr(backend, "value", backend) -def _validate_inference_training_forward( +def _validate_match_generation_forward( config: PolicyConfig, inference_cfg: dict[str, Any], out: ZeroTrainGenValidation ) -> None: """The training MoE forward runs generation's FlashInfer kernel; TE recompute backs it.""" megatron_cfg = config["megatron_cfg"] - prefix = "moe_inference_training_forward=true requires" + prefix = "moe_match_generation_forward=true requires" if megatron_cfg.get("transformer_impl") != "inference_optimized": out.violations.append( f"{prefix} policy.megatron_cfg.transformer_impl='inference_optimized' " diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index b395459ea64..d052f57d5db 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -502,11 +502,7 @@ class MegatronConfig(TypedDict): # inference_grouped_gemm_backend='flashinfer', batch_invariant_mode, and # activation_checkpointing with recompute_granularity='selective' and 'moe' in # recompute_modules. Generation never sets it. - moe_inference_training_forward: NotRequired[bool] - # Token bound per rank for the NVLS symmetric heap of that training forward. - # Unset: max(train_micro_batch_size, logprob_batch_size) * padded - # max_total_sequence_length, or the packing / dynamic batching mb_tokens if larger. - moe_inference_training_max_tokens_per_rank: NotRequired[int | None] + moe_match_generation_forward: NotRequired[bool] # InferenceTopKRouter requires moe_router_num_groups=None # (used when transformer_impl='inference_optimized') moe_router_num_groups: NotRequired[int | None] diff --git a/pyproject.toml b/pyproject.toml index a568ba70c0b..971163a166b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,8 +236,22 @@ mcore = [ "mamba-ssm", "causal-conv1d", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@17cfb817bccec3a9c247013360cc550c2bac441e", - "flashinfer-python==0.7.0.post1", - "flashinfer-cubin==0.7.0.post1", + "flashinfer-python==0.7.2.post1", + "flashinfer-cubin==0.7.2.post1", + # Megatron's FlashInfer MoE backend refuses to start without it (it ships the prebuilt + # TRT-LLM fused-MoE, routing and MXFP8/FP4 quantization modules). + "flashinfer-jit-cache==0.7.2.post1", + # Its per-arch wheels, listed directly: [tool.uv.sources] index pins only attach to direct + # requirements, so without these uv looks for them on PyPI. + "flashinfer-jit-cache-sm80==0.7.2.post1", + "flashinfer-jit-cache-sm89==0.7.2.post1", + "flashinfer-jit-cache-sm90a==0.7.2.post1", + "flashinfer-jit-cache-sm100a==0.7.2.post1", + "flashinfer-jit-cache-sm103a==0.7.2.post1", + "flashinfer-jit-cache-sm120f==0.7.2.post1", + # Required by the aarch64 meta wheel only; neither has an x86_64 build. + "flashinfer-jit-cache-sm110a==0.7.2.post1 ; platform_machine == 'aarch64'", + "flashinfer-jit-cache-sm121a==0.7.2.post1 ; platform_machine == 'aarch64'", "flash_mla ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", "nvidia-cutlass-dsl[cu13]==4.6.2 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", "torch-memory-saver>=0.0.9.post1", @@ -375,6 +389,16 @@ nvidia-modelopt = { git = "https://github.com/NVIDIA/Model-Optimizer", rev = "43 nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", tag = "v0.7.0" } flashinfer-cubin = { index = "flashinfer" } flashinfer-jit-cache = { index = "flashinfer-cu130" } +# Since 0.7.0 flashinfer-jit-cache is a meta-package over these per-arch wheels; each is also a +# direct requirement of the mcore extra so that these entries apply. +flashinfer-jit-cache-sm80 = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm89 = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm90a = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm100a = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm103a = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm120f = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm110a = { index = "flashinfer-cu130" } +flashinfer-jit-cache-sm121a = { index = "flashinfer-cu130" } [tool.uv.workspace] members = [ @@ -618,6 +642,22 @@ flash-attn = { MAX_JOBS = "4" } name = "flash-attn" requires-dist = ["torch", "einops", "setuptools", "psutil", "ninja"] +# uv reads one wheel's metadata per version, and the aarch64 meta wheel also requires +# sm110a/sm121a, which have no x86_64 build. Declare the union with their real markers. +[[tool.uv.dependency-metadata]] +name = "flashinfer-jit-cache" +version = "0.7.2.post1+cu130" +requires-dist = [ + "flashinfer-jit-cache-sm80==0.7.2.post1+cu130", + "flashinfer-jit-cache-sm89==0.7.2.post1+cu130", + "flashinfer-jit-cache-sm90a==0.7.2.post1+cu130", + "flashinfer-jit-cache-sm100a==0.7.2.post1+cu130", + "flashinfer-jit-cache-sm103a==0.7.2.post1+cu130", + "flashinfer-jit-cache-sm120f==0.7.2.post1+cu130", + "flashinfer-jit-cache-sm110a==0.7.2.post1+cu130 ; platform_machine == 'aarch64'", + "flashinfer-jit-cache-sm121a==0.7.2.post1+cu130 ; platform_machine == 'aarch64'", +] + [[tool.uv.dependency-metadata]] name = "flash-mla" version = "1.0.0+b7643bd" diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index dc09a4c6989..2d43c913cc5 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -102,9 +102,3 @@ research/template_project/tests/test_suites/llm/single_update_1n8g.sh # Full Flow-GRPO recipe requires an externally served GenRM OCR judge. research/flow_grpo/tests/test_suites/flow_grpo_qwen_image_ocr.sh - -# FlashInfer MXFP8 / per-token NVFP4 zero-KL via moe_inference_training_forward: needs -# Megatron-LM PR #7798 in the submodule pin (NVFP4 also flashinfer-python>=0.7.0) and -# has not run end to end yet. -tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh index ac0bbe943fd..45fa3d47f72 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh @@ -1,7 +1,7 @@ #!/bin/bash # Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with MXFP8 # routed experts on generation's FlashInfer routed MXFP8 kernel. Each training MoE forward -# runs that kernel (moe_inference_training_forward); the TE recompute pass supplies the +# runs that kernel (moe_match_generation_forward); the TE recompute pass supplies the # backward. SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh index 947547652cc..22212199b18 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh @@ -1,7 +1,7 @@ #!/bin/bash # Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with # per-token NVFP4 routed experts on FlashInfer. Training and refit stay BF16; each -# training MoE forward runs generation's NVFP4 kernel (moe_inference_training_forward) +# training MoE forward runs generation's NVFP4 kernel (moe_match_generation_forward) # and the TE recompute pass supplies the backward. SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 350a8326854..03f4e402188 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -25,6 +25,8 @@ tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh # Zero train/generation KL (batch-invariant Megatron training + Megatron inference) tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh # BF16 FlashInfer TRTLLM + NCCL reshard numerical correctness tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 16aeab5c0cc..52ae348ef5d 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -995,7 +995,7 @@ def test_moe_grouped_gemm_absent_keeps_default(self): assert not hasattr(model_cfg, "moe_grouped_gemm") assert not hasattr(model_cfg, "inference_flashinfer_mxfp8_token_capacity") assert not hasattr(model_cfg, "inference_moe_nvfp4_recipe") - assert not hasattr(model_cfg, "moe_inference_training_forward") + assert not hasattr(model_cfg, "moe_match_generation_forward") def test_flashinfer_mxfp8_token_capacity_forwarded(self): """Without this key MCore runs FlashInfer MXFP8 MoE over the full dispatcher buffer.""" @@ -1042,79 +1042,33 @@ def test_inference_optimized_moe_requires_fp32_router( _apply_moe_config(model_cfg, config) assert model_cfg.moe_router_dtype == moe_router_dtype - def test_inference_training_forward_forwarded_with_token_bound(self): - """The NVLS heap bound covers the larger of the train and log-prob microbatches.""" + def test_match_generation_forward_forwarded(self): from nemo_rl.models.megatron.setup import _apply_moe_config model_cfg = MagicMock() megatron_cfg = self._base_moe_megatron_cfg() megatron_cfg["inference_moe_nvfp4_recipe"] = "per_token" - megatron_cfg["moe_inference_training_forward"] = True - config = { - "megatron_cfg": megatron_cfg, - "train_micro_batch_size": 1, - "logprob_batch_size": 4, - "max_total_sequence_length": 1000, - "make_sequence_length_divisible_by": 64, - } - - _apply_moe_config(model_cfg, config) - - assert model_cfg.inference_moe_nvfp4_recipe == "per_token" - assert model_cfg.moe_inference_training_forward is True - assert model_cfg.moe_inference_training_max_tokens_per_rank == 4 * 1024 - - @pytest.mark.parametrize("batching_key", ["sequence_packing", "dynamic_batching"]) - def test_inference_training_forward_token_bound_covers_packed_microbatches( - self, batching_key - ): - from nemo_rl.models.megatron.setup import ( - moe_inference_training_max_tokens_per_rank, - ) - - config = { - "megatron_cfg": {}, - "train_micro_batch_size": 1, - "logprob_batch_size": 1, - "max_total_sequence_length": 1000, - "make_sequence_length_divisible_by": 64, - batching_key: { - "enabled": True, - "train_mb_tokens": 4096, - "logprob_mb_tokens": 8192, - }, - } - - assert moe_inference_training_max_tokens_per_rank(config) == 8192 - - def test_inference_training_forward_keeps_an_explicit_token_bound(self): - from nemo_rl.models.megatron.setup import _apply_moe_config - - model_cfg = MagicMock() - megatron_cfg = self._base_moe_megatron_cfg() - megatron_cfg["moe_inference_training_forward"] = True - megatron_cfg["moe_inference_training_max_tokens_per_rank"] = 300 + megatron_cfg["moe_match_generation_forward"] = True config = {"megatron_cfg": megatron_cfg} _apply_moe_config(model_cfg, config) - assert model_cfg.moe_inference_training_max_tokens_per_rank == 300 + assert model_cfg.inference_moe_nvfp4_recipe == "per_token" + assert model_cfg.moe_match_generation_forward is True - def test_inference_training_forward_false_resets_a_copied_provider(self): + def test_match_generation_forward_false_resets_a_copied_provider(self): """Generation can start from the training provider and must clear the flag.""" from nemo_rl.models.megatron.setup import _apply_moe_config model_cfg = MagicMock() - model_cfg.moe_inference_training_forward = True - model_cfg.moe_inference_training_max_tokens_per_rank = None + model_cfg.moe_match_generation_forward = True megatron_cfg = self._base_moe_megatron_cfg() - megatron_cfg["moe_inference_training_forward"] = False + megatron_cfg["moe_match_generation_forward"] = False config = {"megatron_cfg": megatron_cfg} _apply_moe_config(model_cfg, config) - assert model_cfg.moe_inference_training_forward is False - assert model_cfg.moe_inference_training_max_tokens_per_rank is None + assert model_cfg.moe_match_generation_forward is False def test_hybridep_input_prepadding_wins_after_bridge_validation(self): from nemo_rl.models.megatron import setup diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 62ae69acb79..35ddefbf2ce 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -1178,6 +1178,39 @@ def test_megatron_forward_backward_calls_forward_backward_func(self, mock_get_fb assert call_kwargs["seq_length"] == 128 assert call_kwargs["micro_batch_size"] == 2 + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_megatron_forward_backward_releases_match_generation_weights( + self, mock_get_fb + ): + """Every model chunk drops its value-pass weights once the call returns.""" + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + chunks = [MagicMock(), MagicMock()] + mock_fb_func = MagicMock( + side_effect=lambda **_: [ + c.release_match_generation_weights.assert_not_called() for c in chunks + ] + ) + mock_get_fb.return_value = mock_fb_func + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + megatron_forward_backward( + model=chunks, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + ) + + for chunk in chunks: + chunk.release_match_generation_weights.assert_called_once_with() + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") def test_megatron_forward_backward_forward_only(self, mock_get_fb): """Test megatron_forward_backward with forward_only=True.""" diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 6764b9dbfc7..49fb3bbccf4 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -503,7 +503,7 @@ def test_validate_rejects_flashinfer_grouped_gemm_with_mxfp8(): assert any("flashinfer" in v for v in _validate(config).violations) -def _inference_training_forward_config(precision: str) -> dict[str, Any]: +def _match_generation_forward_config(precision: str) -> dict[str, Any]: """FlashInfer experts on both sides, with the training forward running that kernel.""" config = _policy_config() config["megatron_cfg"].update( @@ -513,7 +513,7 @@ def _inference_training_forward_config(precision: str) -> dict[str, Any]: activation_checkpointing=True, recompute_granularity="selective", recompute_modules=["moe"], - moe_inference_training_forward=True, + moe_match_generation_forward=True, ) if precision == "mxfp8": config["megatron_cfg"]["fp8_cfg"] = dict(_MXFP8) @@ -523,15 +523,15 @@ def _inference_training_forward_config(precision: str) -> dict[str, Any]: @pytest.mark.parametrize("precision", ["mxfp8", "nvfp4"]) -def test_validate_accepts_flashinfer_with_inference_training_forward(precision): - assert _validate(_inference_training_forward_config(precision)).violations == [] +def test_validate_accepts_flashinfer_with_match_generation_forward(precision): + assert _validate(_match_generation_forward_config(precision)).violations == [] -def test_validate_rejects_flashinfer_nvfp4_without_inference_training_forward(): - config = _inference_training_forward_config("nvfp4") - config["megatron_cfg"]["moe_inference_training_forward"] = False +def test_validate_rejects_flashinfer_nvfp4_without_match_generation_forward(): + config = _match_generation_forward_config("nvfp4") + config["megatron_cfg"]["moe_match_generation_forward"] = False assert any( - "requires policy.megatron_cfg.moe_inference_training_forward=true" in v + "requires policy.megatron_cfg.moe_match_generation_forward=true" in v for v in _validate(config).violations ) @@ -585,20 +585,20 @@ def _set_generation(config: dict[str, Any], key: str, value: Any) -> None: ), ], ) -def test_validate_rejects_unsupported_inference_training_forward(mutate, expected): - config = _inference_training_forward_config("nvfp4") +def test_validate_rejects_unsupported_match_generation_forward(mutate, expected): + config = _match_generation_forward_config("nvfp4") mutate(config) assert any(expected in v for v in _validate(config).violations) -def test_generation_config_never_runs_the_inference_training_forward(): +def test_generation_config_never_runs_the_match_generation_forward(): from nemo_rl.models.generation.megatron.config import ( merged_inference_megatron_cfg, ) - config = _inference_training_forward_config("mxfp8") + config = _match_generation_forward_config("mxfp8") assert ( - merged_inference_megatron_cfg(config)["moe_inference_training_forward"] is False + merged_inference_megatron_cfg(config)["moe_match_generation_forward"] is False ) @@ -608,7 +608,7 @@ def test_dedicated_generation_policy_passes_validation(precision): merged_inference_megatron_cfg, ) - config = _inference_training_forward_config(precision) + config = _match_generation_forward_config(precision) generation_policy = { **config, "megatron_cfg": merged_inference_megatron_cfg(config), @@ -616,7 +616,7 @@ def test_dedicated_generation_policy_passes_validation(precision): assert _validate(generation_policy, generation_only=True).violations == [] assert any( - "requires policy.megatron_cfg.moe_inference_training_forward=true" in v + "requires policy.megatron_cfg.moe_match_generation_forward=true" in v for v in _validate(generation_policy).violations ) diff --git a/uv.lock b/uv.lock index 98d21c1ec66..75df71311c9 100644 --- a/uv.lock +++ b/uv.lock @@ -148,6 +148,11 @@ requires-dist = ["torch", "einops", "setuptools", "psutil", "ninja"] name = "flash-mla" version = "1.0.0+b7643bd" +[[manifest.dependency-metadata]] +name = "flashinfer-jit-cache" +version = "0.7.2.post1+cu130" +requires-dist = ["flashinfer-jit-cache-sm80==0.7.2.post1+cu130", "flashinfer-jit-cache-sm89==0.7.2.post1+cu130", "flashinfer-jit-cache-sm90a==0.7.2.post1+cu130", "flashinfer-jit-cache-sm100a==0.7.2.post1+cu130", "flashinfer-jit-cache-sm103a==0.7.2.post1+cu130", "flashinfer-jit-cache-sm120f==0.7.2.post1+cu130", "flashinfer-jit-cache-sm110a==0.7.2.post1+cu130 ; platform_machine == 'aarch64'", "flashinfer-jit-cache-sm121a==0.7.2.post1+cu130 ; platform_machine == 'aarch64'"] + [[manifest.dependency-metadata]] name = "logsage" version = "0.1.7" @@ -1775,14 +1780,14 @@ wheels = [ [[package]] name = "flashinfer-cubin" -version = "0.7.0.post1" +version = "0.7.2.post1" source = { registry = "https://flashinfer.ai/whl/" } resolution-markers = [ "platform_machine == 'x86_64' and sys_platform == 'linux'", "platform_machine == 'aarch64' and sys_platform == 'linux'", ] wheels = [ - { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_cubin-0.7.0.post1-py3-none-any.whl", hash = "sha256:f5b5fbf6e914b358219b845ed3ee586ffc1ae3c9768c8b135115ac14afbceed3" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_cubin-0.7.2.post1-py3-none-any.whl", hash = "sha256:1998c7006bcc06ee5dd3208a682af247382c6710d6b370a961381262c04c11f0" }, ] [[package]] @@ -1811,6 +1816,99 @@ wheels = [ { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.18/flashinfer_jit_cache-0.6.18+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:428a47a554ade93c30a818e142b781df58582bf056bab94611fb4c906cc366bf" }, ] +[[package]] +name = "flashinfer-jit-cache" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "flashinfer-jit-cache-sm100a" }, + { name = "flashinfer-jit-cache-sm103a" }, + { name = "flashinfer-jit-cache-sm110a", marker = "platform_machine != 'x86_64'" }, + { name = "flashinfer-jit-cache-sm120f" }, + { name = "flashinfer-jit-cache-sm121a", marker = "platform_machine != 'x86_64'" }, + { name = "flashinfer-jit-cache-sm80" }, + { name = "flashinfer-jit-cache-sm89" }, + { name = "flashinfer-jit-cache-sm90a" }, +] +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8449156043036d240f3581c25e6883810f7ea6d656f56e78183b39eee9066877" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:5a0a92ce56e11c19cec31948a4828b245314377eb587fe17307730b096225500" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm100a" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm100a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ed6c128c4033de6fb6edbd1743a008ee70d9a3d4c8d1b60b45d821bd7e4143a2" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm100a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2d118fdb9a4a403ae22c08192830885eb27aba6d363ffa9c88758c7a51d1c126" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm103a" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm103a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8883759aaf289b8d494f714bb28766c9bef2e84fe6115d439d10e2417a8c7ae3" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm103a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c9c8430cd8c49d80edc44af820d4817c1c1b99cbe738c40034b41a44403ec5c7" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm110a" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm110a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ff08fecf03bec168d37c8518c679e2332ba77d79824e806cc3cdd67c4f6b7a75" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm120f" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm120f-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:1d3582708053a70b383c6d4df842a18c86f19518e230eaff6c1f689ce279dda5" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm120f-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ab05a6e987bf79e18594d43ccc89ccdc8dbeacb3241c031ede8bc054a2baa9a7" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm121a" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm121a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:bcbd9b7d51cefcaaac9629cec260d37e253f6a8493876d045b6869ddf42b49c7" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm80" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm80-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:f84fd89aac415b9cf79494e5606a307ecff2e421bcff132b62c6dba88f447a3b" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm80-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b5076a39be7b63d15df44ecdfb9ce5be7206d39629b6d317b06d6e314f0b4cd6" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm89" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm89-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:18fa31a010e3942a00ed4b6ffb4def344b5e10fc1ff9f56a538ce646bc316a7c" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm89-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:61ae80c3a78f20ebd8ef350bfe16bde947242bb9ac93f1cafb63e26c04c009dc" }, +] + +[[package]] +name = "flashinfer-jit-cache-sm90a" +version = "0.7.2.post1+cu130" +source = { registry = "https://flashinfer.ai/whl/cu130" } +wheels = [ + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm90a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:3bd375e978255d3215ee69da3434ca4b9dcce7a2d655c5d3445cd048fcc62096" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache_sm90a-0.7.2.post1+cu130-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:52e5cc064b52290dad2f8889ee0eb5cd59621dfac8351841e4b3987d42d034a3" }, +] + [[package]] name = "flashinfer-python" version = "0.6.14" @@ -1878,7 +1976,7 @@ wheels = [ [[package]] name = "flashinfer-python" -version = "0.7.0.post1" +version = "0.7.2.post1" source = { registry = "https://flashinfer.ai/whl/" } resolution-markers = [ "platform_machine == 'x86_64' and sys_platform == 'linux'", @@ -1904,7 +2002,7 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.0.post1/flashinfer_python-0.7.0.post1-py3-none-any.whl", hash = "sha256:c7adf826568d61fc1b7d3aadd4cae387a35a138bfc08e2deca2f34ecfa280716" }, + { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_python-0.7.2.post1-py3-none-any.whl", hash = "sha256:bfdc18b9e1e4d8d0730fe4343dc0e284cd0c160086de80f99888bc4867922d49" }, ] [[package]] @@ -3181,8 +3279,8 @@ dependencies = [ { name = "einops" }, { name = "fast-hadamard-transform" }, { name = "flash-linear-attention" }, - { name = "flashinfer-cubin", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, - { name = "flashinfer-python", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-cubin", version = "0.7.2.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-python", version = "0.7.2.post1", source = { registry = "https://flashinfer.ai/whl/" } }, { name = "hydra-core" }, { name = "megatron-core", extra = ["dev", "inference", "mlm", "te"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "mistral-common" }, @@ -3225,8 +3323,8 @@ requires-dist = [ { name = "einops" }, { name = "fast-hadamard-transform", specifier = ">=1.1.0" }, { name = "flash-linear-attention" }, - { name = "flashinfer-cubin", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl/" }, - { name = "flashinfer-python", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-cubin", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-python", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/" }, { name = "hydra-core", specifier = ">=1.3.4,<1.4" }, { name = "librosa", marker = "extra == 'audio'" }, { name = "mamba-ssm", marker = "extra == 'ssm'", specifier = "==2.3.1" }, @@ -3319,7 +3417,7 @@ dev = [ { name = "emerging-optimizers" }, { name = "fast-hadamard-transform" }, { name = "fastapi", extra = ["standard"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "flashinfer-python", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-python", version = "0.7.2.post1", source = { registry = "https://flashinfer.ai/whl/" } }, { name = "hypercorn" }, { name = "megatron-energon", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "multi-storage-client" }, @@ -3367,7 +3465,7 @@ requires-dist = [ { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "flash-linear-attention", marker = "extra == 'ssm'", specifier = "==0.5.1" }, - { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.8.0" }, + { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.7.2,<0.8.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "gigatoken", marker = "extra == 'mlm'" }, @@ -3379,7 +3477,7 @@ requires-dist = [ { name = "nemo-lens", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'otel'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", specifier = "==1.29.1" }, + { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", specifier = "==1.32.0" }, { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'", specifier = ">=0.47" }, { name = "nvidia-resiliency-ext", marker = "extra == 'dev'", specifier = "==0.7.0" }, { name = "omegaconf", marker = "extra == 'mlm'" }, @@ -4262,8 +4360,17 @@ mcore = [ { name = "flash-attn" }, { name = "flash-attn-4" }, { name = "flash-mla" }, - { name = "flashinfer-cubin", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, - { name = "flashinfer-python", version = "0.7.0.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-cubin", version = "0.7.2.post1", source = { registry = "https://flashinfer.ai/whl/" } }, + { name = "flashinfer-jit-cache", version = "0.7.2.post1+cu130", source = { registry = "https://flashinfer.ai/whl/cu130" } }, + { name = "flashinfer-jit-cache-sm100a" }, + { name = "flashinfer-jit-cache-sm103a" }, + { name = "flashinfer-jit-cache-sm110a", marker = "platform_machine != 'x86_64'" }, + { name = "flashinfer-jit-cache-sm120f" }, + { name = "flashinfer-jit-cache-sm121a", marker = "platform_machine != 'x86_64'" }, + { name = "flashinfer-jit-cache-sm80" }, + { name = "flashinfer-jit-cache-sm89" }, + { name = "flashinfer-jit-cache-sm90a" }, + { name = "flashinfer-python", version = "0.7.2.post1", source = { registry = "https://flashinfer.ai/whl/" } }, { name = "mamba-ssm" }, { name = "megatron-bridge", extra = ["ssm", "te"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "megatron-energon", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, @@ -4387,12 +4494,21 @@ requires-dist = [ { name = "flash-attn", marker = "extra == 'mcore'", specifier = "==2.8.1" }, { name = "flash-attn-4", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'mcore')", specifier = "==4.0.0b19" }, { name = "flash-mla", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'mcore')", git = "https://github.com/deepseek-ai/FlashMLA?rev=b7643bd54521f563b839b98289b5cd048c062ba2" }, - { name = "flashinfer-cubin", marker = "extra == 'mcore'", specifier = "==0.7.0.post1", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-cubin", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/" }, { name = "flashinfer-cubin", marker = "extra == 'sglang'", specifier = "==0.6.14", index = "https://flashinfer.ai/whl/" }, { name = "flashinfer-cubin", marker = "extra == 'vllm'", specifier = "==0.6.18", index = "https://flashinfer.ai/whl/" }, + { name = "flashinfer-jit-cache", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, { name = "flashinfer-jit-cache", marker = "extra == 'sglang'", specifier = "==0.6.14", index = "https://flashinfer.ai/whl/cu130" }, { name = "flashinfer-jit-cache", marker = "extra == 'vllm'", specifier = "==0.6.18", index = "https://flashinfer.ai/whl/cu130" }, - { name = "flashinfer-python", marker = "extra == 'mcore'", specifier = "==0.7.0.post1" }, + { name = "flashinfer-jit-cache-sm100a", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm103a", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm110a", marker = "platform_machine == 'aarch64' and extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm120f", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm121a", marker = "platform_machine == 'aarch64' and extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm80", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm89", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-jit-cache-sm90a", marker = "extra == 'mcore'", specifier = "==0.7.2.post1", index = "https://flashinfer.ai/whl/cu130" }, + { name = "flashinfer-python", marker = "extra == 'mcore'", specifier = "==0.7.2.post1" }, { name = "flashinfer-python", marker = "extra == 'sglang'", specifier = "==0.6.14" }, { name = "flashinfer-python", marker = "extra == 'vllm'", specifier = "==0.6.18" }, { name = "hydra-core" }, From 6051078e149bee4dfb7c9beecaa1ea724beb4827 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 09:17:53 -0800 Subject: [PATCH 08/26] Megatron-LM #7798 6e77cf3: moe_training_matches_inference, all-to-all value pass Rename moe_match_generation_forward to moe_training_matches_inference and require the training all-to-all dispatcher, which the value pass now dispatches through. Drop the per-call kernel-weight release: Megatron-LM drops them on each optimizer step, and the log-prob pass and the step's microbatches share one build. Signed-off-by: YigongQin --- .../Megatron-Bridge-workspace/Megatron-Bridge | 2 +- ...noncolocated-zero-kl-mxfp8-flashinfer.yaml | 4 +-- ...noncolocated-zero-kl-nvfp4-flashinfer.yaml | 4 +-- .../te_precision/nanov3_experts_mxfp8.yaml | 2 +- nemo_rl/models/generation/megatron/config.py | 4 +-- nemo_rl/models/megatron/setup.py | 6 ++-- nemo_rl/models/megatron/train.py | 5 --- .../megatron/zero_train_gen_mismatch.py | 24 ++++++++----- nemo_rl/models/policy/__init__.py | 2 +- ...n-noncolocated-zero-kl-mxfp8-flashinfer.sh | 2 +- ...n-noncolocated-zero-kl-nvfp4-flashinfer.sh | 2 +- .../models/megatron/test_megatron_setup.py | 16 ++++----- tests/unit/models/megatron/test_train.py | 33 ----------------- .../megatron/test_zero_train_gen_mismatch.py | 35 ++++++++++--------- 14 files changed, 55 insertions(+), 86 deletions(-) mode change 100755 => 100644 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh mode change 100755 => 100644 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 9a9b39a52f9..21efc1550b3 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 9a9b39a52f9bf94d8773f017f9c666ad4c6f0a39 +Subproject commit 21efc1550b32256f109b9c24b086ffa99d23b646 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml index 2a36cd698bb..781a8a77420 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml @@ -1,6 +1,6 @@ # Nemotron Nano v3 zero train/generation KL with MXFP8 routed experts on generation's # FlashInfer routed MXFP8 kernel. The FlashInfer kernel does not match TE training, so -# each training MoE forward runs that kernel (moe_match_generation_forward) and the +# each training MoE forward runs that kernel (moe_training_matches_inference) and the # TE recompute pass of the MoE layer supplies the backward. # # Needs Megatron-LM PR #7798 on top of the zero-kl-mxfp8 parent's requirements. @@ -12,7 +12,7 @@ policy: transformer_impl: inference_optimized inference_grouped_gemm_backend: flashinfer inference_moe_token_dispatcher_type: nvls - moe_match_generation_forward: true + moe_training_matches_inference: true te_precision_config_file: examples/te_precision/nanov3_experts_mxfp8.yaml activation_checkpointing: true recompute_granularity: selective diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index 2748ab2797a..19b86d51424 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -1,7 +1,7 @@ # Nemotron Nano v3 zero train/generation KL with per-token NVFP4 routed experts on # FlashInfer. Training and refit stay BF16; generation quantizes the experts to NVFP4, # and each training MoE forward runs the same kernel on the same quantization -# (moe_match_generation_forward) with the backward from the TE recompute pass. +# (moe_training_matches_inference) with the backward from the TE recompute pass. # # Needs Megatron-LM PR #7798 and flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 # MoE kernels). @@ -52,7 +52,7 @@ policy: inference_grouped_gemm_backend: flashinfer inference_moe_token_dispatcher_type: nvls inference_moe_nvfp4_recipe: per_token - moe_match_generation_forward: true + moe_training_matches_inference: true activation_checkpointing: true recompute_granularity: selective recompute_modules: diff --git a/examples/te_precision/nanov3_experts_mxfp8.yaml b/examples/te_precision/nanov3_experts_mxfp8.yaml index 340cf778a5c..25f4e8c86f6 100644 --- a/examples/te_precision/nanov3_experts_mxfp8.yaml +++ b/examples/te_precision/nanov3_experts_mxfp8.yaml @@ -1,6 +1,6 @@ # MXFP8 on the routed experts (GroupedLinear linear_fc1 / linear_fc2) of Nemotron Nano v3 # only; every other linear stays BF16, including linear_qkv and linear_proj. For -# transformer_impl='inference_optimized' on the training side (moe_match_generation_forward): +# transformer_impl='inference_optimized' on the training side (moe_training_matches_inference): # its log-prob forward runs under eval(), where the inference linears skip TE, so a TE MXFP8 # qkv/proj weight would run a BF16 GEMM while generation runs MXFP8. # diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index fbb48055b96..d945b4df73a 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -216,8 +216,8 @@ def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any] } # Training-only: generation already runs the inference kernel, and MCore # rejects the mode without the MoE recompute that generation disables. - if merged.get("moe_match_generation_forward"): - merged["moe_match_generation_forward"] = False + if merged.get("moe_training_matches_inference"): + merged["moe_training_matches_inference"] = False # inference_optimized layers hard-require SP with TP>1. Raise with the # config key: the colocated build bypasses validate_and_set_config, so this # merge is the only spot the inference cfg gets a named error instead of a diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index cfdec14cb26..1f968f68b08 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1397,9 +1397,9 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.inference_moe_nvfp4_recipe = config["megatron_cfg"][ "inference_moe_nvfp4_recipe" ] - if "moe_match_generation_forward" in config["megatron_cfg"]: - model_cfg.moe_match_generation_forward = config["megatron_cfg"][ - "moe_match_generation_forward" + if "moe_training_matches_inference" in config["megatron_cfg"]: + model_cfg.moe_training_matches_inference = config["megatron_cfg"][ + "moe_training_matches_inference" ] if "moe_router_num_groups" in config["megatron_cfg"]: model_cfg.moe_router_num_groups = config["megatron_cfg"][ diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 6f41c2e0f18..45d514f2772 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -587,11 +587,6 @@ def megatron_forward_backward( finally: if use_router_replay: clear_router_replay(model) - # The next call rebuilds the moe_match_generation_forward kernel weights - # on its first microbatch, so keeping them between calls only holds memory. - for model_chunk in model if isinstance(model, list) else [model]: - if hasattr(model_chunk, "release_match_generation_weights"): - model_chunk.release_match_generation_weights() class LossPostProcessor: diff --git a/nemo_rl/models/megatron/zero_train_gen_mismatch.py b/nemo_rl/models/megatron/zero_train_gen_mismatch.py index 7a901d9143b..6e6c5e38708 100644 --- a/nemo_rl/models/megatron/zero_train_gen_mismatch.py +++ b/nemo_rl/models/megatron/zero_train_gen_mismatch.py @@ -16,7 +16,7 @@ The recipe sets every required value; this module only checks them. Supported precisions are BF16 and MXFP8 (the same on train and generation), and FlashInfer -MXFP8 or per-token NVFP4 experts with moe_match_generation_forward. +MXFP8 or per-token NVFP4 experts with moe_training_matches_inference. """ from __future__ import annotations @@ -106,7 +106,7 @@ def validate_zero_train_gen_mismatch( ``generation_only`` is for a dedicated generation policy, whose ``megatron_cfg`` is the merged inference config: the training-only - ``moe_match_generation_forward`` is cleared there, so the checks on the + ``moe_training_matches_inference`` is cleared there, so the checks on the training side's value are left to the driver, which sees the full config. """ out = ZeroTrainGenValidation() @@ -508,17 +508,17 @@ def _validate_backend( and impl == "inference_optimized" and backend == "flashinfer" and quantized_experts - and not config["megatron_cfg"].get("moe_match_generation_forward") + and not config["megatron_cfg"].get("moe_training_matches_inference") ): out.violations.append( "zero_train_gen_mismatch with inference_grouped_gemm_backend='flashinfer' " "and MXFP8 or NVFP4 experts requires " - "policy.megatron_cfg.moe_match_generation_forward=true: the FlashInfer " + "policy.megatron_cfg.moe_training_matches_inference=true: the FlashInfer " "kernel is batch-invariant but not bitwise identical to TE training. " "Otherwise use 'torch' or 'vllm' with MXFP8." ) - if config["megatron_cfg"].get("moe_match_generation_forward"): - _validate_match_generation_forward(config, inference_cfg, out) + if config["megatron_cfg"].get("moe_training_matches_inference"): + _validate_training_matches_inference(config, inference_cfg, out) def _grouped_gemm_backend(cfg: dict[str, Any]) -> Any: @@ -526,12 +526,12 @@ def _grouped_gemm_backend(cfg: dict[str, Any]) -> Any: return getattr(backend, "value", backend) -def _validate_match_generation_forward( +def _validate_training_matches_inference( config: PolicyConfig, inference_cfg: dict[str, Any], out: ZeroTrainGenValidation ) -> None: """The training MoE forward runs generation's FlashInfer kernel; TE recompute backs it.""" megatron_cfg = config["megatron_cfg"] - prefix = "moe_match_generation_forward=true requires" + prefix = "moe_training_matches_inference=true requires" if megatron_cfg.get("transformer_impl") != "inference_optimized": out.violations.append( f"{prefix} policy.megatron_cfg.transformer_impl='inference_optimized' " @@ -570,11 +570,17 @@ def _validate_match_generation_forward( f"{prefix} a training cuda_graph_impl other than 'local', which disables " "MoE-layer recompute." ) + if megatron_cfg.get("moe_token_dispatcher_type") != "alltoall": + out.violations.append( + f"{prefix} policy.megatron_cfg.moe_token_dispatcher_type='alltoall': the " + "value pass dispatches through training's dispatcher, and only the " + "all-to-all combine adds the EP ranks in generation's order " + f"(got {megatron_cfg.get('moe_token_dispatcher_type')!r})." + ) mismatched = [ key for key in ( "inference_moe_nvfp4_recipe", - "inference_moe_token_dispatcher_type", "expert_model_parallel_size", ) if megatron_cfg.get(key) != inference_cfg.get(key) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index d052f57d5db..83b314abc88 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -502,7 +502,7 @@ class MegatronConfig(TypedDict): # inference_grouped_gemm_backend='flashinfer', batch_invariant_mode, and # activation_checkpointing with recompute_granularity='selective' and 'moe' in # recompute_modules. Generation never sets it. - moe_match_generation_forward: NotRequired[bool] + moe_training_matches_inference: NotRequired[bool] # InferenceTopKRouter requires moe_router_num_groups=None # (used when transformer_impl='inference_optimized') moe_router_num_groups: NotRequired[int | None] diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh old mode 100755 new mode 100644 index 45fa3d47f72..c2282b21994 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh @@ -1,7 +1,7 @@ #!/bin/bash # Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with MXFP8 # routed experts on generation's FlashInfer routed MXFP8 kernel. Each training MoE forward -# runs that kernel (moe_match_generation_forward); the TE recompute pass supplies the +# runs that kernel (moe_training_matches_inference); the TE recompute pass supplies the # backward. SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh old mode 100755 new mode 100644 index 22212199b18..aa6dd27b397 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh @@ -1,7 +1,7 @@ #!/bin/bash # Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with # per-token NVFP4 routed experts on FlashInfer. Training and refit stay BF16; each -# training MoE forward runs generation's NVFP4 kernel (moe_match_generation_forward) +# training MoE forward runs generation's NVFP4 kernel (moe_training_matches_inference) # and the TE recompute pass supplies the backward. SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 52ae348ef5d..40e9b56baac 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -995,7 +995,7 @@ def test_moe_grouped_gemm_absent_keeps_default(self): assert not hasattr(model_cfg, "moe_grouped_gemm") assert not hasattr(model_cfg, "inference_flashinfer_mxfp8_token_capacity") assert not hasattr(model_cfg, "inference_moe_nvfp4_recipe") - assert not hasattr(model_cfg, "moe_match_generation_forward") + assert not hasattr(model_cfg, "moe_training_matches_inference") def test_flashinfer_mxfp8_token_capacity_forwarded(self): """Without this key MCore runs FlashInfer MXFP8 MoE over the full dispatcher buffer.""" @@ -1042,33 +1042,33 @@ def test_inference_optimized_moe_requires_fp32_router( _apply_moe_config(model_cfg, config) assert model_cfg.moe_router_dtype == moe_router_dtype - def test_match_generation_forward_forwarded(self): + def test_training_matches_inference_forwarded(self): from nemo_rl.models.megatron.setup import _apply_moe_config model_cfg = MagicMock() megatron_cfg = self._base_moe_megatron_cfg() megatron_cfg["inference_moe_nvfp4_recipe"] = "per_token" - megatron_cfg["moe_match_generation_forward"] = True + megatron_cfg["moe_training_matches_inference"] = True config = {"megatron_cfg": megatron_cfg} _apply_moe_config(model_cfg, config) assert model_cfg.inference_moe_nvfp4_recipe == "per_token" - assert model_cfg.moe_match_generation_forward is True + assert model_cfg.moe_training_matches_inference is True - def test_match_generation_forward_false_resets_a_copied_provider(self): + def test_training_matches_inference_false_resets_a_copied_provider(self): """Generation can start from the training provider and must clear the flag.""" from nemo_rl.models.megatron.setup import _apply_moe_config model_cfg = MagicMock() - model_cfg.moe_match_generation_forward = True + model_cfg.moe_training_matches_inference = True megatron_cfg = self._base_moe_megatron_cfg() - megatron_cfg["moe_match_generation_forward"] = False + megatron_cfg["moe_training_matches_inference"] = False config = {"megatron_cfg": megatron_cfg} _apply_moe_config(model_cfg, config) - assert model_cfg.moe_match_generation_forward is False + assert model_cfg.moe_training_matches_inference is False def test_hybridep_input_prepadding_wins_after_bridge_validation(self): from nemo_rl.models.megatron import setup diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 35ddefbf2ce..62ae69acb79 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -1178,39 +1178,6 @@ def test_megatron_forward_backward_calls_forward_backward_func(self, mock_get_fb assert call_kwargs["seq_length"] == 128 assert call_kwargs["micro_batch_size"] == 2 - @patch("nemo_rl.models.megatron.train.get_forward_backward_func") - def test_megatron_forward_backward_releases_match_generation_weights( - self, mock_get_fb - ): - """Every model chunk drops its value-pass weights once the call returns.""" - from nemo_rl.models.megatron.train import ( - LossPostProcessor, - megatron_forward_backward, - ) - - chunks = [MagicMock(), MagicMock()] - mock_fb_func = MagicMock( - side_effect=lambda **_: [ - c.release_match_generation_weights.assert_not_called() for c in chunks - ] - ) - mock_get_fb.return_value = mock_fb_func - post_processor = LossPostProcessor( - loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} - ) - - megatron_forward_backward( - model=chunks, - data_iterator=iter([]), - num_microbatches=1, - seq_length=64, - mbs=1, - post_processing_fn=post_processor, - ) - - for chunk in chunks: - chunk.release_match_generation_weights.assert_called_once_with() - @patch("nemo_rl.models.megatron.train.get_forward_backward_func") def test_megatron_forward_backward_forward_only(self, mock_get_fb): """Test megatron_forward_backward with forward_only=True.""" diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 49fb3bbccf4..9befbbe2ca5 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -503,17 +503,18 @@ def test_validate_rejects_flashinfer_grouped_gemm_with_mxfp8(): assert any("flashinfer" in v for v in _validate(config).violations) -def _match_generation_forward_config(precision: str) -> dict[str, Any]: +def _training_matches_inference_config(precision: str) -> dict[str, Any]: """FlashInfer experts on both sides, with the training forward running that kernel.""" config = _policy_config() config["megatron_cfg"].update( transformer_impl="inference_optimized", inference_grouped_gemm_backend="flashinfer", inference_moe_token_dispatcher_type="nvls", + moe_token_dispatcher_type="alltoall", activation_checkpointing=True, recompute_granularity="selective", recompute_modules=["moe"], - moe_match_generation_forward=True, + moe_training_matches_inference=True, ) if precision == "mxfp8": config["megatron_cfg"]["fp8_cfg"] = dict(_MXFP8) @@ -523,15 +524,15 @@ def _match_generation_forward_config(precision: str) -> dict[str, Any]: @pytest.mark.parametrize("precision", ["mxfp8", "nvfp4"]) -def test_validate_accepts_flashinfer_with_match_generation_forward(precision): - assert _validate(_match_generation_forward_config(precision)).violations == [] +def test_validate_accepts_flashinfer_with_training_matches_inference(precision): + assert _validate(_training_matches_inference_config(precision)).violations == [] -def test_validate_rejects_flashinfer_nvfp4_without_match_generation_forward(): - config = _match_generation_forward_config("nvfp4") - config["megatron_cfg"]["moe_match_generation_forward"] = False +def test_validate_rejects_flashinfer_nvfp4_without_training_matches_inference(): + config = _training_matches_inference_config("nvfp4") + config["megatron_cfg"]["moe_training_matches_inference"] = False assert any( - "requires policy.megatron_cfg.moe_match_generation_forward=true" in v + "requires policy.megatron_cfg.moe_training_matches_inference=true" in v for v in _validate(config).violations ) @@ -568,8 +569,8 @@ def _set_generation(config: dict[str, Any], key: str, value: Any) -> None: "cuda_graph_impl other than 'local'", ), ( - lambda c: _set_generation(c, "inference_moe_token_dispatcher_type", "nccl"), - "inference_moe_token_dispatcher_type", + lambda c: c["megatron_cfg"].update(moe_token_dispatcher_type="flex"), + "moe_token_dispatcher_type='alltoall'", ), ( lambda c: _set_generation(c, "inference_moe_nvfp4_recipe", None), @@ -585,20 +586,20 @@ def _set_generation(config: dict[str, Any], key: str, value: Any) -> None: ), ], ) -def test_validate_rejects_unsupported_match_generation_forward(mutate, expected): - config = _match_generation_forward_config("nvfp4") +def test_validate_rejects_unsupported_training_matches_inference(mutate, expected): + config = _training_matches_inference_config("nvfp4") mutate(config) assert any(expected in v for v in _validate(config).violations) -def test_generation_config_never_runs_the_match_generation_forward(): +def test_generation_config_never_runs_the_training_matches_inference(): from nemo_rl.models.generation.megatron.config import ( merged_inference_megatron_cfg, ) - config = _match_generation_forward_config("mxfp8") + config = _training_matches_inference_config("mxfp8") assert ( - merged_inference_megatron_cfg(config)["moe_match_generation_forward"] is False + merged_inference_megatron_cfg(config)["moe_training_matches_inference"] is False ) @@ -608,7 +609,7 @@ def test_dedicated_generation_policy_passes_validation(precision): merged_inference_megatron_cfg, ) - config = _match_generation_forward_config(precision) + config = _training_matches_inference_config(precision) generation_policy = { **config, "megatron_cfg": merged_inference_megatron_cfg(config), @@ -616,7 +617,7 @@ def test_dedicated_generation_policy_passes_validation(precision): assert _validate(generation_policy, generation_only=True).violations == [] assert any( - "requires policy.megatron_cfg.moe_match_generation_forward=true" in v + "requires policy.megatron_cfg.moe_training_matches_inference=true" in v for v in _validate(generation_policy).violations ) From 570d61d5df334f9f4721f54f80d7fcede4a6c2dd Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 09:31:29 -0800 Subject: [PATCH 09/26] test_suites: restore the executable bit on the FlashInfer zero-KL suites Signed-off-by: YigongQin --- ...g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh | 0 ...g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh | 0 2 files changed, 0 insertions(+), 0 deletions(-) mode change 100644 => 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh mode change 100644 => 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh old mode 100644 new mode 100755 diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh old mode 100644 new mode 100755 From 74bb493ba7fa3ec5960ccf09019b2056a1f40d9d Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 10:24:52 -0800 Subject: [PATCH 10/26] Megatron-Bridge: merge Bridge main c88217d under the Megatron-LM 0756e48 pin #7798 6e77cf3 sits on MLM main past #8213 (OptimizerConfig.__post_init__ -> validate()/finalize_optimizer_config()); the old Bridge base still called super().__post_init__() and failed in OptimizerConfig.finalize(). Signed-off-by: YigongQin --- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 21efc1550b3..5c240020ac0 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 21efc1550b32256f109b9c24b086ffa99d23b646 +Subproject commit 5c240020ac0ed7468fb35b002a241ed381a18471 From bcbf2aef7aac0b5cda441045aab4d5265c57c1d9 Mon Sep 17 00:00:00 2001 From: "Shuai Zhang (Worldwide Field Ops)" Date: Fri, 23 Oct 2026 10:06:45 -0700 Subject: [PATCH 11/26] recipe: Qwen3-30B-A3B per-token NVFP4 zero-KL on FlashInfer (cherry picked from commit 8bae295a1525a0a13513e84c214dc97172c8a6cb) Signed-off-by: YigongQin --- ...noncolocated-zero-kl-nvfp4-flashinfer.yaml | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml new file mode 100644 index 00000000000..32306877f1d --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -0,0 +1,28 @@ +# Qwen3-30B-A3B zero train/generation KL with per-token NVFP4 routed experts on FlashInfer +# (SwiGLU): the BF16 zero-KL recipe plus the Nano per-token NVFP4 FlashInfer knobs. +defaults: grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.yaml +checkpointing: + checkpoint_dir: results/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer +policy: + megatron_cfg: + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls + moe_inference_training_forward: true + inference_moe_nvfp4_recipe: per_token + activation_checkpointing: true + recompute_granularity: selective + recompute_modules: + - moe + te_precision_config_file: null + fp8_cfg: + enabled: false + generation: + mcore_generation_config: + inference_grouped_gemm_backend: flashinfer + fp8_cfg: + enabled: false +logger: + log_dir: logs/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + wandb: + name: grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer From 41e8536f085f97c4b8d2b82c8cb913c422542a03 Mon Sep 17 00:00:00 2001 From: "Shuai Zhang (Worldwide Field Ops)" Date: Fri, 23 Oct 2026 10:32:10 -0700 Subject: [PATCH 12/26] megatron: inference layer spec for GPT training with moe_inference_training_forward (cherry picked from commit af560c26b4501a7f83f3a2a93747402f68e9361c) Signed-off-by: YigongQin --- nemo_rl/models/policy/workers/megatron_policy_worker.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index be2ae650589..391f97f6b4c 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -948,12 +948,13 @@ def __init__( ) self.megatron_cfg = runtime_config.megatron_cfg - if ( + if self.megatron_cfg.model.transformer_impl == "inference_optimized" and ( not init_optimizer - and self.megatron_cfg.model.transformer_impl == "inference_optimized" + or getattr(self.megatron_cfg.model, "moe_inference_training_forward", False) ): # Bridge's generic GPT provider defaults to its training layer spec. - # Dedicated inference workers need MCore inference linears instead. + # Dedicated inference workers need MCore inference linears instead, and + # so does training when its MoE forward runs the inference expert kernel. # A model-specific provider already handles this itself. _configure_inference_optimized_layer_spec(self.megatron_cfg.model) self.dtype = runtime_config.dtype From 6824b7d7b4a420dacf5d56486b698b3df45b2dd4 Mon Sep 17 00:00:00 2001 From: "Shuai Zhang (Worldwide Field Ops)" Date: Fri, 23 Oct 2026 10:59:23 -0700 Subject: [PATCH 13/26] recipe: Qwen3 NVFP4 zero-KL on the DAPO-Math 8n4g setup (cherry picked from commit 26ba31a57d45c3c9e13917b04b4b82ed2fc93368) Signed-off-by: YigongQin --- ...ation-noncolocated-zero-kl-nvfp4-dapo.yaml | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml new file mode 100644 index 00000000000..32c37e2a62f --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml @@ -0,0 +1,75 @@ +# Qwen3-30B-A3B-Base per-token NVFP4 zero train/generation KL on the DAPO-Math setup of +# grpo-qwen3-30ba3b-base-8n4g-megatron-te-nvfp4-pertoken.yaml (same data, lengths, batch, +# dynamic sampling, loss and optimizer), for comparison with the vLLM NVFP4 runs. +# 8 nodes: 4 for training (EP4), 4 for Megatron generation. +defaults: grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +grpo: + max_num_steps: 1000 + num_prompts_per_step: 32 + num_generations_per_prompt: 16 + val_at_start: true + val_period: 20 + use_dynamic_sampling: true + batch_multiplier: 3 + dynamic_sampling_max_gen_batches: 20 + use_leave_one_out_baseline: false + reward_scaling: + enabled: true + target_min: -1.0 + reward_shaping: + enabled: true + overlong_buffer_length: 512 + max_response_length: 20480 +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: true + truncated_importance_sampling_type: tis + truncated_importance_sampling_ratio: 2 + ratio_clip_max: 0.28 + ratio_clip_c: 10 +checkpointing: + checkpoint_dir: results/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo + save_period: 20 + keep_top_k: 2 + metric_name: null +policy: + model_name: Qwen/Qwen3-30B-A3B-Base + train_global_batch_size: 512 + train_micro_batch_size: 1 + max_total_sequence_length: 22528 + logprob_chunk_size: 2048 + megatron_cfg: + optimizer: + lr: 1.0e-06 + min_lr: 1.0e-07 + weight_decay: 0.1 + scheduler: + lr_warmup_iters: 10 + lr_warmup_init: 1.0e-08 + generation: + max_new_tokens: 20480 + colocated: + resources: + num_nodes: 4 + mcore_generation_config: + max_tokens: 8192 + max_requests: 512 + buffer_size_gb: 80 +data: + max_input_seq_length: 2048 + train: + dataset_name: DAPOMath17K + validation: + dataset_name: DAPOMathAIME2024 + default: + prompt_file: null +env: + math: + num_workers: 16 + math_verify_impl: dapo_math_verify +cluster: + num_nodes: 8 +logger: + log_dir: logs/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo + wandb: + name: grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo From df267755b6d9f797817d994b78beb8e431751b03 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Fri, 30 Oct 2026 16:02:19 -0700 Subject: [PATCH 14/26] Qwen3 NVFP4 zero-KL: rename to moe_match_generation_forward, enable packing in the DAPO recipe Signed-off-by: YigongQin --- ...n_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml | 2 +- ...egatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml | 7 +++++++ nemo_rl/models/policy/workers/megatron_policy_worker.py | 2 +- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index 32306877f1d..ddf2fe0571b 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -8,7 +8,7 @@ policy: transformer_impl: inference_optimized inference_grouped_gemm_backend: flashinfer inference_moe_token_dispatcher_type: nvls - moe_inference_training_forward: true + moe_match_generation_forward: true inference_moe_nvfp4_recipe: per_token activation_checkpointing: true recompute_granularity: selective diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml index 32c37e2a62f..0ac973b86a9 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml @@ -38,6 +38,13 @@ policy: train_micro_batch_size: 1 max_total_sequence_length: 22528 logprob_chunk_size: 2048 + # grpo_math_qwen30ba3b_megatron turns packing off; without it every microbatch is one + # response padded to the batch's longest. logprob_mb_tokens is set explicitly because the + # default scales with logprob_batch_size (4) to ~90K tokens. + sequence_packing: + enabled: true + train_mb_tokens: ${policy.max_total_sequence_length} + logprob_mb_tokens: ${policy.max_total_sequence_length} megatron_cfg: optimizer: lr: 1.0e-06 diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 391f97f6b4c..80f7d0a8dd1 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -950,7 +950,7 @@ def __init__( self.megatron_cfg = runtime_config.megatron_cfg if self.megatron_cfg.model.transformer_impl == "inference_optimized" and ( not init_optimizer - or getattr(self.megatron_cfg.model, "moe_inference_training_forward", False) + or getattr(self.megatron_cfg.model, "moe_match_generation_forward", False) ): # Bridge's generic GPT provider defaults to its training layer spec. # Dedicated inference workers need MCore inference linears instead, and From 2f2a32a32b3eb396206f274fc533c145dba05db4 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Mon, 2 Nov 2026 18:46:47 -0800 Subject: [PATCH 15/26] Qwen3 NVFP4 zero-KL: enable chunked prefill in generation Signed-off-by: YigongQin --- ...gatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index ddf2fe0571b..23a657fc538 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -19,6 +19,8 @@ policy: enabled: false generation: mcore_generation_config: + # The BF16 parent turns it off; batch-invariant generation stays bitwise-exact with it. + enable_chunked_prefill: true inference_grouped_gemm_backend: flashinfer fp8_cfg: enabled: false From 928a934b1b800532efac1cd6c32acd2d46d3bdc0 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 09:18:09 -0800 Subject: [PATCH 16/26] Qwen3 NVFP4 zero-KL: rename to moe_training_matches_inference Signed-off-by: YigongQin --- ...gatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml | 2 +- nemo_rl/models/policy/workers/megatron_policy_worker.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index 23a657fc538..e5a23252fbd 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -8,7 +8,7 @@ policy: transformer_impl: inference_optimized inference_grouped_gemm_backend: flashinfer inference_moe_token_dispatcher_type: nvls - moe_match_generation_forward: true + moe_training_matches_inference: true inference_moe_nvfp4_recipe: per_token activation_checkpointing: true recompute_granularity: selective diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 80f7d0a8dd1..0d62c72a5aa 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -950,7 +950,7 @@ def __init__( self.megatron_cfg = runtime_config.megatron_cfg if self.megatron_cfg.model.transformer_impl == "inference_optimized" and ( not init_optimizer - or getattr(self.megatron_cfg.model, "moe_match_generation_forward", False) + or getattr(self.megatron_cfg.model, "moe_training_matches_inference", False) ): # Bridge's generic GPT provider defaults to its training layer spec. # Dedicated inference workers need MCore inference linears instead, and From bc5b8453da675ac4b125bc48772708c12ebbe3fc Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 11:27:52 -0800 Subject: [PATCH 17/26] test_suites: Qwen3-30B-A3B per-token NVFP4 FlashInfer zero-KL 5-step suite Signed-off-by: YigongQin --- ...n-noncolocated-zero-kl-nvfp4-flashinfer.sh | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100755 tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh new file mode 100755 index 00000000000..178d2a5c0ed --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# Five-step GB200 test for zero train/generation KL on Qwen3-30B-A3B with per-token +# NVFP4 routed experts on FlashInfer. Training and refit stay BF16; each training MoE +# forward runs generation's NVFP4 kernel (moe_training_matches_inference) and the TE +# recompute pass supplies the backward. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=4 +STEPS_PER_RUN=5 +MAX_STEPS=5 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run --no-sync examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# The recipe must have switched the batch-invariant kernels on. +grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +# Zero train/generation KL means bitwise agreement: gen_kl_error is exactly 0 and +# token_mult_prob_error (max multiplicative probability error) is exactly 1 on every +# step. Both min and max are checked so a negative KL estimate cannot slip through. +uv run --no-sync tests/check_metrics.py $JSON_METRICS \ + 'min(data["train/num_valid_samples"]) > 0' \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/gen_kl_error"]) == 0' \ + 'max(data["train/gen_kl_error"]) == 0' \ + 'min(data["train/token_mult_prob_error"]) == 1' \ + 'max(data["train/token_mult_prob_error"]) == 1' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" From 918664e8bd25ba22642f50bced48015b9723b1c1 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 15:05:14 -0800 Subject: [PATCH 18/26] bump Megatron-Bridge to 2fdfbb35 (Megatron-LM 64428de2: batch-invariant a2a CP, NVIDIA/Megatron-LM#8935) Signed-off-by: YigongQin --- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 5c240020ac0..2fdfbb35db9 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 5c240020ac0ed7468fb35b002a241ed381a18471 +Subproject commit 2fdfbb35db9f7d8c7753c35adaaec1f223918b1a From b96c5bb452206fa424abd7846546c40e8121f80d Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 15:05:14 -0800 Subject: [PATCH 19/26] zero-KL: training-side a2a context parallelism with generation at CP1 Plumb policy.megatron_cfg.cp_comm_type to the model provider. Batch-invariant mode now accepts training context_parallel_size > 1 with cp_comm_type='a2a' and no longer requires training and generation CP to match. Add the Nano MXFP8 zero-KL recipe and 3n4g suite with training at TP2 x CP4 x EP4. Signed-off-by: YigongQin --- ...ration-noncolocated-zero-kl-mxfp8-cp4.yaml | 19 ++++++ nemo_rl/models/megatron/setup.py | 2 + .../megatron/zero_train_gen_mismatch.py | 12 ++-- nemo_rl/models/policy/__init__.py | 3 + ...neration-noncolocated-zero-kl-mxfp8-cp4.sh | 67 +++++++++++++++++++ tests/test_suites/nightly_gb200.txt | 1 + .../models/megatron/test_megatron_setup.py | 20 ++++++ .../megatron/test_zero_train_gen_mismatch.py | 21 +++++- 8 files changed, 140 insertions(+), 5 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml new file mode 100644 index 00000000000..c4945f048f0 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml @@ -0,0 +1,19 @@ +# The zero-kl-mxfp8 recipe with context parallelism on the training side only: +# training runs TP2 x CP4 with EP4 on two nodes, generation stays at CP=1 on the +# third. All-to-all CP gives each rank whole sequences for a quarter of the heads, +# so training still reproduces generation's log-probs bitwise. +# +# Needs Megatron-LM PR #8935 (batch-invariant a2a context parallelism). +defaults: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4 +policy: + megatron_cfg: + context_parallel_size: 4 + cp_comm_type: a2a +logger: + log_dir: logs/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4 + wandb: + name: grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4 +cluster: + num_nodes: 3 diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 1f968f68b08..4c62134d65d 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1312,6 +1312,8 @@ def _apply_parallelism_config(model_cfg: Any, config: PolicyConfig) -> None: ] model_cfg.sequence_parallel = config["megatron_cfg"]["sequence_parallel"] model_cfg.context_parallel_size = config["megatron_cfg"]["context_parallel_size"] + if "cp_comm_type" in config["megatron_cfg"]: + model_cfg.cp_comm_type = config["megatron_cfg"]["cp_comm_type"] if model_cfg.context_parallel_size > 1: # Either NeMo-RL does the packing+CP-sharding itself (classic mcore diff --git a/nemo_rl/models/megatron/zero_train_gen_mismatch.py b/nemo_rl/models/megatron/zero_train_gen_mismatch.py index 6e6c5e38708..b0900c5ef51 100644 --- a/nemo_rl/models/megatron/zero_train_gen_mismatch.py +++ b/nemo_rl/models/megatron/zero_train_gen_mismatch.py @@ -147,10 +147,15 @@ def validate_batch_invariant_mode(config: PolicyConfig) -> ZeroTrainGenValidatio "policy.megatron_cfg.tensor_model_parallel_size <= 2: at TP >= 3 " "Megatron-Core's TP reductions sum in a message-layout dependent order." ) - if megatron_cfg.get("context_parallel_size") != 1: + # Generation always runs at CP=1. Megatron-Core's a2a CP gives each rank whole + # sequences of a subset of heads, which reproduces the CP=1 forward bitwise + # (NVIDIA/Megatron-LM#8935); the other CP types split attention by sequence. + cp_size = megatron_cfg.get("context_parallel_size") or 1 + if cp_size > 1 and megatron_cfg.get("cp_comm_type") != "a2a": out.violations.append( - "batch_invariant_mode=True currently requires training context " - "parallel size 1." + "batch_invariant_mode=True with training context_parallel_size > 1 " + "requires policy.megatron_cfg.cp_comm_type='a2a' " + f"(got {megatron_cfg.get('cp_comm_type')!r})." ) token_multiple = batch_invariant_token_multiple( megatron_cfg["tensor_model_parallel_size"] @@ -234,7 +239,6 @@ def validate_batch_invariant_mode(config: PolicyConfig) -> ZeroTrainGenValidatio ) matching_fields = ( "tensor_model_parallel_size", - "context_parallel_size", "batch_invariant_mode", "batch_invariant_backend", "batch_invariant_collective", diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 83b314abc88..d22af364f6d 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -426,6 +426,9 @@ class MegatronConfig(TypedDict): num_layers_in_first_pipeline_stage: int | None num_layers_in_last_pipeline_stage: int | None context_parallel_size: int + # Megatron-Core cp_comm_type; omit to keep the model provider's value. With + # batch_invariant_mode and context_parallel_size > 1 it must be "a2a". + cp_comm_type: NotRequired[Literal["p2p", "a2a", "all_gather", "a2a+p2p"]] # Nemotron Omni RADIO/provider booleans. Omit any field to retain the model # provider's checkpoint/default value. radio_force_eval_mode: NotRequired[bool] diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh new file mode 100755 index 00000000000..b224a4c7921 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B (hybrid +# Mamba + squared-ReLU MoE) with MXFP8 linear_qkv, linear_proj and routed experts (attention +# core, shared experts and Mamba mixers BF16): Megatron training (TP2, a2a CP4, EP4) against +# non-colocated Megatron inference at CP1 (torch MXFP8 grouped GEMM, chunked prefill) with the +# batch-invariant kernel stack. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=3 +GPUS_PER_NODE=4 +STEPS_PER_RUN=5 +MAX_STEPS=5 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run --no-sync examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# The recipe must have switched the batch-invariant kernels on. +grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG + +# The per-module MXFP8 recipe must have been loaded. +grep -F -q "TE per-module precision recipe loaded from" $RUN_LOG + +# Training must have run with context parallelism; generation is always CP1. +grep -F -q "'context_parallel_size': 4" $RUN_LOG + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +# Zero train/generation KL means bitwise agreement: gen_kl_error is exactly 0 and +# token_mult_prob_error (max multiplicative probability error) is exactly 1 on every +# step. Both min and max are checked so a negative KL estimate cannot slip through. +uv run --no-sync tests/check_metrics.py $JSON_METRICS \ + 'min(data["train/num_valid_samples"]) > 0' \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/gen_kl_error"]) == 0' \ + 'max(data["train/gen_kl_error"]) == 0' \ + 'min(data["train/token_mult_prob_error"]) == 1' \ + 'max(data["train/token_mult_prob_error"]) == 1' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 03f4e402188..4f39d53fb68 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -25,6 +25,7 @@ tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh # Zero train/generation KL (batch-invariant Megatron training + Megatron inference) tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 40e9b56baac..16cd1bd3464 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -815,6 +815,26 @@ def test_context_parallel_with_sequence_packing(self): assert model_cfg.context_parallel_size == 4 + def test_cp_comm_type_is_applied_only_when_set(self): + from nemo_rl.models.megatron.setup import _apply_parallelism_config + + megatron_cfg = { + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "num_layers_in_first_pipeline_stage": None, + "num_layers_in_last_pipeline_stage": None, + "sequence_parallel": False, + "context_parallel_size": 2, + } + model_cfg = MagicMock(cp_comm_type="p2p") + config = {"megatron_cfg": dict(megatron_cfg), "sequence_packing": {"enabled": True}} + _apply_parallelism_config(model_cfg, config) + assert model_cfg.cp_comm_type == "p2p" + + config["megatron_cfg"]["cp_comm_type"] = "a2a" + _apply_parallelism_config(model_cfg, config) + assert model_cfg.cp_comm_type == "a2a" + @pytest.mark.mcore class TestApplyMultimodalConfig: diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 9befbbe2ca5..582c9176f56 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -40,6 +40,7 @@ ZERO_KL_EXEMPLARS = [ "recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.yaml", "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml", + "recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml", "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml", "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml", ] @@ -235,7 +236,10 @@ def test_batch_invariant_reports_a_train_generation_tp_mismatch(): @pytest.mark.parametrize( ("mutate", "expected"), [ - (lambda c: c["megatron_cfg"].update(context_parallel_size=2), "context"), + ( + lambda c: c["megatron_cfg"].update(context_parallel_size=2), + "cp_comm_type='a2a'", + ), ( lambda c: c["megatron_cfg"].update(use_fused_linear_logprobs=True), "use_fused_linear_logprobs", @@ -266,6 +270,21 @@ def test_batch_invariant_rejects_unsupported_settings(mutate, expected): assert any(expected in v for v in result.violations), result.violations +def test_batch_invariant_accepts_a2a_training_context_parallel(): + # Generation stays at CP=1, so the CP sizes are not required to match. + config = _policy_config() + config["megatron_cfg"].update(context_parallel_size=2, cp_comm_type="a2a") + assert zgm.validate_batch_invariant_mode(config).violations == [] + + +@pytest.mark.parametrize("cp_comm_type", ["p2p", "all_gather", "a2a+p2p"]) +def test_batch_invariant_rejects_non_a2a_context_parallel(cp_comm_type): + config = _policy_config() + config["megatron_cfg"].update(context_parallel_size=2, cp_comm_type=cp_comm_type) + result = zgm.validate_batch_invariant_mode(config) + assert any("cp_comm_type='a2a'" in v for v in result.violations), result.violations + + def test_batch_invariant_accepts_sequence_packing(): config = _policy_config() config["sequence_packing"] = {"enabled": True} From 7cc35c450b006c0181f5a835aed68caae2982a0b Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 15:23:06 -0800 Subject: [PATCH 20/26] zero-KL: run the Nano per-token NVFP4 FlashInfer recipe at training a2a CP4 (3n4g) Signed-off-by: YigongQin --- ...oncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml} | 15 +++++++++------ ...-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh} | 8 ++++++-- tests/test_suites/nightly_gb200.txt | 2 +- .../megatron/test_zero_train_gen_mismatch.py | 2 +- 4 files changed, 17 insertions(+), 10 deletions(-) rename examples/configs/recipes/llm/{grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml => grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml} (81%) rename tests/test_suites/llm/{grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh => grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh} (89%) diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml similarity index 81% rename from examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml rename to examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml index 19b86d51424..b2858ec0d1c 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml @@ -2,9 +2,10 @@ # FlashInfer. Training and refit stay BF16; generation quantizes the experts to NVFP4, # and each training MoE forward runs the same kernel on the same quantization # (moe_training_matches_inference) with the backward from the TE recompute pass. +# Training runs TP2 x a2a CP4 x EP4 on two nodes; generation stays at CP=1 on the third. # -# Needs Megatron-LM PR #7798 and flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 -# MoE kernels). +# Needs Megatron-LM PRs #7798 and #8935 (batch-invariant a2a context parallelism) and +# flashinfer-python>=0.7.0 (the per-token NVFP4 Relu2 MoE kernels). defaults: ../../grpo_math_1B.yaml grpo: max_num_steps: 5 @@ -12,7 +13,7 @@ grpo: num_generations_per_prompt: 4 checkpointing: enabled: false - checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + checkpoint_dir: results/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4 policy: model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 tokenizer: @@ -43,6 +44,8 @@ policy: TRITON_CACHE_AUTOTUNING: '1' bias_activation_fusion: false tensor_model_parallel_size: 2 + context_parallel_size: 4 + cp_comm_type: a2a expert_model_parallel_size: 4 sequence_parallel: true moe_router_dtype: fp32 @@ -81,12 +84,12 @@ policy: gpus_per_node: ${cluster.gpus_per_node} num_nodes: 1 logger: - log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + log_dir: logs/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4 wandb_enabled: true tensorboard_enabled: true wandb: project: nemo-rl - name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + name: grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4 cluster: gpus_per_node: 4 - num_nodes: 2 + num_nodes: 3 diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh similarity index 89% rename from tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh rename to tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh index aa6dd27b397..5cafe7070fa 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh @@ -2,12 +2,13 @@ # Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with # per-token NVFP4 routed experts on FlashInfer. Training and refit stay BF16; each # training MoE forward runs generation's NVFP4 kernel (moe_training_matches_inference) -# and the TE recompute pass supplies the backward. +# and the TE recompute pass supplies the backward. Training runs TP2 x a2a CP4 x EP4 on +# two nodes against generation at CP1 on the third. SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env # ===== BEGIN CONFIG ===== -NUM_NODES=2 +NUM_NODES=3 GPUS_PER_NODE=4 STEPS_PER_RUN=5 MAX_STEPS=5 @@ -38,6 +39,9 @@ uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # The recipe must have switched the batch-invariant kernels on. grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG +# Training must have run with context parallelism; generation is always CP1. +grep -F -q "'context_parallel_size': 4" $RUN_LOG + MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 4f39d53fb68..cb41bea5a36 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -27,7 +27,7 @@ tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-ze tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.sh tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh # BF16 FlashInfer TRTLLM + NCCL reshard numerical correctness tests/test_suites/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.sh diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 582c9176f56..600f251e30a 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -42,7 +42,7 @@ "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml", "recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml", "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml", - "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml", + "recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml", ] # Installed-package versions that satisfy every gate in `_validate_packages`. From fcc5454918388fdf663f02694a528fa66a3165d2 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Tue, 3 Nov 2026 15:58:00 -0800 Subject: [PATCH 21/26] zero-KL: drop the Nano MXFP8 FlashInfer recipe and suite for now Signed-off-by: YigongQin --- ...noncolocated-zero-kl-mxfp8-flashinfer.yaml | 27 -------- ...n-noncolocated-zero-kl-mxfp8-flashinfer.sh | 63 ------------------- tests/test_suites/nightly_gb200.txt | 1 - .../megatron/test_zero_train_gen_mismatch.py | 1 - 4 files changed, 92 deletions(-) delete mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml delete mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml deleted file mode 100644 index 781a8a77420..00000000000 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# Nemotron Nano v3 zero train/generation KL with MXFP8 routed experts on generation's -# FlashInfer routed MXFP8 kernel. The FlashInfer kernel does not match TE training, so -# each training MoE forward runs that kernel (moe_training_matches_inference) and the -# TE recompute pass of the MoE layer supplies the backward. -# -# Needs Megatron-LM PR #7798 on top of the zero-kl-mxfp8 parent's requirements. -defaults: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml -checkpointing: - checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer -policy: - megatron_cfg: - transformer_impl: inference_optimized - inference_grouped_gemm_backend: flashinfer - inference_moe_token_dispatcher_type: nvls - moe_training_matches_inference: true - te_precision_config_file: examples/te_precision/nanov3_experts_mxfp8.yaml - activation_checkpointing: true - recompute_granularity: selective - recompute_modules: - - moe - generation: - mcore_generation_config: - inference_grouped_gemm_backend: flashinfer -logger: - log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer - wandb: - name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh deleted file mode 100755 index c2282b21994..00000000000 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh +++ /dev/null @@ -1,63 +0,0 @@ -#!/bin/bash -# Five-step GB200 test for zero train/generation KL on Nemotron Nano v3 30B-A3B with MXFP8 -# routed experts on generation's FlashInfer routed MXFP8 kernel. Each training MoE forward -# runs that kernel (moe_training_matches_inference); the TE recompute pass supplies the -# backward. -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source $SCRIPT_DIR/common.env - -# ===== BEGIN CONFIG ===== -NUM_NODES=2 -GPUS_PER_NODE=4 -STEPS_PER_RUN=5 -MAX_STEPS=5 -NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=120 -SNAPSHOT_MEGATRON_BRIDGE=1 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -cd $PROJECT_ROOT -uv run --no-sync examples/run_grpo.py \ - --config $CONFIG_PATH \ - grpo.max_num_steps=$MAX_STEPS \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=True \ - logger.wandb.project=nemo-rl \ - logger.wandb.name=$EXP_NAME \ - logger.monitor_gpus=True \ - logger.tensorboard_enabled=True \ - checkpointing.enabled=False \ - checkpointing.checkpoint_dir=$CKPT_DIR \ - $@ \ - 2>&1 | tee $RUN_LOG - -uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -# The recipe must have switched the batch-invariant kernels on. -grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG - -# The per-module MXFP8 recipe must have been loaded. -grep -F -q "TE per-module precision recipe loaded from" $RUN_LOG - -MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) -if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then - echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" - exit 1 -fi - -# Zero train/generation KL means bitwise agreement: gen_kl_error is exactly 0 and -# token_mult_prob_error (max multiplicative probability error) is exactly 1 on every -# step. Both min and max are checked so a negative KL estimate cannot slip through. -uv run --no-sync tests/check_metrics.py $JSON_METRICS \ - 'min(data["train/num_valid_samples"]) > 0' \ - 'all_finite(data["train/loss"])' \ - 'all_finite(data["train/grad_norm"])' \ - 'min(data["train/gen_kl_error"]) == 0' \ - 'max(data["train/gen_kl_error"]) == 0' \ - 'min(data["train/token_mult_prob_error"]) == 1' \ - 'max(data["train/token_mult_prob_error"]) == 1' - -# Clean up checkpoint directory after successful run to save space. -rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index cb41bea5a36..941277e189a 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -26,7 +26,6 @@ tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-ep16tp2-mxfp8-trtllm.sh tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.sh tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.sh tests/test_suites/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.sh # BF16 FlashInfer TRTLLM + NCCL reshard numerical correctness diff --git a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py index 600f251e30a..9a9e41376aa 100644 --- a/tests/unit/models/megatron/test_zero_train_gen_mismatch.py +++ b/tests/unit/models/megatron/test_zero_train_gen_mismatch.py @@ -41,7 +41,6 @@ "recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.yaml", "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml", "recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-mxfp8-cp4.yaml", - "recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer.yaml", "recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml", ] From 701d288fa7ce371538982e384e50c8f6439b2d76 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Wed, 4 Nov 2026 14:20:55 -0800 Subject: [PATCH 22/26] update config changes Signed-off-by: YigongQin --- ...tion-noncolocated-te-nvfp4-flashinfer.yaml | 83 +++++++++++++++++++ ...noncolocated-zero-kl-nvfp4-flashinfer.yaml | 72 +++++++++++++++- ...ation-noncolocated-zero-kl-nvfp4-dapo.yaml | 2 - ...ration-noncolocated-te-nvfp4-flashinfer.sh | 54 ++++++++++++ 4 files changed, 207 insertions(+), 4 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml new file mode 100644 index 00000000000..8efde808b47 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml @@ -0,0 +1,83 @@ +# Nemotron Nano v3 with TE per-token NVFP4 routed experts in training and FlashInfer +# per-token NVFP4 routed experts in Megatron generation, without batch-invariant mode: +# train/generation KL is small but nonzero. Refit stays BF16 and generation requantizes +# the experts. Attention, Mamba and shared experts stay BF16 on both sides. +# Training runs TP2 x EP4 on one node; generation runs TP2 x EP4 on the other. +defaults: ../../grpo_math_1B.yaml +grpo: + max_num_steps: 5 + num_prompts_per_step: 8 + num_generations_per_prompt: 4 +checkpointing: + enabled: false + checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer +policy: + model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 + tokenizer: + name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + hf_config_overrides: + router_aux_loss_coef: 0 + train_global_batch_size: 32 + train_micro_batch_size: 1 + logprob_batch_size: 4 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 128 + automodel_cfg: + enabled: false + megatron_cfg: + enabled: true + env_vars: + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False + NVTE_NVFP4_ROW_SCALED_ACTIVATION: '1' + NVTE_NVFP4_DISABLE_RHT: '1' + NVTE_NVFP4_DISABLE_2D_QUANTIZATION: '1' + NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING: '1' + NVTE_BACKWARD_OVERRIDE: dequantized + bias_activation_fusion: false + tensor_model_parallel_size: 2 + expert_model_parallel_size: 4 + sequence_parallel: true + moe_router_dtype: fp32 + defer_fp32_logits: true + use_fused_weighted_squared_relu: true + fp4_cfg: + enabled: true + te_precision_config_file: examples/te_precision/attn_bf16_mlp_nvfp4.yaml + generation: + backend: megatron + refit_transport: mcore + mcore_generation_config: + refit_backend: nccl + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls + inference_moe_nvfp4_recipe: per_token + # Training's TE FP4 setup is merged into generation otherwise; the inference + # model quantizes its BF16 experts itself. + fp4_cfg: + enabled: false + te_precision_config_file: null + tensor_model_parallel_size: 2 + expert_model_parallel_size: 4 + sequence_parallel: true + moe_router_dtype: fp32 + logprobs_mode: raw_logprobs + mamba_inference_ssm_states_dtype: float32 + max_tokens: 256 + max_requests: 32 + moe_pad_experts_for_cuda_graph_inference: false + colocated: + enabled: false + resources: + gpus_per_node: ${cluster.gpus_per_node} + num_nodes: 1 +logger: + log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer +cluster: + gpus_per_node: 4 + num_nodes: 2 diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index e5a23252fbd..fb469415131 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -1,9 +1,52 @@ -# Qwen3-30B-A3B zero train/generation KL with per-token NVFP4 routed experts on FlashInfer -# (SwiGLU): the BF16 zero-KL recipe plus the Nano per-token NVFP4 FlashInfer knobs. +# Qwen3-30B-A3B-Base zero train/generation KL with per-token NVFP4 routed experts on FlashInfer +# (SwiGLU): the BF16 zero-KL recipe plus the Nano per-token NVFP4 FlashInfer knobs, on the +# DAPO-Math setup of grpo-qwen3-30ba3b-base-8n4g-megatron-te-nvfp4-pertoken.yaml. +# 2 nodes: 1 for training (EP4), 1 for Megatron generation. The batch is 1/4 of +# grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml so the +# per-GPU training and generation load matches that 8-node run. defaults: grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.yaml +grpo: + max_num_steps: 1000 + num_prompts_per_step: 8 + num_generations_per_prompt: 16 + val_at_start: false + val_period: 20 + use_dynamic_sampling: true + batch_multiplier: 3 + dynamic_sampling_max_gen_batches: 20 + use_leave_one_out_baseline: false + reward_scaling: + enabled: true + target_min: -1.0 + reward_shaping: + enabled: true + overlong_buffer_length: 512 + max_response_length: 20480 +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: true + truncated_importance_sampling_type: tis + truncated_importance_sampling_ratio: 2 + ratio_clip_max: 0.28 + ratio_clip_c: 10 checkpointing: checkpoint_dir: results/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer + save_period: 20 + keep_top_k: 2 + metric_name: null policy: + model_name: Qwen/Qwen3-30B-A3B-Base + train_global_batch_size: 128 + train_micro_batch_size: 1 + max_total_sequence_length: 22528 + logprob_chunk_size: 2048 + # grpo_math_qwen30ba3b_megatron turns packing off; without it every microbatch is one + # response padded to the batch's longest. logprob_mb_tokens is set explicitly because the + # default scales with logprob_batch_size (4) to ~90K tokens. + sequence_packing: + enabled: true + train_mb_tokens: ${policy.max_total_sequence_length} + logprob_mb_tokens: ${policy.max_total_sequence_length} megatron_cfg: transformer_impl: inference_optimized inference_grouped_gemm_backend: flashinfer @@ -17,13 +60,38 @@ policy: te_precision_config_file: null fp8_cfg: enabled: false + optimizer: + lr: 1.0e-06 + min_lr: 1.0e-07 + weight_decay: 0.1 + scheduler: + lr_warmup_iters: 10 + lr_warmup_init: 1.0e-08 generation: + max_new_tokens: 20480 mcore_generation_config: # The BF16 parent turns it off; batch-invariant generation stays bitwise-exact with it. enable_chunked_prefill: true inference_grouped_gemm_backend: flashinfer fp8_cfg: enabled: false + # FlashInfer MoE processes the full round_up(max_tokens) / TP * EP dispatcher + # buffer every step, so keep it small; ~96 rollouts land on each generation GPU. + max_tokens: 512 + max_requests: 128 + buffer_size_gb: 80 +data: + max_input_seq_length: 2048 + train: + dataset_name: DAPOMath17K + validation: + dataset_name: DAPOMathAIME2024 + default: + prompt_file: null +env: + math: + num_workers: 16 + math_verify_impl: dapo_math_verify logger: log_dir: logs/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer wandb: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml index 0ac973b86a9..c28f1c86c4a 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-base-8n4g-megatron_generation-noncolocated-zero-kl-nvfp4-dapo.yaml @@ -59,8 +59,6 @@ policy: resources: num_nodes: 4 mcore_generation_config: - max_tokens: 8192 - max_requests: 512 buffer_size_gb: 80 data: max_input_seq_length: 2048 diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.sh new file mode 100755 index 00000000000..7ea5921b21e --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.sh @@ -0,0 +1,54 @@ +#!/bin/bash +# Five-step GB200 test for Nemotron Nano v3 30B-A3B with TE per-token NVFP4 routed +# experts in training and FlashInfer per-token NVFP4 routed experts in Megatron +# generation, without batch-invariant mode. Refit stays BF16. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=4 +STEPS_PER_RUN=5 +MAX_STEPS=5 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=120 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run --no-sync examples/run_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Training must have run TE FP4. +grep -F -q "[fp4_cfg] Megatron FP4 training enabled: fp4=e2m1 recipe=nvfp4 fp4_param=False" $RUN_LOG + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +uv run --no-sync tests/check_metrics.py $JSON_METRICS \ + 'min(data["train/num_valid_samples"]) > 0' \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'all_finite(data["train/gen_kl_error"])' \ + 'max(data["train/gen_kl_error"]) < 0.05' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" From 45374c4dad64a8239b7d6687332ed17253536a4b Mon Sep 17 00:00:00 2001 From: YigongQin Date: Wed, 4 Nov 2026 15:06:11 -0800 Subject: [PATCH 23/26] NVFP4 recipes: bound decode-only MoE rows with the FlashInfer token capacity Signed-off-by: YigongQin --- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- ...gatron_generation-noncolocated-te-nvfp4-flashinfer.yaml | 3 ++- ...neration-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml | 3 ++- ...n_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml | 7 +++++-- 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 2fdfbb35db9..abde8697e7c 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 2fdfbb35db9f7d8c7753c35adaaec1f223918b1a +Subproject commit abde8697e7c3a359982b437be7c09dfcbb917207 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml index 8efde808b47..88eae3ea332 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml @@ -63,8 +63,9 @@ policy: moe_router_dtype: fp32 logprobs_mode: raw_logprobs mamba_inference_ssm_states_dtype: float32 - max_tokens: 256 + max_tokens: 2048 max_requests: 32 + inference_flashinfer_mxfp8_token_capacity: 128 moe_pad_experts_for_cuda_graph_inference: false colocated: enabled: false diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml index b2858ec0d1c..fc93a0b2693 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml @@ -74,8 +74,9 @@ policy: moe_router_dtype: fp32 logprobs_mode: raw_logprobs mamba_inference_ssm_states_dtype: float32 - max_tokens: 256 + max_tokens: 2048 max_requests: 32 + inference_flashinfer_mxfp8_token_capacity: 128 enable_prefix_caching: false moe_pad_experts_for_cuda_graph_inference: false colocated: diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml index fb469415131..f71c91a23c6 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer.yaml @@ -76,9 +76,12 @@ policy: fp8_cfg: enabled: false # FlashInfer MoE processes the full round_up(max_tokens) / TP * EP dispatcher - # buffer every step, so keep it small; ~96 rollouts land on each generation GPU. - max_tokens: 512 + # buffer on steps where any EP rank prefills; steps where every EP rank decodes + # are bounded to the token capacity, which must cover max_requests * EP. + # ~96 rollouts land on each generation GPU. + max_tokens: 8192 max_requests: 128 + inference_flashinfer_mxfp8_token_capacity: 512 buffer_size_gb: 80 data: max_input_seq_length: 2048 From 9c1de7e9756a2d020d1dacefb45cf1fa4bee8546 Mon Sep 17 00:00:00 2001 From: YigongQin Date: Wed, 4 Nov 2026 16:05:23 -0800 Subject: [PATCH 24/26] use p prefix caching as default Signed-off-by: YigongQin --- ...-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml | 5 ++++- ...generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml index 270a8e8ed49..22f42622883 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8.yaml @@ -64,7 +64,10 @@ policy: mamba_inference_ssm_states_dtype: float32 max_tokens: 256 max_requests: 32 - enable_prefix_caching: false + enable_prefix_caching: true + # Only full blocks are cached, and batch-invariant Mamba caching needs a multiple + # of the 128-token SSM chunk; 128 lets ~100-token math prompts share blocks. + block_size_tokens: 128 moe_pad_experts_for_cuda_graph_inference: false fp8_cfg: enabled: true diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml index fc93a0b2693..24a8d94b105 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml @@ -77,7 +77,10 @@ policy: max_tokens: 2048 max_requests: 32 inference_flashinfer_mxfp8_token_capacity: 128 - enable_prefix_caching: false + enable_prefix_caching: true + # Only full blocks are cached, and batch-invariant Mamba caching needs a multiple + # of the 128-token SSM chunk; 128 lets ~100-token math prompts share blocks. + block_size_tokens: 128 moe_pad_experts_for_cuda_graph_inference: false colocated: enabled: false From cf6ceed547c67aa3aedc623b536790679a5e3f5d Mon Sep 17 00:00:00 2001 From: YigongQin Date: Thu, 5 Nov 2026 11:13:28 -0800 Subject: [PATCH 25/26] zero kl example for gym single controller Signed-off-by: YigongQin --- ...oken_capture-zero-kl-nvfp4-flashinfer.yaml | 84 +++++++++++++++++ .../policy/workers/megatron_policy_worker.py | 6 +- ...-token_capture-zero-kl-nvfp4-flashinfer.sh | 89 +++++++++++++++++++ .../policy/test_megatron_split_state.py | 10 +-- 4 files changed, 182 insertions(+), 7 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.yaml create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.yaml new file mode 100644 index 00000000000..dec32609435 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.yaml @@ -0,0 +1,84 @@ +# Zero train/generation KL sibling of the SingleController token-capture nightly, with +# per-token NVFP4 routed experts on FlashInfer. Training and refit stay BF16; generation +# quantizes the experts to NVFP4, and each training MoE forward runs the same kernel on +# the same quantization (moe_training_matches_inference). Training runs TP2 x EP4 on one +# node; generation runs TP2 x EP4 on the other. +# +# Needs Megatron-LM PRs #7798 and #8691 and flashinfer-python>=0.7.0 (the per-token +# NVFP4 Relu2 MoE kernels). +# Launch with: +# uv run examples/run_grpo_single_controller.py \ +# --config examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.yaml +defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-single-controller-token_capture.yaml + +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer + +policy: + hf_config_overrides: + router_aux_loss_coef: 0 + make_sequence_length_divisible_by: 128 + megatron_cfg: + zero_train_gen_mismatch: true + batch_invariant_mode: true + batch_invariant_backend: te_native + batch_invariant_collective: ordered + attention_backend: flash + flash_attention_version: 4 + moe_permute_fusion: false + env_vars: + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False + CUBLASLT_WORKSPACE_SIZE: '0' + CUBLAS_WORKSPACE_CONFIG: :0:0 + MAMBA_DETERMINISTIC: '1' + TRITON_CACHE_AUTOTUNING: '1' + tensor_model_parallel_size: 2 + context_parallel_size: 1 + expert_model_parallel_size: 4 + sequence_parallel: true + moe_router_dtype: fp32 + defer_fp32_logits: true + use_fused_weighted_squared_relu: true + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls + inference_moe_nvfp4_recipe: per_token + moe_training_matches_inference: true + activation_checkpointing: true + recompute_granularity: selective + recompute_modules: + - moe + generation: + mcore_generation_config: + transformer_impl: inference_optimized + inference_grouped_gemm_backend: flashinfer + inference_moe_token_dispatcher_type: nvls + tensor_model_parallel_size: 2 + expert_model_parallel_size: 4 + sequence_parallel: true + moe_router_dtype: fp32 + logprobs_mode: raw_logprobs + mamba_inference_ssm_states_dtype: float32 + # 8 prompts x 16 generations over two TP2 engines is ~64 requests per engine. + max_requests: 64 + inference_flashinfer_mxfp8_token_capacity: 256 + moe_pad_experts_for_cuda_graph_inference: false + colocated: + resources: + gpus_per_node: ${cluster.gpus_per_node} + num_nodes: 1 + +# Rollouts from an older policy version cannot match the training forward, so stay +# strictly on-policy. +async_rl: + sampler: + max_lookahead_versions: 0 + +logger: + log_dir: logs/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer + wandb: + name: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer + +cluster: + gpus_per_node: 4 + num_nodes: 2 diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 0d62c72a5aa..7b2528bfa63 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1319,7 +1319,9 @@ def _normalize_in_loss_seq_filter( # Finish any overlap on the comm stream before touching the reduced # gradient buffers (including distributed-optimizer gradient shards). torch.cuda.synchronize() - self.model.scale_gradients(token_factor) + self.model.scale_gradients( + torch.tensor(token_factor, device=kept_toks.device) + ) metrics = [ rescale_loss_metrics( m, @@ -2302,7 +2304,7 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: # below sees the rescaled grads; for all_reduce the result is the # global mean grad; for reduce_scatter (dist-opt) it's the shard. # Either way, opt.step sees the right-normalized gradient. - self.model.scale_gradients(inv_n) + self.model.scale_gradients(torch.tensor(inv_n, device=n_safe.device)) if draft_step_state.active: draft_step_state.correct_main_grads( self.model.parameters(), diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.sh new file mode 100755 index 00000000000..17a24b4784b --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-async-gym-single-controller-token_capture-zero-kl-nvfp4-flashinfer.sh @@ -0,0 +1,89 @@ +#!/bin/bash +# Zero train/generation KL on the SingleController token-capture path, with per-token +# NVFP4 routed experts on FlashInfer, on GB200. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=4 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=150 +SNAPSHOT_MEGATRON_BRIDGE=1 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Token capture is only wired up on the SingleController path. +uv run --no-sync examples/run_grpo_single_controller.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + checkpointing.save_period=3 \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run --no-sync tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# The recipe must have switched the batch-invariant kernels on. +grep -F -q "[zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4" $RUN_LOG + +MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS) +if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then + echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP" + exit 1 +fi + +# Zero train/generation KL means bitwise agreement: gen_kl_error is exactly 0 and +# token_mult_prob_error is exactly 1 on every step. The finalize/* gates are the +# capture health signal: a rejected or poisoned rollout means capture or recovery lost +# a rollout's tokens. +uv run --no-sync tests/check_metrics.py $JSON_METRICS \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/gen_kl_error"]) == 0' \ + 'max(data["train/gen_kl_error"]) == 0' \ + 'min(data["train/token_mult_prob_error"]) == 1' \ + 'max(data["train/token_mult_prob_error"]) == 1' \ + 'max(data["train/finalize/invalid_row_rate"]) == 0' \ + 'max(data["train/finalize/capture_poisoned_rollouts"]) == 0' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/unit/models/policy/test_megatron_split_state.py b/tests/unit/models/policy/test_megatron_split_state.py index c24344321c0..719ba977b1a 100644 --- a/tests/unit/models/policy/test_megatron_split_state.py +++ b/tests/unit/models/policy/test_megatron_split_state.py @@ -681,7 +681,7 @@ def test_rescales_grads_with_inv_n(self, mock_module_symbols): w.finish_train_step() # scale_gradients should have been called with some 1/N scalar < 1 w.model.scale_gradients.assert_called_once() - arg = w.model.scale_gradients.call_args.args[0] + arg = float(w.model.scale_gradients.call_args.args[0]) assert 0 < arg <= 1.0 def test_restores_first_step_param_sync_after_success(self, mock_module_symbols): @@ -791,7 +791,7 @@ def test_picks_global_valid_toks_for_token_level_loss(self, mock_module_symbols) w.finish_train_step() # local_valid_toks accumulated = 2048; with mocked all_reduce as no-op, # global_valid_toks == 2048 → inv_n = 1/2048 - arg = w.model.scale_gradients.call_args.args[0] + arg = float(w.model.scale_gradients.call_args.args[0]) assert arg == pytest.approx(1.0 / 2048.0, rel=1e-4) def test_picks_global_valid_seqs_for_sequence_level_loss(self, mock_module_symbols): @@ -800,7 +800,7 @@ def test_picks_global_valid_seqs_for_sequence_level_loss(self, mock_module_symbo w = self._setup_open_step(mock_module_symbols, LossType.SEQUENCE_LEVEL) w.finish_train_step() # local_valid_seqs = 8 → inv_n = 1/8 - arg = w.model.scale_gradients.call_args.args[0] + arg = float(w.model.scale_gradients.call_args.args[0]) assert arg == pytest.approx(1.0 / 8.0, rel=1e-4) @pytest.mark.parametrize("token_level_loss", [False, True]) @@ -1012,7 +1012,7 @@ def test_mtp_grads_use_token_denominator_under_sequence_level_loss( w.finish_train_step() # The main loss gets 1/8; MTP additionally gets 8/2048, for net 1/2048. - assert w.model.scale_gradients.call_args.args[0] == pytest.approx(1.0 / 8.0) + assert float(w.model.scale_gradients.call_args.args[0]) == pytest.approx(1.0 / 8.0) expected_mtp_grad = torch.full((4,), 8.0 / 2048.0) torch.testing.assert_close(mtp_param.main_grad, expected_mtp_grad) torch.testing.assert_close(grad_seen_by_finalize[0], expected_mtp_grad) @@ -1766,7 +1766,7 @@ def zero_grad_buffers() -> None: for buffer in buffers: buffer.grad_data.zero_() - def scale_gradients(scale: float) -> None: + def scale_gradients(scale: torch.Tensor) -> None: for buffer in buffers: buffer.grad_data.mul_(scale) From b369a948ecf9f376b6436590eb371837ec54fd99 Mon Sep 17 00:00:00 2001 From: Yigong Qin Date: Thu, 5 Nov 2026 14:00:09 -0800 Subject: [PATCH 26/26] fix refit hang Signed-off-by: Yigong Qin --- .../models/generation/megatron/megatron_worker.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 45bd5f411b9..95e2e1a6129 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -1645,10 +1645,16 @@ def init_collective_mcore_generation( nccl_store, global_rank, world_size, nccl_options ) nccl_backend._set_sequence_number_for_group() - # Create the group-wide NCCL communicator now, on every rank. - nccl_backend.eager_connect_single_device( - torch.device("cuda", torch.cuda.current_device()) - ) + # NCCLCopyService connects on its first run, including idle ranks. + # Connecting here as well can reuse stale scalable-rendezvous IDs + # in large refit groups and hang with connection refused. M2N has + # no first-run hook, so it still connects eagerly. + # Set NRL_MCORE_NCCL_SINGLE_INIT=0 to restore eager connection. + single_init = os.environ.get("NRL_MCORE_NCCL_SINGLE_INIT", "1") != "0" + if refit_backend == "nccl_m2n" or not single_init: + nccl_backend.eager_connect_single_device( + torch.device("cuda", torch.cuda.current_device()) + ) pg._register_backend( torch.device("cuda"), ProcessGroup.BackendType.NCCL,