feat(refit): export native MXFP8 training weights by seonjinn · Pull Request #3908 · NVIDIA-NeMo/RL

seonjinn · 2026-08-29T03:38:20Z

Summary

Megatron can keep trained parameters in native MXFP8 storage when fp8_param=true, while refit previously expected one ordinary tensor per parameter. This PR sends Transformer Engine's live E4M3 values and E8M0 scales through the ordered component contract from #3907 without dequantizing the weight.

Design

Validate native MXFP8 storage against fp8_cfg and te_precision_config_file.
Preserve first/last BF16 layers and keep co-trained MTP parameters on the normal tensor path.
Ask Megatron-Bridge #5917 to build globally ordered PP/VP/EP-aware export tasks.
Synchronize every optimizer-owned model chunk before refit, including overlap-buffer configurations.
Restore shared MXFP8 buffers before policy and reference logprob passes.

Dependencies

feat(refit): support ordered NCCL Reshard components #3907: ordered refit components.
Megatron-Bridge #5917: native MXFP8 export tasks.
fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs #3281 and feat(megatron): support te precision config recipes #3654: optimizer and TE precision configuration plumbing.

Validation

Unit coverage includes values/scales, mixed BF16/MXFP8 storage, MTP exclusion, optimizer synchronization, TE-recipe precedence, and Bridge delegation.
Targeted Megatron-Bridge Linux/GB200 suite: 288 passed, 2 unrelated Apex-dependent cases deselected.
Qwen3.5-35B-A3B completed 20/20 Async steps with gen_kl_error mean 0.001264 (range 0.000838-0.001549).
Latest Megatron-Bridge head completed a 2/2-step GB200 smoke with gen_kl_error=0.00130165 at step 2. The nightly used the NCCL Reshard exact-transfer Python fallback, so this is correctness evidence rather than native-op performance evidence.
20-step W&B | latest-head smoke

Signed-off-by: seonjinn <sna@nvidia.com>

copy-pr-bot · 2026-08-29T03:38:23Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-12T03:02:00Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 4ece17c (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

…0-20260922 Signed-off-by: seonjinn <sna@nvidia.com> # Conflicts: # nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py

…3908-mcore7300-20260922 Signed-off-by: seonjinn <sna@nvidia.com> # Conflicts: # nemo_rl/models/megatron/setup.py # tests/unit/models/megatron/test_megatron_setup.py # tests/unit/models/policy/test_megatron_worker.py # uv.lock

github-actions · 2026-09-22T14:04:00Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 628f55b (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T14:46:21Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 7900726 (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T14:50:47Z

✅ Submodule Fast-Forward Check Results

Check based on commit: d3eeda6 (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T15:19:23Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 94bd46d (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

…0-20260922 Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T15:22:34Z

✅ Submodule Fast-Forward Check Results

Check based on commit: de34cbe (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

…0-20260922 Signed-off-by: seonjinn <sna@nvidia.com>

…3908-mcore7300-20260922 Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T15:30:13Z

✅ Submodule Fast-Forward Check Results

Check based on commit: e8fa878 (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T16:39:50Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 2e4203b (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T17:29:29Z

✅ Submodule Fast-Forward Check Results

Check based on commit: f77a586 (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T18:15:31Z

✅ Submodule Fast-Forward Check Results

Check based on commit: dc3d61c (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T19:04:52Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 3c575fc (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

seonjinn · 2026-09-22T19:53:29Z

Exact-head GB200 validation completed on 3c575fcd with the pinned Megatron-Bridge/Megatron-LM/Transformer Engine revisions.

Dependency pins: 11 passed
NVRx async-checkpoint compatibility: 5 passed
Native MXFP8 producer and transfer: 63 passed
First/last BF16 scope and NCCL component preservation: 2 passed
Public Megatron-Bridge export API: 2 passed
Real MCore selective MXFP8 parameter construction: 3 passed

All 86 selected tests passed. The run used one GB200 node and the current NeMo-RL nightly image.

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions · 2026-09-22T20:50:09Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 9ebba9a (PR #3908 from sna/native-mxfp8-megatron-source)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

seonjinn added 25 commits August 28, 2026 20:21

test(refit): define ordered component contract

835885c

Signed-off-by: seonjinn <sna@nvidia.com>

refactor(refit): add ordered weight components

ed66cbf

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): enforce native component ordering

9b75eb1

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): cover NCCL plan validation

7fc328e

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): validate plans before NCCL setup

01d6a83

Signed-off-by: seonjinn <sna@nvidia.com>

style(refit): format component transport changes

a517a8e

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): define native MXFP8 source extraction

77cfef6

Signed-off-by: seonjinn <sna@nvidia.com>

feat(refit): extract native MXFP8 components

1d8f951

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): require pinned MXFP8 dtype metadata

11c625f

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): validate pinned MXFP8 dtype metadata

43975e2

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): reject forged MXFP8 dtype bindings

8fe8aa5

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): bind MXFP8 dtype to transformer engine

beda248

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): define native MXFP8 source mapping

9db6597

Signed-off-by: seonjinn <sna@nvidia.com>

feat(refit): expose native MXFP8 source components

6bfa619

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): cover native MXFP8 review failures

bd6f890

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): repair native MXFP8 grouped sources

eebf9f7

Signed-off-by: seonjinn <sna@nvidia.com>

test(refit): cover mixed native BF16 producer storage

985aca4

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): partition native MXFP8 sources by storage

