[YigongQin]

What does this PR do ?

Add policy.megatron_cfg.zero_train_gen_mismatch: a gate that validates, and then enables, a configuration that puts Megatron training and Megatron Inference (M-Inf) generation on the same batch-invariant kernel stack, so that in a synchronous (on-policy) GRPO loop the generation logprobs and the training-side logprobs of the same tokens are bitwise identical (train/gen_kl_error == 0, train/token_mult_prob_error == 1).

The recipe sets every required value; the module only checks them (one ValueError lists every violation) and switches the MCore batch-invariant kernels on at worker init.

Scope of this PR: BF16, Blackwell, TP=1/CP=1, non-colocated Megatron generation (inference_optimized). See "Support status" for exactly what is and is not covered.

Issues

• Training-generation mismatch Tracking: zero train-gen KL on Megatron (BF16 / MXFP8 / NVFP4) #3743 - zero train/gen mismatch

Support status

	Status
Verified by a run	Qwen3-30B-A3B (MoE), BF16, GB200: 1 training node (TP1/PP1/CP1, EP4, ETP1) + 1 dedicated generation node (inference_optimized, CUDA graphs on). 5 GRPO steps: train/gen_kl_error exactly 0 and train/token_mult_prob_error exactly 1 at every step (W&B).
Required in the recipe, enforced by the validator	megatron_cfg: 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 CUBLASLT_WORKSPACE_SIZE=0 and CUBLAS_WORKSPACE_CONFIG=:0:0. Generation: backend=megatron, merged transformer_impl=inference_optimized, logprobs_mode=raw_logprobs, enable_chunked_prefill=false, and temperature=1.0, top_p=1.0, top_k=null (raw logprobs are scored).
Also enforced	Blackwell (SM100+); policy.precision=bfloat16 and no FP8 on either side; tensor_model_parallel_size=1 and context_parallel_size=1; sequence_packing off; use_fused_linear_logprobs off; make_sequence_length_divisible_by (and dynamic_batching.sequence_length_round when enabled) a multiple of MCore's batch-invariant token multiple; train and generation agree on TP/CP and all batch_invariant_*, attention_backend, flash_attention_version fields. The resolved model provider must not be MLA, an experimental attention variant, or an SSM/Mamba hybrid. The driver runs the config-only checks; workers also check packages and the GPU.
Packages required on workers	MCore with the batch_invariant_* / flash_attention_version config fields, transformer_engine>=2.18, flash_attn>=2.8.1, flash-attn-4 (pinned 4.0.0b19), nvidia-cutlass-dsl>=4.6.0.dev0.
Not supported / not validated	Colocated generation (see Known issues). FP8 / MXFP8. TP>1 and CP>1. Sequence packing and dynamic batching. MLA and SSM-hybrid models (e.g. Nemotron-Nano-v3 is not covered by this PR). transformer_impl=transformer_engine generation. Hopper / FA3. Dense models and other MoE models have not been run.
Not expected to be zero by design	Async GRPO with trajectory age > 0 or in-flight weight updates: gen_kl_error compares generation logprobs with the trainer's current-weights logprobs, so policy lag shows up as KL even with bit-exact kernels. Use the synchronous loop.

What changed

