Skip to content

feat(megatron): support selective MXFP8 inference parameters - #4126

Draft
wdykas wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
wdykas:feat/megatron-selective-mxfp8-inference
Draft

wdykas wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
wdykas:feat/megatron-selective-mxfp8-inference

Conversation

@wdykas

@wdykas wdykas commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

Summary

Align NeMo RL selective MXFP8 generation with the construction-time TE recipe design in NVIDIA/Megatron-LM#7300. No batch-invariance changes.

  • Use the existing policy.megatron_cfg.te_precision_config_file for training and Megatron generation. Remove this PR's obsolete inference parameter regex filters and refit keyword arguments.
  • Forward first_last_layers_bf16, num_layers_at_start_in_bf16, and num_layers_at_end_in_bf16 to the provider. These previously were not applied by NeMo RL setup.
  • Apply the precision recipe and boundary settings before colocated inference model construction, just as dedicated workers do.
  • Keep refit storage-based: BF16 destinations stay BF16; Torch, FlashInfer, and vLLM-selected MXFP8 use MCore's canonical layout. Extend the existing mixed-precision/idempotent refit test to vLLM.
  • Document the routed-experts-only recipe, BF16 first 2 / last 4 layers, and backend limitations.

MCore handles parameter selection before checkpoint loading and automatically inherits the enclosing MXFP8 init context where appropriate. Users do not need a checkpoint helper or an explicit inherit_model_init_context flag. The recipe keeps shared experts, attention, and MTP in BF16.

Dependency

Requires NVIDIA/Megatron-LM#7300 and a NeMo RL environment containing that MCore revision. Keep this draft until the MCore change and dependency update land. The current pinned MCore is not sufficient for the new selective-storage behavior.

Validation

Using this PR's NeMo RL checkout with MCore 8faece7a713e53e6a85b54cff42c8d5083897d61, cached Torch 2.11.0+cu130 / TE 2.18.0+27486e03 environment on GB200:

  • Precision/config propagation: 17 passed.
  • MXFP8/M-to-N refit checks: 16 passed.
  • Config loading, validation, and reference configs: 746 passed, 6 skipped.
  • Ruff 0.9.9 and changed type-checked files: passed.
  • Full pyrefly reports three errors in unchanged nemo_rl/weight_sync/xferdtensor.py against the cached NCCL M-to-N dependency. This PR does not modify that file; this is not a full clean-environment CI pass.

End-to-end smoke test passed (exit 0): Qwen/Qwen3-30B-A3B on 4 GB200 GPUs, TP1/EP4, colocated Megatron generation with native in-place refit, vLLM-selected grouped GEMM, NVLS, global MXFP8 parameter storage, routed-experts-only TE recipe, first 2 / last 4 layers BF16, and eager execution. Both short GRPO steps completed rollout, policy logprobs, backward, and optimizer step.

Reported generation KL errors were 0.0039 / 0.0011. The smoke run had zero reward/loss; this is execution coverage, not a convergence, performance, or zero-logprob-parity claim. Batch invariance and R3 router replay were disabled. Nano and non-colocated M-to-N end-to-end runs were not repeated for this PR.

Signed-off-by: wdykas <wdykas@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 14, 2026

Copy link
Copy Markdown

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 github-actions Bot added the Documentation Improvements or additions to documentation label Sep 14, 2026
@wdykas

wdykas commented Sep 14, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b64a8e3

@wdykas wdykas added the CI:L1 Run doctests, unit tests, and functional tests label Sep 14, 2026
Signed-off-by: wdykas <wdykas@nvidia.com>
@wdykas

wdykas commented Sep 15, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test de7b76c

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L1 Run doctests, unit tests, and functional tests Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant