feat(megatron): NVFP4 0 train-gen mismatch on flashinfer backend by YigongQin · Pull Request #4791 · NVIDIA-NeMo/RL

YigongQin · 2026-10-19T17:42:46Z

What does this PR do ?

Extend policy.megatron_cfg.zero_train_gen_mismatch to FlashInfer MoE experts: per-token NVFP4 (NVFP4 weights × per-token NVFP4 activations) and MXFP8. Train and generation are still bitwise identical (train/gen_kl_error == 0, train/token_mult_prob_error == 1) in a synchronous GRPO loop.

The FlashInfer expert kernel is batch-invariant but not bitwise identical to TE training, so the previous preset rejected it. Megatron-Core's new moe_inference_training_forward makes training use it instead. Each training MoE forward runs generation's FlashInfer kernel, on the same quantization rebuilt from the live parameters on every call. The backward comes from a TE recompute of the MoE layer (a straight-through estimator). For NVFP4, training weights and refit stay BF16; only generation's expert GEMMs (and the matching training forward) run NVFP4.

Builds on #4683 (Nano MXFP8 zero-KL, merged). Shipped: a Nemotron Nano v3 per-token NVFP4 recipe and nightly, and an MXFP8-on-FlashInfer recipe and nightly. Both nightlies stay in disabled.txt until the submodule pins below are on upstream main.

Important

