[nxphi47]

Summary and decisions

Adds opt-in FP8 forward / BF16 backward QAT for dense/grouped projections and routed experts, plus independent DeepSeek V4 FP8 KV simulation. Dense QAT, expert QAT, and SR are configured independently; protected FP32 optimization/reduction defaults are unchanged.

Stacked on #4067, with base feat/fp8-stochastic-rounding. That PR owns the shared SR state, parameter identities, Philox helper, global expert coordinates, and SGD support. QAT reuses them for consistent masks across FSDP/CP replicas, EP shards, tied weights, accumulation microsteps, forward/backward, and checkpoint replay. Merge #4067 first, then rebase/retarget this PR to main.

This selectively adapts grouped-projection and KV ideas from #3627 rather than adopting its backward/RNG behavior. QAT computes dX = dY @ decoded(W_fp8) using the exact forward weight decisions, and dW = dY.T @ X_bf16; GEMMs are BF16. KV fake quantization uses identity STE. QAT preserves grouped output projections and fuses inverse RoPE with FP8 activation casting.

Use ordinary FP8+SR when SR is the only requested change. QAT is for the requested backward precision and audited vLLM forward conventions. The short validation does not establish a quality advantage.

Committed report documents decisions, the exact SR contract, portable configs, full results, failed attempts, and limitations. Machine-readable results retain every step's loss, gradient norm, timing, and W&B link.

BIG NOTICE — TRAINER / vLLM QUANTIZATION CONVENTIONS DIFFER

The audited serving target is vLLM 0.31.0. QAT follows its numerical conventions, with an explicit SR extension using #4067's shared contract.

Convention	Difference and decision
Weight floors	Ordinary trainer FP8 floors scale at 1e-4; vLLM floors amax at 1e-4 before /448. QAT follows vLLM's rule and operation order.
Two-stage weights	vLLM online-casts, then requantizes decoded weights to UE8M0 on Hopper and Blackwell. QAT reproduces both stages. SR is at the first stage; deterministic requantization can change final unbiasedness.
Negative zero	Preserve E4M3 -0 (0x80) separately from +0 (0x00), including tiny negative values rounding to zero. Checks compare payload bytes and decoded BF16 bits.
SR agreement	Stock vLLM's audited online cast uses RN. #3627's rank-local RNG cannot provide the shared rank/replay contract. QAT and ordinary FP8+SR share the same stateless seed/salt/counter and lower/upper-magnitude rule. Serving SR must implement that exact contract; matching a seed alone is insufficient.
Activation floors	Generic vLLM CUDA quantization floors scale; the fused grouped output projection floors amax. QAT preserves the respective rules (2^-33 versus 2^-42 minimum UE8M0 scales).
Compact formats	Hopper FP32 scale containers and Blackwell packed UE8M0 exponent bytes have different layouts/strides/padding. DeepGEMM handles packing at the GEMM boundary. B300 execution/packing remains unvalidated.
KV/indexer	Preserve vLLM's separate SWA amax and compressed-KV scale floors, FP8/BF16 KV channel split, and indexer amax floor. KV fake quant returns decoded values, not physical cache bytes.
Remaining differences	Compressor BF16 boundaries and pooling/reduction implementations differ. Matching quantizers does not establish whole-model logit parity. Shared masks do not imply equal final weights when scale/stage conventions differ.

Results

All three full 43-layer DeepSeek V4 Flash Base runs completed five real-data SFT steps on two nodes / 16 H200s. Slurm 4082 completed with exit 0. All 15 losses and gradient norms were finite; all NaN-loss counts were zero.

Common settings: pinned base checkpoint/data, EP16, sequence 512, SGD lr=1e-5, optimizer-state-only offload, FP32 optimization/reduction, compilation, activation checkpointing/offloading, FP8 KV simulation, and reduced NCCL buffers. These are SGD validation runs, not AdamW comparisons.

Mode	DP / CP / EP	Batch	Final loss	Rank-0 peak reserved GiB	W&B
QAT SR	16 / 1 / 16	16	1.01928	138.172	SR/CP1
QAT RN	16 / 1 / 16	16	1.01822	138.162	RN/CP1
QAT SR	8 / 2 / 16	8	1.07499	137.996	SR/CP2

CP1 last-two-step medians were 10.531 s RN / 11.239 s SR (+6.7%) on the same nodes. Only two samples, warmup, sequential execution, and allocator retries make this preliminary. CP2 uses a different batch and had a final-step timing outlier. Recovered allocation failures show that capacity is tight; PyTorch reserved memory excludes external allocations such as NCCL.

Operator forward/backward measurements (dense / grouped, milliseconds): BF16 0.279 / 0.596, ordinary FP8 RN 0.813 / 2.832, ordinary FP8 SR 0.940 / 4.851, QAT RN 0.558 / 1.455, QAT SR 0.652 / 4.183. Exact shapes/methodology are in the report. These are short microbenchmarks, not a general speed or quality claim.

Validation completed

• 101 H200 GPU tests: actual vLLM byte/scales references, independent Philox reference, negative zero, cross-mode masks at matching scales, wide counters, BF16 gradients, grouped/empty experts, tied weights, checkpoint replay, and inverse RoPE.
• Three distributed cases passed on each of two ranks: FSDP, CP token splitting, and EP shard/whole-weight agreement, including input gradients.
• 245 config tests and 30 MoE regression tests passed; lint/format checks passed on the tested implementation.
• Nine SM103 kernel variants compile, which is not B300 execution evidence.
• Runtime results are from f1eee78b725b3dafc2281144b8169be780806c64. The subsequent parent-history merge and report commit preserve runtime sources, packages, tests, dependencies, and lockfiles byte-for-byte; this was verified before publishing. Report links and copied result data were also checked.

Remaining work

• B300 hardware validation is blocked by availability. PTX KV work has not started and remains gated on the requested validation, including B300.
• Earlier full CPU optimizer attempts completed zero steps: one logged optimizer timeouts with node failures; the retry exited with SIGSEGV. Root causes are unconfirmed, and OOM is not established. Successful SGD runs do not validate AdamW/full offload.
• Longer optimizer-matched studies are needed before claiming convergence or a quality benefit. No end-to-end RL/serving or whole-model logit-parity result is claimed.