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

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

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

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

seonjinn · 2026-08-18T06:41:17Z

/ok to test 88721ce

seonjinn

Self-review at head `88721ced8`

Second team self-review (the first covered 7705d22a3; the only delta since is the 2026-08-17 main merge). Six domain reviewers over the unique diff vs merge-base 3d873ebc2, then a devil's-advocate verification pass that re-read every citation: 12 findings confirmed, 2 refuted, 3 downgraded, 2 new. Blocking items are posted inline; the rest is summarized here.

Blocking (inline)

[Critical] VllmQuantInternalWorkerExtension.prepare_refit_info still takes one data argument while both driver workers now RPC two - TypeError on the first refit for every ModelOpt-backed vLLM engine.
[High] The refit_slim_offload_after branch drops the offload_optimizer_for_refit gate and moves the optimizer to CPU against explicit user config; the new unit test encodes the missing gate as expected behavior.
[High] The per-batch MXFP8 *_scale_from_checkpoint presence check false-positives whenever a weight/scale pair straddles a transfer-batch boundary - deterministic crash for oversized params on the IPC path, probabilistic on packed collective. Small-CI models fit one buffer, so CI cannot catch it; 235B-class targets hit it every refit.

Rebase-blocking collisions with `main` (not fixable on this branch as-is)

The branch merged main@3d873ebc2 but main has moved 28 commits; a trial merge conflicts in 14 files. Two are semantic, not mechanical:

#3477 (36b5999a7): main's NCCL-reshard receiver imports quantize_mxfp8_weight from fp8.py; this branch has no such function and instead inlines the equivalent quantize + zero-scale clamp. Taking "ours" on those hunks breaks main's import. The two implementations must be unified (the shared helper is the natural home for this PR's trainer-side prequantizer too).
#3391 (85e02cca3): main passes a CollectiveSenderSpec (buffer_size_bytes, num_buffers, nccl_peer) into init_collective; this branch passes a scalar buffer_size_bytes at the same call sites. Two incompatible designs on one interface - needs a design decision, not a merge resolution.

This confirms the 2026-08-18 assessment that the branch must not be mechanically refreshed. Verdicts on that comment's specific claims: the #3477-family overlap is real and now concrete (above); the "#3630/#3638" numbers do not appear in main history (likely misattributed - #3477/#2744 are the actual drivers); the "CuTe scales" claim was unverifiable (no such reference in either branch), though *_scale_from_checkpoint is genuinely contested ground between this branch and main's receiver-side quantize path.

The PR body's "#3477 remains separate; NCCL Reshard changes are not included here" is stale - #3477 merged 2026-08-19, and the scope/interaction statement needs rewriting after rebase.

Non-blocking notes

checkpoint_engine.py refit now unconditionally runs the full _weight_update_lifecycle finalize where base only ran the FP8-KV-cache conditional post-process. Looks like intentional harmonization; confirming intent since it adds per-refit work for non-FP8-KV deployments.
The prequant handshake mutates trainer state (_refit_prequant_names) before the second generation.prepare_refit_info push; all 10 call sites verified fail-loud today, so this is latent - worth a comment or commit-after-push ordering.
interfaces.py's rejection message says vllm_cfg.refit_prequantize; the config path users set is policy.generation.vllm_cfg.refit_prequantize.
Four reference configs carry refit_cache_loader_routes with no counterpart in their real source configs (grpo_math_1B.yaml lacks the key).
Loader-route cache resets silently when vLLM post-load recreates Parameters; one log line would make a quietly-inert perf knob observable.

Verified clean (spot-checked by the adversarial pass)

E8M0 zero-scale clamp parity trainer/receiver; _mxfp8_e4m3_quantize_torch math incl. 2D/3D reshape; persistent IPC buffer reuse and oversized-solo fallback; loader-route cache record/replay/fail-closed reset under nccl_reshard's in-place receive (an earlier MAJOR claim here was refuted - no .data reassignment exists in that path); buffer_size_bytes threading; YAML recipe/test matrix; GB200 wire-parity test is wired into L1 CI on target hardware.

seonjinn · 2026-08-24T22:06:25Z


