perf(refit): batch MXFP8 expert prequantization

[seonjinn]

Summary

Batch matching MoE expert weights before trainer-side MXFP8 prequantization. The implementation quantizes at most 16 experts per call, reuses stream-local stacking buffers within each export pass, and releases them before training resumes. Stream handoffs synchronize pending inputs, batched outputs, and per-tensor fallbacks before packing. The path keeps the existing weight and scale entries consumed by vLLM.

Non-expert weights and expert groups that cannot be batched keep the existing per-tensor path.

Depends on #3294, which adds trainer-side prequantized refit. Receiver-side MXFP8 layout conversion remains in #3478.

Why

The current prequantized path launches one quantization operation per expert tensor. Large MoE models repeat that operation hundreds of times per refit. Batching tensors with the same layer, projection, shape, dtype, and device reduces launch and synchronization overhead without changing the refit wire format.

Prior prototype results

Qwen3-30B-A3B, GB200, synchronous colocated GRPO, CUDA Graph enabled, 20 steps:

Metric	Per-expert	Batched	Change
Transfer/update	4.506 s	2.200 s	-51.2%
Total refit	8.431 s	5.935 s	-29.6%
E2E step	188.88 s	185.32 s	-1.89%
Throughput	2,199.0 tok/s/GPU	2,243.0 tok/s/GPU	+2.00%

The 20-step prototype run completed with unchanged generation KL error (0.00398). The current head adds stricter export-order, scale-shape, memory-lifetime, and CUDA-stream handling. Its end-to-end performance rerun is still pending.

Validation

• Current head 1d1970194b868553d997a8d234976e3caff1b337 passed 17 MXFP8/layout tests and 1 Megatron worker integration test on GB200
• Unit coverage for realistic Megatron export order, bounded expert batching, value and scale parity, fallback behavior, grad-enabled parameter views, stream-local scratch reuse, and all packed-stream handoff boundaries
• Worker-level coverage that routes enabled prequantized refit through the batched iterator
• Formatting and lint checks on all changed files

Current-head Qwen3-30B-A3B performance revalidation and Qwen3-235B end-to-end memory and performance validation are still pending, so this PR remains a draft.

[seonjinn]

Self-review at head 1d1970194

The GitHub diff shows 78 files because the branch stacks on #3294 (88721ced8 is an ancestor of this head); the actual review surface is 7 commits / 4 files (+574/-2): the batched expert prequantization in fp8_train_utils.py, its wiring in megatron_policy_worker.py, and two test files. Everything below is scoped to that delta - #3294's findings (including the blocking prepare_refit_info signature drift and the per-batch scale-check straddle) apply to the base layer and are already posted there.

Process: three domain reviewers (numerical correctness, stream/memory lifetime, tests), then an adversarial verification pass that re-read every citation and fetched Megatron-Bridge at the pinned submodule commit: 8 confirmed, 5 downgraded, 2 refuted, 1 new. No Critical or High correctness bug survived verification - the core batching math, grouping, fallback routing, wire pairing, and all four stream-sync fix commits check out.

Gates for undrafting

