v0.72.2 — NF4 layer streaming: fine-tune Llama-3.1-8B on a 4 GB card

What's New

NF4 layer streaming — fine-tune Llama-3.1-8B on a 4 GB laptop GPU.

Layer streaming (v0.72.0) keeps the frozen base in CPU RAM and feeds it to the GPU one decoder layer at a time, so peak VRAM is bounded by one layer instead of the whole model. It was bf16-only, which capped it near 3B on a small card. Quantising the streamed base to NF4 shrinks it ~4×, and that is what brings 8B within reach.

One line to enable it:

training:
  stream_layers: true
  quantization: 4bit    # NF4
  batch_size: 1

Measured on a 4 GB RTX 3050 Laptop (Windows, batch 1, S=512, gradient checkpointing, 50 steps after 10 warm-up, PagedAdamW8bit):

Model	tok/s	Peak VRAM	RAM store	GPU util
Llama-3.1-8B-Instruct	119.6	3.32 GB	3.60 GB (page-locked)	100%
Qwen2.5-3B	264.2	1.76 GB	1.43 GB (page-locked)	100%

For scale, 1M training tokens is about 2.3 h at 8B on that card — arithmetic from the measured rate, not a separate measurement.

Qwen2.5-3B is also 1.85× faster than the bf16 streaming path (264.2 vs 143.1 tok/s), and the reason is not arithmetic: a 1.43 GB store fits under the machine's page-locked memory ceiling where a 5.55 GB one did not, which restores asynchronous host-to-device copies and lifts GPU utilisation from 79.3% to 100%.

Correctness. A streamed NF4 run is bit-exact against a resident NF4 run — the same quantised bytes through the same bitsandbytes kernels. Logit equality, non-zero gradients at layer 0, and a matching multi-step loss curve are all regression tests that run on CPU in CI, not one-off measurements.

The base is quantised once, offline, and cached under ~/.soup/layer-stream/. The cache is keyed to the quantisation and the source checkpoint, so switching between none and 4bit, or retraining a base in place, re-shards instead of silently streaming the wrong bytes.

Also fixed: a streamed 4-bit run reported its parameter count ~6.5× too high (SmolLM2-135M printed "878,154,048 total" against a true 134,515,008). Display only — training was unaffected — but at 8B it would have read ~52 B.

This release also carries the soup --help startup fix already on main: 6.0 s → 1.15 s, after PyTorch leaked onto the light CLI import path in v0.71.41.

Install / Upgrade

pip install --upgrade "soup-cli[train]"

Security

• index.json is treated as a trust boundary. Its shape / blocksize / quant_type / dtype fields are read back off disk and reach bitsandbytes' dequantise kernels, which allocate and read prod(shape) elements without bounds-checking the packed tensor. Deserialisation now validates every field, and the runtime cross-checks the index's claim against the bytes actually present in the shard, so a corrupted or tampered index fails as a clean Python exception rather than an out-of-bounds read in native code.
• The shard cache key now covers the quantisation, the double-quant flag and the quantisation device, alongside the existing dtype and source-checkpoint fingerprint. Reusing a bf16 shard set for an NF4 request would otherwise have fed full-precision bytes to matmul_4bit.
• Layer-uniformity is now checked by tensor shape, not just parameter names, because layer 0's quantisation state is reused for every layer.

Known Limitations

• 8B is the largest size measured. 14B NF4 (~7.5 GB store) exceeds the reference box's measured 7.12 GB page-locked ceiling and would fall back to the pageable path. It was not run, so no 14B claim is made.
• Untied embed_tokens + lm_head stay resident and unquantised — 2.10 GB of the 8B row's 3.32 GB peak. That is why 8B sits close to this card's ceiling; treating them as streamed large layers is a v0.72.3 item.
• training.bnb_4bit_use_double_quant is ignored — by the streaming path and by every resident 4-bit load in Soup (double-quant is always on). Streaming deliberately matches the resident behaviour rather than diverging, because the streamed-vs-resident bit-exactness guarantee depends on both sides quantising identically.
• The pre-flight RAM estimate is deliberately coarse — it charges the NF4 rate to every parameter although layernorms and embeddings stay unquantised, which errs low. The authoritative check runs after sharding, on real shard sizes.
• The shard cache costs disk roughly equal to the quantised model (~5.7 GB for 8B).
• Apple Silicon (MPS) is untested. bitsandbytes ships no 4-bit MPS kernels, so NF4 streaming is measured on CUDA and CPU only. The two end-to-end training tests skip on an MPS runner rather than assert something unverified.
• Numbers are Windows/WDDM and therefore systematically pessimistic versus Linux. This card's boost clock varies ~13% between sessions, so any fraction-of-ceiling figure is quoted only with the clock it was taken at.
• Scope is unchanged and still BETA: RAM tier, task: sft, Llama/Qwen, batch_size: 1, no gradient accumulation, no --resume. The disk tier, larger batches, accumulation, checkpoint/resume and more architectures land in v0.72.3; preference losses in v0.72.4.