feat(megatron): NVFP4 0 train-gen mismatch on flashinfer backend

.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

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

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-te-nvfp4-flashinfer.yaml

@@ -0,0 +1,84 @@
+# 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: 2048
+      max_requests: 32
+      inference_flashinfer_mxfp8_token_capacity: 128
+      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

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

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

examples/configs/recipes/llm/grpo-nanov3-30BA3B-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4.yaml

@@ -0,0 +1,99 @@
+# 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_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 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
+  num_prompts_per_step: 8
+  num_generations_per_prompt: 4
+checkpointing:
+  enabled: false
+  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:
+    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
+    context_parallel_size: 4
+    cp_comm_type: a2a
+    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:
+    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
+      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: 2048
+      max_requests: 32
+      inference_flashinfer_mxfp8_token_capacity: 128
+      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
+      resources:
+        gpus_per_node: ${cluster.gpus_per_node}
+        num_nodes: 1
+logger:
+  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-3n4g-megatron_generation-noncolocated-zero-kl-nvfp4-flashinfer-cp4
+cluster:
+  gpus_per_node: 4
+  num_nodes: 3