• nemo_rl/models/megatron/zero_train_gen_mismatch.py (new): validate_zero_train_gen_kl(config, *, check_environment) validates the train + generation configs (plus validate_batch_invariant_mode for the batch-invariant requirements, and validate_zero_train_gen_model_provider for the resolved architecture), and enable_batch_invariant_kernels turns the MCore kernels on and lets an installed flash-attn-4 below MCore's version gate be used (the pinned b19 vs MCore's >=b20 flag).
• nemo_rl/models/megatron/setup.py: enable_batch_invariant_mode() (validate with check_environment=True, then enable the kernels), the provider check, and plumbing of the batch-invariant fields onto the MCore model config. flash_attention_version defaults to 2 when unset.
• nemo_rl/distributed/model_utils.py: when the TP group has a single rank, scoring uses the same fused log_softmax as Megatron Inference instead of the distributed max/exp/sum/log decomposition (mathematically equal, not bitwise equal). This closes the log-softmax gap between train and generation for TP=1.
• nemo_rl/models/megatron/batch_invariant.py + nemo_rl/models/policy/lm_policy.py: align the eager scoring token count with MCore's batch-invariant token rounding so TE kernels see the same shape class as generation buckets. The cuBLAS workspace variables are required in the recipe's env_vars and checked by the validator.
• nemo_rl/models/generation/megatron/megatron_worker.py: raw logprobs under batch-invariant mode and no fused RoPE.
• nemo_rl/models/policy/__init__.py, generation/megatron/config.py: new config fields (batch_invariant_*, flash_attention_version, zero_train_gen_mismatch, generation transformer_impl / logprobs_mode).
• Base configs (grpo_math_1B, ppo_math_1B, sft, dpo, distillation_math, vlm_mpo) and their reference configs pin flash_attention_version: 2: flash-attn-4 is now installed, so TE would otherwise pick FA4. Zero-KL recipes set 4 explicitly.
• pyproject.toml / uv.lock: flash-attn-4==4.0.0b19 (b20+ need apache-tvm-ffi>=0.1.12, but the shared override keeps 0.1.11).
• Recipe examples/configs/recipes/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.yaml and GB200 nightly test tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.sh (listed in nightly_gb200.txt).

Usage

Use the recipe above as the template: it sets every required value. The switch and the main settings:

policy:
  precision: bfloat16
  sequence_packing: {enabled: false}
  dynamic_batching: {enabled: false}
  make_sequence_length_divisible_by: 64
  megatron_cfg:
    zero_train_gen_mismatch: true
    tensor_model_parallel_size: 1
    context_parallel_size: 1
    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:
      CUBLASLT_WORKSPACE_SIZE: "0"
      CUBLAS_WORKSPACE_CONFIG: ":0:0"
  generation:
    backend: megatron
    temperature: 1.0
    top_p: 1.0
    top_k: null
    colocated: {enabled: false}
    mcore_generation_config:
      transformer_impl: inference_optimized
      logprobs_mode: raw_logprobs
      enable_chunked_prefill: false

Run the nightly test with tools/launch tests/test_suites/llm/grpo-qwen3-30ba3b-2n4g-megatron_generation-noncolocated-zero-kl.sh (needs a worker environment that has flash-attn-4).

Validation

• Functional (nightly test): the script asserts the kernel-enabled marker line ([zero_train_gen_mismatch] batch-invariant kernels enabled: backend=te_native collective=ordered flash_attention_version=4), finite loss / grad norm, min == max == 0 for train/gen_kl_error and min == max == 1 for train/token_mult_prob_error. Passing run: Qwen3-30B-A3B, 2 nodes x 4 GB200, 5 steps, 32 valid samples per step, all checks exact.
• Unit: tests/unit/models/megatron/test_zero_train_gen_mismatch.py (resolver and validator, including that the shipped recipe validates, that the driver check does not touch the GPU, and that TP>1 / packing / FP8 / Hopper are rejected). CI passes.

Known issues / follow-ups

• Colocated Megatron generation is not supported yet. A trial on 1 node x 4 GPUs ran end to end (second inference model + MCore in-place reshard), was bit-exact for steps 1-3, then 1-2 of 32 sequences diverged from step 4 (gen_kl_error 0.35 / 0.42, token_mult_prob_error ~1e15) while the rest stayed exact. Root cause not identified; leading suspect is the colocated refit / sleep-wake path (MCore swap_model_weights and the persisted inference state), not the kernels. The recipe is not shipped here. (Colocated on one node also needs the optimizer kept on the GPU to avoid a host-RAM OOM.)
• FP8 / MXFP8 and SSM-hybrid models (Nemotron-Nano-v3) are follow-ups.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[shanmugamr1992]

@YigongQin can you elaborate more on this MR ? maybe attach some sample runs that you did , how the performance changes etc.

[rybakov]

zero mismatch between training and rollouts is important feature for RL debugging, please submit for review

[ZhiyuLi-Nvidia]

There are lots of patches do you think we can minimize these and only keep the ones are really necessary?

Otherwise it would be really hard to maintain in the long run.

