Skip to content

perf: Reduce MXFP8 refit weight-transfer overheads - #3294

Open
seonjinn wants to merge 69 commits into
NVIDIA-NeMo:mainfrom
seonjinn:sna/pr-mxfp8-refit-optimization
Open

perf: Reduce MXFP8 refit weight-transfer overheads#3294
seonjinn wants to merge 69 commits into
NVIDIA-NeMo:mainfrom
seonjinn:sna/pr-mxfp8-refit-optimization

Conversation

@seonjinn

@seonjinn seonjinn commented Jul 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

Reduce repeated BF16-to-MXFP8 refit work for colocated Sync RL.

  • Quantize eligible trainer weights before the Sync RL transfer.
  • Reuse CUDA IPC buffers, validated vLLM loader routes, and MXFP8 shuffle scratch buffers.
  • Keep the existing path as the default when the optimization is disabled.

The Qwen Sync MXFP8 performance recipes enable the optimization. Async recipes explicitly leave it disabled; Async MXFP8 uses the NCCL Reshard path configured by #3865.

Safety

The prequantized path exchanges metadata before transfer and requires each E4M3 tensor to include its matching *_scale_from_checkpoint. It rejects incompatible trainer FP8 storage and NCCL Reshard configurations.

Results

Historical integrated B200 measurements:

Model Transfer + update Total refit
Qwen3-30B-A3B 9.67 to 2.98 s (-69.2%) 17.2 to 13.1 s (-23.8%)
Qwen3-235B-A22B 18.40 to 7.26 s (-60.5%) 40.2 to 30.4 s (-24.4%)
Nemotron Nano 3 3.56 to 1.69 s (-52.6%) 13.2 to 11.6 s (-12.0%)

These measurements include the batched MXFP8 MoE shuffle now provided by #3478. A matched current-main A/B is still needed to isolate the remaining changes in this PR.

Validation

  • python3 -m pytest -q tests/test_mxfp8_rollout_recipes.py (22 passed)
  • Unit coverage for the prequantization handshake, wire format, incompatible configuration rejection, IPC buffer reuse, loader-route caching, and MXFP8 MoE layout handling

@seonjinn
seonjinn requested review from a team as code owners July 20, 2026 19:32
@copy-pr-bot

copy-pr-bot Bot commented Jul 20, 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.

@seonjinn seonjinn self-assigned this Jul 20, 2026
@seonjinn 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

seonjinn commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

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

seonjinn commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 740c825

@seonjinn 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

seonjinn commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 3f60f5d

@seonjinn

seonjinn commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

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

seonjinn commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 04af034

@seonjinn 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
@seonjinn

seonjinn commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test e06f7a8

@seonjinn

seonjinn commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 3ebdf9d

@seonjinn
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
seonjinn force-pushed the sna/pr-mxfp8-refit-optimization branch from 3ebdf9d to b88aa2c Compare July 21, 2026 23:16
@seonjinn

seonjinn commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b88aa2c

@seonjinn

seonjinn commented Jul 22, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 438e20d

@seonjinn

seonjinn commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 96ba299

@seonjinn

seonjinn commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 96ba299

@seonjinn
seonjinn requested a review from terrykong July 27, 2026 19:01
@seonjinn 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
@seonjinn
fix: resolve CI unit-test failures after main merge
c162cfd 
- Guard offload_after_refit against configs without megatron_cfg
- Skip fp8 module import in prepare_refit_info for non-FP8 refits so
  stubbed quant-backend tests can prepare refit info
- Drop duplicate MXFP8 scale clamp already done in quantize_mxfp8_weight
- Fix prequantized-load test to compare tensor contents (reshape breaks
  object identity) and pin the noncolocated PPO mock's refit negotiation

Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
Merge remote branch update, keeping local merge resolution and CI fixes
382fa8b 
Signed-off-by: seonjinn <sna@nvidia.com>
@github-actions github-actions Bot added the CI Relating to CI label Aug 26, 2026
@seonjinn

seonjinn commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 382fa8b

@seonjinn
Merge main into sna/pr-mxfp8-refit-optimization
125f178 
Signed-off-by: seonjinn <sna@nvidia.com>
@github-actions github-actions Bot removed the CI Relating to CI label Aug 26, 2026
@seonjinn

seonjinn commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 125f178

@seonjinn

seonjinn commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test e28239c

@seonjinn
Fix pyrefly no-matching-overload false positive on masked_fill
232da17 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn

seonjinn commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 232da17

