# 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).