-    def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None:
-        super().prepare_refit_info(state_dict_info)
+    def prepare_refit_info(


[Critical] Signature drift: first refit raises TypeError for every ModelOpt-backed vLLM engine.

The base class gained a second positional parameter (vllm_backend.py:460-464, prepare_refit_info(self, state_dict_info, serialized_fp8_config)), and both drivers now RPC two args unconditionally:

vllm_worker.py:1157-1160: collective_rpc("prepare_refit_info", args=(state_dict_info, fp8.serialize_fp8_config()))

vllm_worker_async.py:1412-1415: same

This override still declares prepare_refit_info(self, state_dict_info), so the RPC fails with TypeError: takes 2 positional arguments but 3 were given before the body runs - the _is_real_quant_model()==False fallback is unreachable, so fake-quant configs break too. No test exercises the 2-arg call against this extension.

Fix: accept and forward the new parameter (def prepare_refit_info(self, state_dict_info, serialized_fp8_config=None) and pass it through to super()).

seonjinn · 2026-08-24T22:06:25Z

+                and self.optimizer is not None
+                and not self.optimizer_cpu_offload
+            ):
+                self.move_optimizer("cpu")


[High] Slim path ignores offload_optimizer_for_refit: False and moves the optimizer to CPU anyway.

offload_before_refit gates the optimizer move on both knobs (:2958-2963 includes self.offload_optimizer_for_refit), but this slim branch only checks not self.optimizer_cpu_offload. A user who sets offload_optimizer_for_refit: false to keep the optimizer resident (e.g. CPU-RAM-constrained hosts) gets it silently offloaded whenever refit_slim_offload_after: true - the opposite of their config, with no warning.

The new test test_offload_after_refit_routes_cleanup_by_mode (test_megatron_worker.py:710-765) never sets offload_optimizer_for_refit and asserts move_optimizer.assert_called_once_with("cpu") for the slim mode, so it currently enshrines the missing gate.

Fix: add and self.offload_optimizer_for_refit to this condition and extend the test with the False case.

seonjinn · 2026-08-24T22:06:25Z

+                    "other FP8 trainer scale layouts are not compatible."
+                )
+            scale_name = k + "_scale_from_checkpoint"
+            if global_fp8_config.is_mx and scale_name not in weight_names:


[High] Per-batch scale check crashes valid refits when a weight/scale pair straddles a transfer-batch boundary.

weight_names covers only the current load_weights batch, but every transport batches by size with no pairing constraint, and the trainer yields the E4M3 weight then its scale as two separate entries (megatron_policy_worker.py:2039-2041):

packed collective: the producer breaks after the threshold-crossing tensor (packed_tensor.py:104-108). A prequantized weight is ~32x its scale, so the weight is usually the tensor that closes the buffer - its scale opens the next one, and this check raises Prequantized MXFP8 weight ... is missing ... on a perfectly healthy stream.

IPC ping-pong: same straddle at utils.py:527-535.

IPC oversized-param branch (utils.py:493-524) ships a larger-than-buffer weight alone, so for oversized params (235B-class embed_tokens) the crash is deterministic on every refit.

Small functional-test models fit one buffer, so CI passes while the PR's actual targets hit boundaries every refit. Note the load itself is batch-order-independent - a scale arriving in a later batch loads fine (exactly how non-MX *_scale_inv already behaves); only this validation is wrong.

Fix options: validate against the full state_dict_info manifest instead of the per-batch names, defer the completeness check to refit end (the IPC manifest already tracks it), or make the sender pack weight+scale atomically.

seonjinn · 2026-08-24T22:06:25Z

+                x_scales = x_scales.view(x.size(0), -1)
+    if x_q is None or x_scales is None:
+        x_q, x_scales = _mxfp8_e4m3_quantize_torch(x)
+    x_scales = torch.squeeze(x_scales, dim=-1)


[Minor, latent] Unconditional squeeze(dim=-1) collapses the block dimension when the reduction dim is exactly 32.

For a weight with K == 32, flashinfer's scale output has a trailing dim of size 1 that is the real K/32 block dimension, not padding; squeezing it yields a 1D scale that later fails vLLM's 2-index slicing in process_weights_after_loading_mxfp8_linear. No realistic checkpoint has K == 32, so latent - but a if x_scales.ndim > 2 guard (or reshape to [M, K // 32] explicitly, mirroring _mxfp8_e4m3_quantize_torch) removes the trap.

Related hardening from the same review pass: this sm100+ fast path has no 3D branch, unlike the torch fallback. The adversarial pass verified 3D expert tensors cannot currently reach it (prequant eligibility requires .endswith("weight") on Linear/RoutedExperts HF exports, and nccl_reshard+prequantize is config-rejected), so it is unreachable today - a defensive assert weight.dim() == 2 here would keep it that way when eligibility rules evolve.