This needs an unmerged Megatron-LM PR. For reproduction, .gitmodules points the Bridge submodule at the fork branch YigongQin/Megatron-Bridge:yigongq/nano-minf-training-forward, whose Megatron-LM submodule is YigongQin/Megatron-LM:yigongq/minf-training-forward (d9b0ba0, the head of NVIDIA/Megatron-LM#7798). Both must move back to the upstream repos before merge.

Megatron-LM d9b0ba0 = NVIDIA/Megatron-LM#7798 on top of main at 9314b01, which already has #8262 (per-token NVFP4 squared-ReLU MoE), #7562, #8195 and #6076.
Megatron-Bridge 7c984ce = upstream main (f73818d) + one commit: Megatron-LM pin d9b0ba0 and flashinfer-python/flashinfer-cubin 0.7.0.post1.

Links

Builds on (merged): Nano MXFP8 zero-KL, NVIDIA-NeMo/RL#4683
Merged BF16 zero-KL preset: NVIDIA-NeMo/RL#3208
Tracking issue: Training-generation mismatch #3743
Required: FlashInfer train-gen parity and NVFP4 refit: NVIDIA/Megatron-LM#7798
Already on Megatron-LM main: #8262 (per-token NVFP4 FlashInfer MoE), #7562 (hybrid MoE parity), #8195 (fp64 expert combine), #6076 (deterministic Triton autotune)

What FlashInfer zero-KL needs

Everything the Nano MXFP8 zero-KL preset requires still applies. On top of that:

Requirement	Where it is set	Checked?
Megatron-LM with #7798	submodule pins	no. An older Megatron-Core rejects the config keys at model build.
`flashinfer-python` / `flashinfer-cubin` 0.7.0.post1	NeMo-RL `mcore` extra and Bridge `pyproject.toml`, in lockstep	no
MXFP8 on FlashInfer only: `flashinfer-jit-cache` installed	not in the `mcore` extra (see below)	by Megatron-Core at MoE layer build
`moe_inference_training_forward=true` whenever generation uses FlashInfer with MXFP8 or NVFP4 experts	recipe	yes
`transformer_impl=inference_optimized` and `inference_grouped_gemm_backend=flashinfer` on train and generation	recipe	yes
MoE recompute: `activation_checkpointing`, `recompute_granularity=selective`, `moe` in `recompute_modules`; `cuda_graph_impl` not `local`	recipe	yes. The backward comes from the recompute.
Same `inference_moe_nvfp4_recipe`, `inference_moe_token_dispatcher_type` and `expert_model_parallel_size` on both sides	recipe	yes
MXFP8 experts or `inference_moe_nvfp4_recipe=per_token`	recipe	yes. The FlashInfer BF16 kernel is not batch invariant.
NVFP4: FP8 off on both sides (`fp8_cfg` disabled, no `te_precision_config_file`)	NVFP4 recipe	no
NVLS symmetric-heap bound `moe_inference_training_max_tokens_per_rank`	derived unless set	derived from the largest train or logprob micro-batch at the padded sequence length, or the packing / dynamic-batching `mb_tokens`

flashinfer-jit-cache 0.7.0 is a meta-package over per-architecture wheels (flashinfer-jit-cache-sm80 … sm121a) on the explicit flashinfer-cu130 index, which uv only uses for direct dependencies mapped to it. Rather than list and map all of them, the mcore extra drops flashinfer-jit-cache. Per-token NVFP4 doesn't need it (Megatron-Core skips the check for its CuTe DSL JIT kernels), but the FlashInfer MXFP8 path does, so the MXFP8-on-FlashInfer recipe can't run in the default mcore environment yet.

What changed

nemo_rl/models/megatron/zero_train_gen_mismatch.py:
- FlashInfer with MXFP8 or NVFP4 experts is allowed when moe_inference_training_forward is on, and _validate_inference_training_forward checks the requirements above.
- Dedicated generation workers validate with generation_only=True. Their merged config clears the training-only flag, so the "requires moe_inference_training_forward" rule is left to the driver, which sees the full config.
nemo_rl/models/megatron/setup.py: passes inference_moe_nvfp4_recipe, moe_inference_training_forward and the derived token bound to the model config.
nemo_rl/models/generation/megatron/config.py: merged_inference_megatron_cfg clears moe_inference_training_forward for generation.
nemo_rl/models/policy/__init__.py: the three new MegatronConfig keys; te_precision_config_file accepts null.
nemo_rl/models/policy/workers/megatron_policy_worker.py: passes generation_only=is_refit_destination.
pyproject.toml / uv.lock: FlashInfer 0.7.0.post1 in the mcore extra, without flashinfer-jit-cache; the lock also picks up the Bridge bump (nccl-extensions from Megatron-LM, Pillow 13).
Recipes and nightlies: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer (on the Nano MXFP8 zero-KL recipe) and ...-zero-kl-nvfp4-flashinfer (directly on grpo_math_1B.yaml), both in tests/test_suites/disabled.txt for now.
Unit tests for the new validator rules, the generation-only path and the token bound.

Validation

Unit: tests/unit/models/megatron/test_zero_train_gen_mismatch.py, tests/unit/models/megatron/test_megatron_setup.py.
Functional (GB200, 2 nodes × 4 GPUs, 5 sync GRPO steps, Nemotron Nano v3 per-token NVFP4 nightly, at the pins above): job 1309408 passed (W&B). train/gen_kl_error was exactly 0, and train/token_mult_prob_error, train/probs_ratio_min and train/probs_ratio_max were exactly 1 on all 5 steps.
The MXFP8-on-FlashInfer nightly has not run at these pins; it needs flashinfer-jit-cache (above).

copy-pr-bot · 2026-10-19T17:42:49Z

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.

github-actions · 2026-10-20T17:50:30Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 4c76a57 (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-10-21T01:42:40Z

✅ Submodule Fast-Forward Check Results

Check based on commit: c53e4d9 (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-10-26T19:41:56Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 381b69f (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-10-26T23:18:54Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 19d04be (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-10-27T01:21:12Z

✅ Submodule Fast-Forward Check Results

Check based on commit: e1669aa (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Signed-off-by: YigongQin <yigongq@nvidia.com>

Bridge 7c984ce = Bridge main f73818d + Megatron-LM #7798 (d9b0ba0) + FlashInfer 0.7.0.post1. Signed-off-by: YigongQin <yigongq@nvidia.com>

Signed-off-by: YigongQin <yigongq@nvidia.com>

github-actions · 2026-10-27T20:22:38Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 3c33b9f (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

…perts-only recipe, FlashInfer 0.7.2.post1 + jit-cache, enable FlashInfer nightlies Signed-off-by: YigongQin <yigongq@nvidia.com>

… 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 <yigongq@nvidia.com>

Signed-off-by: YigongQin <yigongq@nvidia.com>

…e48 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 <yigongq@nvidia.com>

(cherry picked from commit 8bae295) Signed-off-by: YigongQin <yigongq@nvidia.com>

…aining_forward (cherry picked from commit af560c2) Signed-off-by: YigongQin <yigongq@nvidia.com>

(cherry picked from commit 26ba31a) Signed-off-by: YigongQin <yigongq@nvidia.com>

…acking in the DAPO recipe Signed-off-by: YigongQin <yigongq@nvidia.com>

Signed-off-by: YigongQin <yigongq@nvidia.com>

…suite Signed-off-by: YigongQin <yigongq@nvidia.com>

…nt a2a CP, NVIDIA/Megatron-LM#8935) Signed-off-by: YigongQin <yigongq@nvidia.com>

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 <yigongq@nvidia.com>

…2a CP4 (3n4g) Signed-off-by: YigongQin <yigongq@nvidia.com>

Signed-off-by: YigongQin <yigongq@nvidia.com>

github-actions · 2026-11-04T00:19:21Z

✅ Submodule Fast-Forward Check Results

Check based on commit: fcc5454 (PR #4791 from yigongq/nano-minf-training-forward)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Signed-off-by: YigongQin <yigongq@nvidia.com>

…apacity Signed-off-by: YigongQin <yigongq@nvidia.com>

Signed-off-by: YigongQin <yigongq@nvidia.com>

Signed-off-by: Yigong Qin <yigongq@oci-hsg-cs-001-vscode-01.cm.cluster>

YigongQin mentioned this pull request Oct 20, 2026

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

Open

6 of 9 tasks

YigongQin mentioned this pull request Oct 21, 2026

feat(megatron): Support 0 train-gen mismatch for MXFP8/Mamba/TP2 #4683

Merged

YigongQin force-pushed the yigongq/nano-minf-training-forward branch from c53e4d9 to 381b69f Compare October 26, 2026 19:40

github-actions Bot added the Documentation Improvements or additions to documentation label Oct 26, 2026

YigongQin force-pushed the yigongq/nano-minf-training-forward branch from 381b69f to 19d04be Compare October 26, 2026 23:17

YigongQin changed the title ~~feat(megatron): MXFP8/NVFP4 zero train-gen mismatch on Megatron flashinfer backend~~ feat(megatron): MXFP8/NVFP4 0 train-gen mismatch on flashinfer backend Oct 26, 2026

YigongQin added 6 commits October 27, 2026 13:18

pertoken nvfp4 zero kl recipe

de1095a

Signed-off-by: YigongQin <yigongq@nvidia.com>

bump flashinfer to 0.7.0.post1 for per-token NVFP4

140a178

Signed-off-by: YigongQin <yigongq@nvidia.com>

pipeclean pertoken nvfp4 changes

96bc310

Signed-off-by: YigongQin <yigongq@nvidia.com>

pin Bridge/Megatron-LM forks for per-token NVFP4 zero-KL

152726e

Bridge 7c984ce = Bridge main f73818d + Megatron-LM #7798 (d9b0ba0) + FlashInfer 0.7.0.post1. Signed-off-by: YigongQin <yigongq@nvidia.com>

uv lock for Bridge 7c984ce (Megatron-LM #7798, FlashInfer 0.7.0.post1)

fd140c3

Signed-off-by: YigongQin <yigongq@nvidia.com>

recipe: base NVFP4 FlashInfer zero-KL recipe on grpo_math_1B

3c33b9f

Signed-off-by: YigongQin <yigongq@nvidia.com>

YigongQin force-pushed the yigongq/nano-minf-training-forward branch from e1669aa to 3c33b9f Compare October 27, 2026 20:20

kahyunnam mentioned this pull request Oct 29, 2026

[RL] FlashInfer MoE kernel as the training forward for zero train/rollout mismatch (NeMo-RL #4791, Megatron-LM #7798, #7830) flashinfer-ai/flashinfer#7939

Open

YigongQin and others added 9 commits October 30, 2026 15:46

match_generation_forward rename, release value-pass weights, MXFP8 ex…

c975434

…perts-only recipe, FlashInfer 0.7.2.post1 + jit-cache, enable FlashInfer nightlies Signed-off-by: YigongQin <yigongq@nvidia.com>

test_suites: restore the executable bit on the FlashInfer zero-KL suites

570d61d

Signed-off-by: YigongQin <yigongq@nvidia.com>

recipe: Qwen3-30B-A3B per-token NVFP4 zero-KL on FlashInfer

bcbf2ae

(cherry picked from commit 8bae295) Signed-off-by: YigongQin <yigongq@nvidia.com>

megatron: inference layer spec for GPT training with moe_inference_tr…

41e8536

…aining_forward (cherry picked from commit af560c2) Signed-off-by: YigongQin <yigongq@nvidia.com>

recipe: Qwen3 NVFP4 zero-KL on the DAPO-Math 8n4g setup

6824b7d

(cherry picked from commit 26ba31a) Signed-off-by: YigongQin <yigongq@nvidia.com>

Qwen3 NVFP4 zero-KL: rename to moe_match_generation_forward, enable p…

df26775

…acking in the DAPO recipe Signed-off-by: YigongQin <yigongq@nvidia.com>

Qwen3 NVFP4 zero-KL: enable chunked prefill in generation

2f2a32a

Signed-off-by: YigongQin <yigongq@nvidia.com>

YigongQin added 2 commits November 3, 2026 12:02

Qwen3 NVFP4 zero-KL: rename to moe_training_matches_inference

928a934

Signed-off-by: YigongQin <yigongq@nvidia.com>

test_suites: Qwen3-30B-A3B per-token NVFP4 FlashInfer zero-KL 5-step …

bc5b845

…suite Signed-off-by: YigongQin <yigongq@nvidia.com>

YigongQin changed the title ~~feat(megatron): MXFP8/NVFP4 0 train-gen mismatch on flashinfer backend~~ feat(megatron): NVFP4 0 train-gen mismatch on flashinfer backend Nov 3, 2026

YigongQin added 4 commits November 3, 2026 15:05

bump Megatron-Bridge to 2fdfbb35 (Megatron-LM 64428de2: batch-invaria…

918664e

…nt a2a CP, NVIDIA/Megatron-LM#8935) Signed-off-by: YigongQin <yigongq@nvidia.com>

zero-KL: run the Nano per-token NVFP4 FlashInfer recipe at training a…

7cc35c4

…2a CP4 (3n4g) Signed-off-by: YigongQin <yigongq@nvidia.com>

zero-KL: drop the Nano MXFP8 FlashInfer recipe and suite for now

fcc5454

Signed-off-by: YigongQin <yigongq@nvidia.com>

YigongQin added 4 commits November 4, 2026 14:20

update config changes

701d288

Signed-off-by: YigongQin <yigongq@nvidia.com>

NVFP4 recipes: bound decode-only MoE rows with the FlashInfer token c…

45374c4

…apacity Signed-off-by: YigongQin <yigongq@nvidia.com>

use p prefix caching as default

9c1de7e

Signed-off-by: YigongQin <yigongq@nvidia.com>

zero kl example for gym single controller

cf6ceed

Signed-off-by: YigongQin <yigongq@nvidia.com>

terrykong added the CI:L1 Run doctests, unit tests, and functional tests label Nov 5, 2026

fix refit hang

b369a94

Signed-off-by: Yigong Qin <yigongq@oci-hsg-cs-001-vscode-01.cm.cluster>

Navigation Menu

Repository navigation

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

feat(megatron): NVFP4 0 train-gen mismatch on flashinfer backend#4791
YigongQin wants to merge 26 commits into
mainfrom
yigongq/nano-minf-training-forward

YigongQin commented Oct 19, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 21, 2026

Uh oh!

github-actions Bot commented Oct 26, 2026

Uh oh!

github-actions Bot commented Oct 26, 2026

Uh oh!

github-actions Bot commented Oct 27, 2026

Uh oh!

github-actions Bot commented Oct 27, 2026

Uh oh!

github-actions Bot commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

YigongQin commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Links

What FlashInfer zero-KL needs

What changed

Validation

Uh oh!

copy-pr-bot Bot commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Oct 21, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Oct 26, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Oct 26, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Oct 27, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Oct 27, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Nov 4, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

YigongQin commented Oct 19, 2026 •

edited

Loading