1. Perf + peak-memory rerun at current head. The body's prototype table predates the four fix commits and is already labeled as pending rerun - just noting the memory side should be measured too: with interleaved per-expert export order, pending extends the lifetime of up to ~46 projection tensors (~540 MiB for 235B-shaped experts) plus one reused ~192 MiB stacking scratch per (device, dtype, stream) key. Much of that is memory Megatron-Bridge's grouped export already materializes, so the true delta needs a torch.cuda.max_memory_allocated() before/after on a real MoE refit rather than arithmetic.
2. One worker-composition test (inline below): the only worker-level test stubs out the real iterator, so an export-naming drift would silently disable batching (correctness preserved via per-tensor fallback, but the PR's entire perf value gone with no failing test).

Verified clean (adversarial pass, own citations)

• Batched vs per-tensor parity: MXFP8 quantization is row-independent, so the flatten-stack-quantize-unflatten path is mathematically identical to per-tensor; the CPU parity test does a genuine bitwise torch.equal against the real per-tensor quantizer (not a tautology) and covers both the multi-block and squeezed single-block scale shapes. E8M0 zero-scale clamp and the squeeze(dim=-1) behavior are inherited verbatim through the shared quantize_fn, not reimplemented.
• Stream handshake complete at every producer-consumer edge (pending-input to stack, stack to quantize output, per-tensor fallback), each with a CUDA-gated regression test targeting exactly that edge. The load-bearing assumption - that export_hf_weights produces on the ambient stream - was verified against Megatron-Bridge at submodule pin 8c46dc425: zero stream primitives in model_bridge.py/param_mapping.py, synchronous collectives, cpu=False. Holds today; worth a comment for future Bridge bumps (inline).
• Grouping/fallback: layer+projection-scoped grouping with flush-on-prefix-change; shape/dtype/device/layout mismatch, lone remainders, duplicate expert ids, fused w13-style names, and shared_experts all route to the per-tensor or passthrough path correctly. Chunk arithmetic has no off-by-one; max_experts_per_batch<=0 raises, ==1 degenerates to exact old-path behavior.
• The scratch-release fix (34da20e46) is the right call: the buffer now lives only for the export pass, and record_stream bookkeeping makes an explicit sync unnecessary.

Refuted by the adversarial pass (so nobody re-flags them)

• "Shape-mismatch fallback never actually triggered in tests" - false; the (3,64) parametrized case drives a real mismatch through the fallback with value assertions.
• "The grad-enabled fixture wouldn't catch a removed no_grad guard" - false; the torch.stack(..., out=) in the test path is real production code and would raise RuntimeError without the guard, which is exactly what commit 09591f0b2 fixed.

Notes

• Race-window tests use fixed torch.cuda._sleep(5_000_000) cycles; the failure mode on faster future GPUs is a spurious pass (missed detection), not flakiness - fine for CI, worth a comment.
• Test niceties, all Low: a zeroed row through the CPU parity test (clamp through the batched reshape), a dtype-mismatch fallback case, a 16/17-expert default-boundary case.

[seonjinn]

[Medium] The only worker-level test stubs out the real iterator, so batching can be silently disabled with no failing test.

This test monkeypatches fp8_train_utils.iter_mxfp8_prequantized_params with a stub (effective because the call site does a deferred from-import on every call, megatron_policy_worker.py:2256-2258), and nothing else composes the worker with the real iterator. The direct unit tests do exercise _EXPERT_WEIGHT_PATTERN against realistic HF export names, so the regex isn't untested - but if Megatron-Bridge's export naming ever drifts, every selected tensor falls back to per-tensor quantize_one: correctness preserved, the PR's entire perf win silently gone, and no test turns red.

Suggested: one composition test that runs the real iter_mxfp8_prequantized_params behind a fake export_hf_weights yielding export-shaped expert names and asserts batched (stacked) quantize calls actually occur.

[seonjinn]

[Low] Two ambient-stream assumptions worth a comment - both hold today, neither is checked.

1. torch.cuda.current_stream() sampled at consumption time is assumed to be the stream export_hf_weights produced on. Verified true at the pinned Megatron-Bridge commit 8c46dc425 (no stream primitives anywhere in model_bridge.py/param_mapping.py, synchronous collectives, ambient-stream export loop) - but a future Bridge bump that introduces a side/copy stream would turn every wait_stream in this file into a silent race (intermittent bad weights, not a crash).
2. The scratch cache key uses the raw cudaStream_t pointer (int(stack_stream.cuda_stream) below). If a torch.cuda.Stream were created and destroyed within one export pass, CUDA can reissue the pointer and two different streams would alias one scratch entry. No current call site does this, but the failure would be silent corruption rather than an error.

A short comment documenting both ("producer must be the ambient stream; streams must outlive the export pass") makes the contract auditable at the next submodule bump.

[seonjinn]

[Low] scratch_cache is per-export-pass by design in production - say so, or the next reviewer re-flags it.

The only production call site (megatron_policy_worker.py:2260-2263) passes no scratch_cache, so cross-refit reuse never happens - and that's intentional: 34da20e46 deliberately removed the worker-level cache so the ~192 MiB stacking buffer doesn't stay resident through training. What remains meaningful is in-pass reuse (one buffer shared across all layers/projections of a single export). The parameter is now test-injection-only API, and test_batched_expert_prequantization_preserves_wire_entries_and_reuses_scratch still asserts cross-call reuse as if it were a production behavior. A one-line docstring note (or reframing that assertion as a library-capability check) closes the trap. Do not wire a persistent cache back in - that would undo the memory-lifetime fix.

[seonjinn]

[Low, latent] Yielded per-expert tensors are views pinning the whole <=16-expert batch.

value[index] / scale[index] are views into the single stacked quantize output, so any downstream consumer that retains one expert's tensor keeps the entire batch's value and scale allocations alive. The current receiver copies into vLLM params promptly, so nothing leaks today - but a persistent-IPC or caching consumer would silently hold up to 16x the expected memory per retained entry. A one-line comment is enough; .contiguous()/clone per entry would defeat the batching win and is not recommended.

[seonjinn]

[Info] _maybe_prequantize_param's quantization branch is now dead in production.

This else branch only runs when self._refit_prequant_names is empty, and _maybe_prequantize_param passes through whenever the name isn't in that set - always true here. All real quantization now flows through iter_mxfp8_prequantized_params; the old helper's quantize branch is reachable only from its direct unit tests. Worth a comment (or folding the helper away) so a future reader doesn't assume it's still on the hot path.

[seonjinn]

/ok to test cd825d2

[seonjinn]

/ok to test 6fbbb27

[seonjinn]

/ok to test 5ccba2f

[seonjinn]

/ok to test e83281e

[seonjinn]

/ok to test 1c9e4b5