perf: Reduce MXFP8 refit weight-transfer overheads by seonjinn · Pull Request #3294 · NVIDIA-NeMo/RL

seonjinn · 2026-07-20T19:32:09Z

Summary

Reduce BF16-to-MXFP8 refit work for colocated synchronous RL by quantizing eligible Megatron weights before transfer. vLLM receives each E4M3 weight with its matching *_scale_from_checkpoint. Unsupported grouped layouts and incompatible refit configurations fail explicitly.

The path is opt-in via policy.generation.vllm_cfg.refit_prequantize. Qwen synchronous MXFP8 performance recipes enable it; asynchronous NCCL Reshard recipes leave it off.

Performance

Historical Qwen3-30B-A3B prequant-only A/B on 4 nodes / 16 GB200 GPUs (one 20-step job per variant, steps 3-19; control, prequant only):

Metric	Off	On	Change
E2E step time	160.349 s	157.892 s	-1.53%
Total refit	10.413 s	8.180 s	-21.44%
Transfer + update	7.909 s	4.859 s	-38.57%

These measurements predate the latest-main merge and current HEAD. They are directional, not a current-HEAD performance claim. A matched current-HEAD A/B remains outstanding.

Validation

OCI-HSG GB200: 275 focused vLLM, 20 MXFP8 prequantization, 33 recipe, and 3 targeted MCore tests passed.
Ruff lint and format checks passed for all 50 Python files changed by the PR.
A broader vLLM run had 288 passing tests and one external HF model-download error (HTTP 401, expired authentication). The focused offline vLLM suite above passed; the broader run did not pass overall.

copy-pr-bot · 2026-07-20T19:32:13Z

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.

seonjinn · 2026-07-20T22:16:38Z

Isolated-branch reproduction (follow-up promised in the PR body).

Setup: Qwen3-30B-A3B GRPO (OpenMathInstruct-2), 2 nodes x 8 B200, grpo-qwen3-30ba3b-4n4g-mxfp8-rollout.yaml, 20 steps, means over steps 3-20. Control = current main behavior (all flags off); treatment = this branch with refit_prequantize, persistent IPC buffers (4 GB), slim offload, cached loaders, pinned reference swap (batched shuffle is default-on). Same cluster, back-to-back runs.

steps 3-20 mean	flags off (main)	this PR	delta
transfer_and_update_weights (s)	9.67	2.98	-69%
prepare_for_generation total (s)	17.2	13.1	-4.1 s
generation / logprob / training (s)	49.8 / 84.6 / 91.0	51.8 / 82.0 / 90.6	unchanged (noise)
total step (s)	256.3	250.9	-2.1%
train reward	0.528	0.529	=

The saving is confined to the refit phase, exactly as claimed; all compute phases and reward are untouched. NRL_MXFP8_SHUFFLE_VERIFY=1 bit-exact assert passed on first refit.

seonjinn · 2026-07-20T22:23:32Z

/ok to test 740c825

seonjinn · 2026-07-21T00:39:41Z

/ok to test 3f60f5d

seonjinn · 2026-07-21T01:46:12Z

Isolated reproduction at 235B scale (follow-up to the 30B numbers above).

Setup: Qwen3-235B-A22B GRPO, 8 nodes x 8 B200 (64 GPUs), grpo-qwen3-235b-16n4g-mxfp8-rollout.yaml, vLLM TP 8, 20 steps, means over steps 3-20. Control = same clean branch with all refit flags off; treatment = refit_prequantize + persistent IPC buffers + slim offload + cached loaders + pinned reference swap (batched shuffle default-on).

steps 3-20 mean	flags off (main)	this PR	delta
transfer_and_update_weights (s)	21.8	7.4	-66%
prepare_for_generation total (s)	42.2	28.8	-13.4 s

The refit-transfer saving grows with model size (30B: 9.7 -> 3.0 s; 235B: 21.8 -> 7.4 s), because larger weights spend proportionally more of the refit in transport - exactly the phase this PR targets. Batched-shuffle bit-exact verify passed on first refit.

seonjinn · 2026-07-21T22:00:19Z

/ok to test 04af034

seonjinn · 2026-07-21T22:56:02Z

/ok to test e06f7a8

seonjinn · 2026-07-21T23:04:01Z

/ok to test 3ebdf9d

Reduce MXFP8 and ModelOpt refit overhead while preserving transport and checkpoint-engine lifecycle correctness. Signed-off-by: sna <sna@nvidia.com>

seonjinn · 2026-07-21T23:16:10Z

/ok to test b88aa2c

seonjinn · 2026-07-22T04:23:57Z

/ok to test 438e20d

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

seonjinn · 2026-07-27T18:49:08Z

/ok to test 96ba299

seonjinn · 2026-07-27T19:00:41Z

/ok to test 96ba299

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

seonjinn · 2026-10-16T19:09:28Z

/ok to test f555a77

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

seonjinn · 2026-10-16T23:00:04Z