┌─────┬──────────────────────────────────────────────────────────────────────────┬─────────┬────────────────────────────────────────────────────────────────┐
│  #  │                                  Target                                  │  Owner  │                              Goal                              │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 1   │ te…cpp_extensions.gemm.get_cublas_workspace_size_bytes → 4               │ TE      │ Starve cuBLAS workspace → force splitK=1 (fixed reduction)     │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 2   │ model_utils._compute_distributed_log_softmax_with_grad                   │ NeMo-RL │ TP=1 use F.log_softmax to match gen raw_logprobs               │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 3   │ moe_utils.unpermute + token_dispatcher.unpermute                         │ mcore   │ Batch-invariant gather combine (the core idea)                 │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 4   │ moe_utils.topk_routing_with_score_function + moe_router.…                │ mcore   │ index_put_ routing maps instead of scatter                     │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 5   │ TextGenerationController._dynamic_step_context_bookkeeping               │ mcore   │ Reconstruct routing before update_requests frees KV blocks     │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 6   │ DynamicInferenceEngine.async_bookkeep                                    │ mcore   │ Apply that routing; drop upstream's finished_routing_block_ids │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 7   │ CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list │ mcore   │ Force graph buckets to 64-multiples                            │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 8   │ DynamicInferenceContext.__setattr__                                      │ mcore   │ Round runtime decode padding to the same 64 quantum            │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 9   │ DynamicInferenceContext.initialize_attention_state                       │ mcore   │ Track graph-capture window                                     │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 10  │ CUDAGraphBatchDimensionBuilder.match_graph_config                        │ mcore   │ Suppress bucket match when graphs are off                      │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 11  │ MegatronGenerationMixin.prepare_for_generation / .finish_generation      │ NeMo-RL │ Track when inference graphs toggle on/off                      │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 12  │ MegatronGenerationMixin._sleep                                           │ NeMo-RL │ Clear expert padding before scoring                            │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 13  │ MambaMixer._ssm_prefill / ._ssm_decode + 5 injected helpers              │ mcore   │ Align train/prefill/decode SSM scan                            │
├─────┼──────────────────────────────────────────────────────────────────────────┼─────────┼────────────────────────────────────────────────────────────────┤
│ 14  │ megatron.core.utils.get_mamba_version                                    │ mcore   │ Avoid importing mamba_ssm (tvm crash on py3.13)                │
└─────┴──────────────────────────────────────────────────────────────────────────┴─────────┴────────────────────────────────────────────────────────────────┘

Plus non-patch global mutations: RouterReplay.global_router_replay_instances (rebuild), metadata.routing_indices_buffer = None, cg_utils.cuda_graph_attr_cache.pop(), and three env vars (CUBLAS_WORKSPACE_CONFIG=:0:0, CUBLASLT_WORKSPACE_SIZE=0, MAMBA_DETERMINISTIC=1).

Grouped by goal:

• Batch invariance (1, 2, 3, 4) — mcore's batch_invariant_kernels.py already owns 1 and 2; 3–4 are the genuine gap
• CUDA-graph shape pinning (7, 8, 9, 10) — all downstream of the unexplained 64
• Router replay plumbing (5, 6) — works around blocks being freed before serialization
• Mamba (13, 14) — 14 is an unrelated py3.13 import workaround riding along
• First-party lifecycle (11, 12) — patching our own code; should be hooks

Three observations that fall out of the table: 3 of 14 patch NeMo-RL's own code (2, 11, 12), which should be plain hooks or edits; #14 is unrelated to zero-KL and belongs in its own PR; and only 3, 4, 5 are genuinely novel — the rest is scaffolding around them.

[ZhiyuLi-Nvidia]

Zero-KL config is validated and built inside workers after GPU allocation; suggest building + validating once in setup() before the cluster.

[ZhiyuLi-Nvidia]

flash-attn-4 in the mcore extra switches default TE flash attention to FA4; suggest pinning flash_attention_version: 2 in base configs and 4 only in the zero-KL recipe.

[ZhiyuLi-Nvidia]

Batch-invariant mode silently ignores temperature/top_p/top_k; suggest validating sampling config instead of overwriting.

[ZhiyuLi-Nvidia]

/ok to test 9877780

[ZhiyuLi-Nvidia]

/ok to test 72d1b84

[ZhiyuLi-Nvidia]

Thank you @YigongQin. LGTM!

[ZhiyuLi-Nvidia]

/ok to test b6bd602

[ZhiyuLi-Nvidia]

/ok to test de3330e