e5b9f28

Signed-off-by: seonjinn <sna@nvidia.com>

fix(megatron): restore MXFP8 shared buffers for logprobs

0c8f2c8

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): synchronize native MXFP8 params before export

f30cee6

Signed-off-by: seonjinn <sna@nvidia.com>

fix(megatron): resolve TE recipes from repo root

176217a

Signed-off-by: seonjinn <sna@nvidia.com>

fix(megatron): preserve BF16 tail with per-module recipes

71ba25a

Signed-off-by: seonjinn <sna@nvidia.com>

chore(megatron): type precision recipe settings

50a5814

Signed-off-by: seonjinn <sna@nvidia.com>

chore: add refit test copyright header

c42272e

Signed-off-by: seonjinn <sna@nvidia.com>

chore: add refit test copyright header

a369cb2

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added the CI:megatron label Aug 29, 2026

seonjinn added 2 commits August 28, 2026 20:43

refactor(megatron): rely on precision recipe loader

3c5586f

Signed-off-by: seonjinn <sna@nvidia.com>

fix(refit): honor disabled FP8 export

4293d8c

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn mentioned this pull request Aug 29, 2026

feat(refit): load native MXFP8 weights into vLLM #3909

Open

seonjinn added CI:L1 Run doctests, unit tests, and functional tests CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:megatron CI:L1 Run doctests, unit tests, and functional tests labels Sep 17, 2026

seonjinn added 2 commits September 22, 2026 15:53

Merge remote-tracking branch 'origin/main' into codex/pr3907-mcore730…

efeaa2d

…0-20260922 Signed-off-by: seonjinn <sna@nvidia.com> # Conflicts: # nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py

build(mcore): align native MXFP8 dependencies

7900726

Signed-off-by: seonjinn <sna@nvidia.com>

build(te): use runtime torch for source builds

d3eeda6

Signed-off-by: seonjinn <sna@nvidia.com>

build(te): declare wheel for source builds

94bd46d

Signed-off-by: seonjinn <sna@nvidia.com>

Merge remote-tracking branch 'origin/main' into codex/pr3908-mcore730…

de34cbe

…0-20260922 Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added 2 commits September 22, 2026 17:28

Merge remote-tracking branch 'origin/main' into codex/pr3907-mcore730…

1fa9c1a

…0-20260922 Signed-off-by: seonjinn <sna@nvidia.com>

Merge commit '1fa9c1a7574addf0f5faed2f0518c96d354a2235' into codex/pr…

e8fa878

…3908-mcore7300-20260922 Signed-off-by: seonjinn <sna@nvidia.com>

fix(mcore): support NVRx-only async checkpoint API

2e4203b

Signed-off-by: seonjinn <sna@nvidia.com>

test(mxfp8): initialize overlap config in worker fixture

f77a586

Signed-off-by: seonjinn <sna@nvidia.com>

test(mxfp8): use public conversion task contract

dc3d61c

Signed-off-by: seonjinn <sna@nvidia.com>

test(mxfp8): complete precision config fixtures

3c575fc

Signed-off-by: seonjinn <sna@nvidia.com>

test(mxfp8): identify grouped storage in cache fixture

9ebba9a

Signed-off-by: seonjinn <sna@nvidia.com>

Conversation

seonjinn commented Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Dependencies

Validation

Uh oh!

copy-pr-bot Bot commented Aug 29, 2026

Uh oh!

github-actions Bot commented Sep 12, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

seonjinn commented Sep 22, 2026

Uh oh!

github-actions Bot commented Sep 22, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

seonjinn commented Aug 29, 2026 •

edited

Loading