@seonjinn
test: update refit cleanup fixture after main merge
6b2ab6e 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
refactor(refit): limit optimized recipes to sync RL
06b6fbb 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn seonjinn mentioned this pull request Aug 27, 2026
Open
@seonjinn
Merge remote-tracking branch 'origin/main' into codex/pr3294-push-202…
7c9b881 
…60904

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

# Conflicts:
#	nemo_rl/algorithms/grpo.py
#	nemo_rl/weight_sync/collective_weight_synchronizer.py
#	nemo_rl/weight_sync/interfaces.py
#	tests/unit/models/generation/test_vllm_fp8_quantization.py
@seonjinn
fix(mxfp8): handle mixed grouped expert refit
bf03429 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): align refit fixtures with vllm
d9246a9 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): model routed expert config
443a9f9 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): reroute grouped MXFP8 scale sidecars
b8159de 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): patch modular kernel factory
8d38b28 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): use block-aligned hidden size
658d385 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): exercise unpadded kernel reuse
172d346 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): use current refit loader contract
546d346 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(mxfp8): update remaining loader fixtures
db15416 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn

seonjinn commented Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

Validation update on db15416d (based on main at 4601ba2c):

Exact PR head

  • MXFP8 quantization unit file: 56 passed, 1 deselected.
  • Separately selected Ray V1 initialization case: 1 passed. The main unit run covers Ray V2.
  • DCO, copyright, secrets, semantic-title, triage, and submodule checks pass.

Combined mixed-refit compatibility run

This run includes PR #3294 plus the dependent mixed BF16/MXFP8 refit changes; it is not a PR #3294-only performance result.

  • Qwen3.5 35B-A3B, 16x GB200, colocated sync CUDA IPC, CUDA Graph, first 3 / last 5 layers BF16, middle-layer QKVO and routed experts MXFP8: 20/20 steps completed.
  • Steps 2-19: 334.55 E2E tokens/s/GPU, 2,386.89 generation tokens/s/GPU, 243.85 s E2E step time, 9.72 s total refit.
  • gen_kl_error: mean 0.001732, range 0.001550-0.001860. No refit, shape, manifest, or OOM error occurred.
  • W&B run

This also exercises the grouped MXFP8 scale-sidecar routing added in this update. The old path failed on the first Qwen refit; this run completed all repeated refits.

seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 5, 2026
@seonjinn
test(mxfp8): repair variant C test signatures
ac30224 
Fixes six gate-0 failures introduced by porting tests from a newer
variant onto NVIDIA-NeMo#3294's fp8.py:

- test_mxfp8_moe_checkpoint_scales_survive_layerwise_reload: stub
  ModelWeightParameter's TP rank/world-size before allocating the
  fake layer's parameters. The earlier stub call ran only inside
  _allocate_moe_checkpoint_scales, after the parameters were already
  built, so parallel_state's uninitialised group tripped construction.
- test_checkpoint_refit_preserves_nonsharded_fp8_path: variant C's
  fp8.load_weights takes (weights, model_runner); drop the stale
  model_load_weights kwarg assertion.
- test_unquantized_nccl_reshard_keeps_existing_refit_lifecycle and
  test_fp8_flashinfer_trtllm_keeps_existing_refit_lifecycle: variant
  C's _weight_update_lifecycle has no _maybe_process_fp8_kv_cache
  hook, so drop the mock and its assertion.
- test_prepare_refit_info_reports_only_fp8_weights[False|True]: give
  vllm_config a non-None quant_config so _uses_unquantized_flashinfer_trtllm
  short-circuits before probing the opaque object() model.

Signed-off-by: seonjinn <sna@nvidia.com>
seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 5, 2026
@seonjinn
Merge latest refit base into batched prequantization
cfbe773 
Use the replayed PR NVIDIA-NeMo#3804 tree on the latest PR NVIDIA-NeMo#3294 and main integration.

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

seonjinn commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test cfbe773

@copy-pr-bot

copy-pr-bot Bot commented Sep 6, 2026

Copy link
Copy Markdown

/ok to test cfbe773

@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/

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

Reviewers

@terrykong terrykong terrykong left review comments
@youngeunkwon0405 youngeunkwon0405 youngeunkwon0405 left review comments
@claude claude[bot] claude[bot] left review comments
@ZhiyuLi-Nvidia ZhiyuLi-Nvidia Awaiting requested review from ZhiyuLi-Nvidia
@parthmannan parthmannan Awaiting requested review from parthmannan
@sharonyu-115 sharonyu-115 Awaiting requested review from sharonyu-115

At least 1 approving review is required to merge this pull request.

Assignees

@seonjinn seonjinn

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@seonjinn @youngeunkwon0405 @terrykong