/ok to test d30a813

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

seonjinn · 2026-10-18T03:25:53Z

/ok to test 5e451db

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

seonjinn · 2026-10-18T03:30:54Z

/ok to test 7bc2096

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

seonjinn · 2026-10-18T05:07:00Z

/ok to test ecb9aa0

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

seonjinn · 2026-10-18T06:17:24Z

/ok to test 501b5d5

copy-pr-bot · 2026-10-18T06:17:27Z

/ok to test 501b5d5

@seonjinn, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

seonjinn · 2026-10-18T06:18:09Z

/ok to test 501b5d5

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

seonjinn · 2026-10-18T07:01:27Z

/ok to test b1a8060

seonjinn requested review from a team as code owners July 20, 2026 19:32

seonjinn self-assigned this Jul 20, 2026

seonjinn changed the title ~~perf: Reduce MXFP8 refit weight-transfer overheads (prequantize, persistent buffers, batched MoE shuffle)~~ perf: Reduce MXFP8 refit weight-transfer overheads Jul 20, 2026

seonjinn added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 21, 2026

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

perf(vllm): optimize quantized refit paths

b88aa2c

Reduce MXFP8 and ModelOpt refit overhead while preserving transport and checkpoint-engine lifecycle correctness. Signed-off-by: sna <sna@nvidia.com>

seonjinn force-pushed the sna/pr-mxfp8-refit-optimization branch from 3ebdf9d to b88aa2c Compare July 21, 2026 23:16

Merge branch 'main' into sna/pr-mxfp8-refit-optimization

438e20d

Merge branch 'main' into sna/pr-mxfp8-refit-optimization

96ba299

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

seonjinn requested a review from terrykong July 27, 2026 19:01

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

terrykong previously approved these changes Oct 16, 2026

View reviewed changes

terrykong enabled auto-merge (squash) October 16, 2026 00:02

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

seonjinn added 2 commits October 16, 2026 12:08

Merge upstream/main into sna/pr-mxfp8-refit-optimization

9e2ab4b

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

Fix MXFP8 CI test collection

f555a77

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

seonjinn dismissed terrykong’s stale review via f555a77 October 16, 2026 19:09

terrykong previously approved these changes Oct 16, 2026

View reviewed changes

seonjinn added 2 commits October 16, 2026 15:59

Merge upstream/main into sna/pr-mxfp8-refit-optimization

21a70a6

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

Install vLLM extra for GB200 MXFP8 CI tests

d30a813

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

seonjinn dismissed terrykong’s stale review via d30a813 October 16, 2026 22:59

Merge latest main and repair MXFP8 GB200 L1 config

5e451db

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

Match pinned Ruff formatter for refit test

7bc2096

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

ZhiyuLi-Nvidia previously approved these changes Oct 18, 2026

View reviewed changes

Fix MXFP8 checkpoint handshake unit fixture after main merge

ecb9aa0

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

seonjinn dismissed ZhiyuLi-Nvidia’s stale review via ecb9aa0 October 18, 2026 05:06

Update MXFP8 refit test for source metadata API

501b5d5

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

test(refit): exercise prequant handshake through real worker

b1a8060

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

terrykong approved these changes Oct 18, 2026

View reviewed changes

terrykong merged commit c74be94 into NVIDIA-NeMo:main Oct 18, 2026
332 of 359 checks passed

yuki-97 mentioned this pull request Oct 19, 2026

cp: cherry-pick #4277 - docs(automodel): rename the DTensor backend to Automodel #4790

Merged

Conversation

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

Uh oh!

Summary

Performance

Validation

Uh oh!

copy-pr-bot Bot commented Jul 20, 2026

Uh oh!

seonjinn commented Jul 20, 2026

Uh oh!

seonjinn commented Jul 20, 2026

Uh oh!

seonjinn commented Jul 21, 2026

Uh oh!

seonjinn commented Jul 21, 2026

Uh oh!

seonjinn commented Jul 21, 2026

Uh oh!

seonjinn commented Jul 21, 2026

Uh oh!

seonjinn commented Jul 21, 2026

Uh oh!

seonjinn commented Jul 21, 2026

Uh oh!

seonjinn commented Jul 22, 2026

Uh oh!

seonjinn commented Jul 27, 2026

Uh oh!

seonjinn commented Jul 27, 2026

Uh oh!

seonjinn commented Oct 16, 2026

Uh oh!

seonjinn commented Oct 16, 2026

Uh oh!

seonjinn commented Oct 18, 2026

Uh oh!

seonjinn commented Oct 18, 2026

Uh oh!

seonjinn commented Oct 18, 2026

Uh oh!

seonjinn commented Oct 18, 2026

Uh oh!

copy-pr-bot Bot commented Oct 18, 2026

Uh oh!

seonjinn commented Oct 18, 2026

Uh oh!

seonjinn commented Oct 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

seonjinn commented Jul 20, 2026 •

edited

Loading