Layer streaming: a checkpoint-visible fused NF4 dequantise-and-multiply (upper bound 9.8% of the step) #842

Why

benchmarks/probe-v0.73.0-what-bounds-streaming.md ablated the streamed Llama-3.1-8B NF4 step on the RTX 3050. The setup was batch 1, S=512, interleaved arms, all at 960 MHz. Removing the NF4 dequantisation took the step from 4.211 s to 3.799 s, -9.80% (lines 155-164). Removing every host-to-device transfer bought 1.44%. The mechanism is that every streamed weight is dequantised into a dense bf16 tensor on every visit: 448 calls per step, about 437 MB of VRAM writes per layer visit (lines 299-304).

9.8% is a ceiling, not a forecast. The ablation arm multiplied by a cached zero weight and produced garbage mathematics (lines 151-153). The record says outright that it did not measure what a real fused kernel would cost (lines 540-542). If a kernel appears to beat 9.8%, the two arms differ in something other than the dequantisation.

The constraint: #331

The dequantisation is there on purpose. MatMul4Bit keeps the packed weight and quant_state on ctx as plain attributes, not through save_for_backward (src/soup_cli/utils/layer_stream_runtime.py:164-168). torch.utils.checkpoint cannot see them, so the reference aliases the buffer pool. Above roughly 165 MiB per layer the gradients came back wrong while the forward and the loss stayed bit-exact (#331).

De-aliasing was measured and rejected: 4 220 -> 19 720 MiB on 32B (lines 177-179). The shipped repair is install_dequant_forward (lines 161-230). It dequantises inside the checkpointed region and calls F.linear, so the dense weight is saved through the ordinary mechanism and recomputed (lines 181-184, 220-223). On real 32B the repair cost +650.9 MiB (2.71 layer-equivalents) and -4.8% throughput (benchmarks/gate-h100-validation.md:2674-2682).

So the requirement is not simply "fuse it". It is: fuse it without putting the packed weight back where checkpointing cannot see it.

A nuance the two records state differently

The probe record calls the dequantisation the only streaming-specific cost (line 304). The H100 record shows that at real training shapes bitsandbytes itself already takes _dequant_linear_fallback for resident NF4. Its fused gemm_4bit runs only at small M, at most 32 on the architectures queried (gate-h100-validation.md:2601-2620; also layer_stream_runtime.py:186-189).

The consequence: a large-M fused NF4 kernel is something neither path has today. It would speed up resident NF4 training at those shapes too, and whether resident adopts it decides what "bit-exact vs resident" means (see Acceptance criteria).

Scope

In: the forward and the input-gradient of streamed Linear4bit projections under NF4 streaming.

Out: bf16 streaming (it has no dequantisation), the per-visit Python and launch overhead (a separate item, and the cheaper one), host-to-device bandwidth (1.4%), and a weight gradient (the base is frozen).

Design sketch

• Before writing a kernel, check whether an existing library already ships a large-M NF4 GEMM that meets the constraints below.

• A custom autograd Function (or torch custom op) at the plug point in _dequant_forward (layer_stream_runtime.py:220-223). Keep that function's dtype handling, which mirrors Linear4bit.forward line for line (lines 186-192).

• Save only through save_for_backward: the packed nibbles, the absmax and the nested statistics. No tensor may be held on ctx as an attribute. That is the whole Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 lesson.

• The backward computes only grad_x = grad_y @ W with a transposed fused kernel. No dense W is ever materialised or saved.

• Deterministic reduction. No atomics and no non-deterministic split-K. The project's gate is equality, and a kernel that is not reproducible against itself destroys it. The same record measured that for Cut Cross-Entropy (lines 404-419).

• Possible side effect, to measure rather than claim: F.linear currently saves one layer's dense weights across the recompute. Not saving them may recover part of the repair's +650.9 MiB.

Acceptance criteria

• Re-measure the dequantisation share on the target card first (arm C of the record). If the share there is small, stop and say so.

• Aliasing control. Compare gradients with recycled buffers against private buffers, with a control arm that keeps the packed weight outside save_for_backward and must produce wrong gradients. Otherwise the harness cannot detect the Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 class. benchmarks/harness/bnb_repro.py has the no-download protocol shape (recycled/private arms, a bf16 control, a --bypass-pool negative control). The layer must be above the ~165 MiB threshold where Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 fired.

• Streamed vs resident NF4, bit-exact, forward and backward stated separately. The repair-disabled control runs in the same process. Both arms run the new kernel; otherwise the comparison measures the kernel, not streaming.

• Kernel vs today's dequantize_4bit + F.linear. Report the difference per shape and per M, with a counter proving which path actually ran. The H100 record's first numerics answer was vacuous because both arms took the same fallback (gate-h100-validation.md:2537-2559).

• tok/s and peak VRAM together, interleaved arms, with a same-session GEMM ceiling and the SM clock reported (the record's method, lines 143-150 and 215-220). Also report memory-controller utilisation, which the record sampled at 44-51% under load (lines 259-260).

• Streamed NF4 is not bit-exact vs resident NF4 on Blackwell (sm_120): 4.9e-04 / 3.9e-03, while quantization=none stays exact #776 caveat: on sm_120, streamed NF4 is currently not bit-exact against resident NF4 even without this change (4.88e-4 fp16 / 3.9e-3 bf16). The streamed-vs-resident NF4 criterion therefore runs on a card where the invariant holds, or after Streamed NF4 is not bit-exact vs resident NF4 on Blackwell (sm_120): 4.9e-04 / 3.9e-03, while quantization=none stays exact #776 is resolved. The aliasing control is unaffected and runs on the RTX 5070.

• The harness is published under benchmarks/harness/ and the result as a gate record in benchmarks/.

Upstream

bitsandbytes#2034 (MatMul4Bit keeps B and quant_state on ctx) is open as of this writing. Close or re-scope this issue if upstream fixes it in a way that survives checkpointing and gives a training-M path that does not materialise a dense weight. A #2034 fix on its own would let streamed weights back through bitsandbytes, but at training M that still means _dequant_linear_fallback, which costs the same as today.

Hardware needed

A CUDA GPU. The RTX 5070 Laptop (8 GB, sm_120) is enough for the aliasing control and for 8B NF4 streaming, whose peak was 3.448 GB on the 3050 (probe record line 103). The kernel toolchain must target sm_120. On Windows, Triton ships as the triton_windows distribution, and that name already broke a version check once (probe record line 331). All numbers above are from the RTX 3050 and an 8x H100 box. A 5070 measurement is a new baseline.

Honest limits

• The 9.8% is one card (RTX 3050, Windows, torch 2.5.1, bitsandbytes 0.49.2; record lines 8-10), one model, one shape.

• A fused kernel almost certainly changes numerics against the shipped arithmetic. Adopting it means redefining the NF4 reference, and that should be decided deliberately, not arrive as a side effect.

• Why Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 fires sharply around 165 MiB per layer is still unexplained (Layer streaming: NF4 above ~165 MiB/layer produces silently wrong gradients (forward stays bit-exact) #331 closing comment), so the control's layer size is an empirical choice.

Related, filed the same day:

• soup bench train: a reproducible training-benchmark contract that fails the run when nothing is training #836 — soup bench train: a reproducible training-benchmark contract that fails the run when nothing is training
• Layer streaming: cut the ~10 ms fixed cost per layer visit, which dominates the step below ~128 tokens #841 — Layer streaming: cut the ~10 ms fixed cost per layer visit, which dominates the step below ~128 tokens