PHASE1_REPORT.md

Phase 1 — Gemma-4-26B-A4B teacher capture: feasibility & proof

Box: pop — RTX 5070 Ti (16 GB, sm_120 Blackwell), 31 GB RAM, torch 2.11+cu128, transformers 5.12, bitsandbytes 0.49.2. Date: 2026-06-18.

Goal (Phase 1 only): make the Gemma-4 text decoder runnable forward-only for activation capture on this hardware, and prove the capture works. No training yet.

TL;DR

• It runs. The 26.5B teacher's text decoder does forward-only capture on the 16 GB card via bf16 layer-streaming (accelerate offload to the fast NVMe). Capture proven correct on a 32,768-token slice: all 30 layers, correct shapes, finite, X≠Y.
• bnb-4bit does NOT work here — neither fully-resident (OOMs at load) nor offloaded (meta quant_state crash). Details below.
• The real blocker is STORAGE, not compute. Storing block I/O (X and Y) for all 30 layers costs 330 KiB/token → 6.15 TiB for 20M tokens, 15.4 TiB for 50M (50M exceeds the 6.3 TB free on the cache drive). And the cache drive is a slow USB HDD (111 MB/s write), which makes the capture write-bound at ~330 tok/s → ~17 h for 20M tokens. This needs a decision before a real run — see NEEDS_MIKEY.txt.

What is captured (verified against the reference forward)

In Gemma4TextDecoderLayer.forward the FF/MoE sub-block is, after the attention residual:

residual = hidden_states                       # X  = FF-block input (post-attention)
h  = pre_feedforward_layernorm(residual)
h  = mlp(h)                                     # shared (dense) expert
if moe:                                         # 8-of-128 routed experts (GeGLU)
    h1 = post_feedforward_layernorm_1(h)
    _, w, idx = router(residual.flatten)
    h2 = experts(pre_feedforward_layernorm_2(residual.flatten), idx, w)
    h  = h1 + post_feedforward_layernorm_2(h2)
h  = post_feedforward_layernorm(h)             # D  = FF-block delta
hidden_states = residual + h                    # Y  = X + D = FF-block output
hidden_states *= layer_scalar                   # (AFTER the block — NOT captured)

Capture hooks: X = forward-PRE hook on pre_feedforward_layernorm (its input); D = forward hook on post_feedforward_layernorm (its output); stored Y = X + D. hidden_size_per_layer_input = 0, so the per-layer-input branch is inert. This is the exact block the student replaces, so (X, Y) is the correct feature-distillation pair.

Load methods evaluated

method	result
bf16 layer-streaming (accelerate, NVMe offload)	✅ WORKS. device_map=auto, max_memory={0:6GiB, cpu:16GiB}, offload_folder on NVMe. 4 layers GPU-resident / 9 CPU / 23 disk. Reuses the real Gemma4 forward.
bnb-4bit, fully resident on GPU	❌ text decoder ~14.4 GiB; load OOMs on the final ~484 MiB expert tensor. No room to finish loading, let alone a forward. (The bnb 4-bit kernel runs fine on sm_120 — verified with a standalone quantize+matmul; it's purely the size.)
bnb-4bit + accelerate CPU/disk offload	❌ loads, but forward dies: offloaded quant_state tensors stay on meta → NotImplementedError: Cannot copy out of meta tensor. (Also needed a monkeypatch just to get past hook-attach, where bnb does offset.item() on a meta tensor.) bnb-4bit is fundamentally not offloadable on this stack.
--cpu-layers (per-layer device_map → CPU bf16)	❌ silently ignored by transformers 5.12's new core_model_loading; the '':0 catch-all wins, GPU use unchanged (identical OOM for cpu-layers 2 and 8).
disabling warmup / serializing the loader (1 worker)	❌ no effect on the 14.4 GiB resident wall.
AWQ / compressed-tensors / FP8 prequantized checkpoint	❌ none exists for this brand-new model; 4-bit AWQ would hit the same 14 GiB wall, FP8 (26 GiB) is far too big.

Two small, documented monkeypatches live in capture_gemma.py for the 4-bit diagnostic path only (patch_disable_warmup, patch_serial_load, patch_bnb_meta_quantstate); the working bf16 path needs none of them.

Throughput (measured)

• Forward-only (no writes), bf16 streaming, ctx=256:
  • batch 16 → 153 tok/s (peak 10.4 GiB)
  • batch 64 → 558 tok/s (peak 12.1 GiB) — bigger batch amortises the ~29 s/forward NVMe weight transfer over more tokens. This is the lever; batch is bounded by the MoE token×top_k activation expansion.
• Full capture incl. writes (batch 16) → 97 tok/s, fwd peak 7.06 GiB (write-bound).
• Key infra: NVMe read 3.0 GB/s (teacher); USB-HDD write 111 MB/s (cache drive).

Capture proof (validation slice)

capture_gemma.py --val --batch 16 --shard-tokens 8192 → /mnt/data/cache/gemma_cap_val (32,768 tokens). Verified by inspect_capture.py:

• 30 layer dirs; each layerNN/{input,output}_NNNNN.pt, 4 shards/tensor (sharding works).
• Shapes [8192, 2816] bf16, all finite; X≠Y with FF-block relative delta 0.93–3.36.
• 330 KiB/token total (X+Y over 30 layers) = 11,264 B/token/layer.

Size & wall-time estimates for a real capture

At 330 KiB/token, and (for wall-time) batch≈64 so forward (≈558 tok/s) exceeds the HDD write rate → write-bound at ≈330 tok/s:

corpus	cache size	wall-time (write-bound)	fits 6.3 TB free?
20M tokens	6.15 TiB	~17 h	barely
50M tokens	15.4 TiB	~43 h	NO

(If left at batch 16 it is forward-bound instead: ~57 h for 20M. Use batch ≥48.)

Recommendation / open decision (see /tmp/NEEDS_MIKEY.txt)

The forward problem is solved; the storage cost is the gate. Options, in order of my preference:

1. Don't pre-materialize all 30 layers' X,Y to disk. Re-stream the teacher during training (compute distillation targets on the fly), or capture/train a few layers per pass. Removes the multi-TB cache entirely. (Phase-2 design change.)
2. fp8 storage halves the cache (165 KiB/tok → 3.1 TiB @ 20M) and roughly doubles the write-bound rate; verify distillation tolerates fp8 targets.
3. Cap the budget at ~15–20M tokens on the existing USB HDD (~6 TiB, ~17 h) — feasible today but uses nearly all free space.
4. Rent a one-off 80 GB GPU for a fast capture pass — fixes speed (full model resident, no offload, hours not days) but not storage; you'd still need ~6–15 TB of fast scratch, so pair it with option 1 or 2.

Files

• capture_gemma.py — capture (bf16-stream default; --method 4bit reproduces the OOM).
• test_stream_bf16.py — standalone streaming throughput probe.
• inspect_capture.py — shard validation + size/time estimator.
• Validation output: /mnt/data/cache/gemma_cap_val/ (meta.json + sample shards).