The b1 step, profiled end to end — 2026-10-16

Owner: “полный профайлинг всего и вся”. This is the whole stack at the b1 preset (1,066,358,209 params with tables; 663,572,929 with Engram off; 57M active per token), on the 5060 Ti, eager, one heavy thing at a time. All numbers from one day of probes; scripts in /tmp/opencode/*.py, logs in ~/logs/; the plan’s port table carries the summary rows.

1. Phase split (b8×s1024, max_iter 4, bf16 factors, fp32 masters)

phase	ms/step	share
data	0.9	0.1%
forward	278	17.5%
backward	404	25.4%
grad clip	58-107	3.7-6.7%
optimizer	530-850	~44-53%

Total 1399 ms/step plain-AdamW/TF32-off, 1201 ms/step with TF32 (8192 tok → 6.8k tok/s). The optimizer is HALF the step; that is the finding.

2. The kernel table (whole step, bnb8bit ON, tf32 OFF — the worst mix)

kernel	self CUDA	share
aten::mm (incl. cutlass sgemm + magma_sgemmEx 211 ms)	1.246 s	28.4%
aten::bmm (expert banks + NS)	1.001 s	24.1%
AdamW8bit.step (bnb)	0.711 s	17.8%
elementwise (vectorized_elementwise<4>, add, mul, add_)	~1.1 s	~25%
Command Buffer Full (launch stalls)	0.299 s	7.5%

magma_sgemmEx = fp32 GEMM still in the path (KDA internals, fp32 head path, NS) — the TF32 knob replaces it with cutlass at 10-bit mantissa.

3. Optimizer internals (TF32 on)

part	ms/step	note
clip (clip_grads, foreach)	107	2228 add_ + norms
opt.step	530	246 mm + 366 bmm calls
— Muon Newton-Schulz matmuls	≈ 320 (mm) + 215 (bmm), overlapped	5 NS iters × every 2-D shape-group
— ColRow + weight decay + AdamW elementwise	≈ 300
zero_grad(set_to_none)	3.7

4. bnb8bit verdict (measured, not assumed)

With the Engram tables OFF the AdamW group is small: bnb costs +174…+200 ms/step and saves little. With the tables ON (+402.7M params → +3.2 GB of fp32 m+v) it is the fit lever, not a speed knob. It is optim_8bit=false by default; the b1 run turns it on only if the tables are in.

5. Shape/VRAM table (engram OFF, tf32 ON, 25-step probes)

shape	tok/s	reserved GiB	verdict
b4·s1024	4,319	11.9	safe, slow
b8·s512	4,318	12.7	safe, slow
b8·s768	5,778	13.5	desktop-safe long-run shape
b16·s512	6,954	14.5	needs desktop closed
b4·s2048	6,937	14.5	needs desktop closed
b8·s1024	6,820	14.6	UNSTABLE — CLI OOM at ~step 16; trainer path peaks ~1.8 GiB above the raw loop

The step cost is ≈950 ms fixed + tokens; tok/s peaks at ≈8192 tok/step. Engram ON adds ≈4.8 GiB (tables’ params+grads+AdamW) → the full “1.07B” preset needs optim_8bit and/or the host-table port (unported) and/or iteration-level activation checkpointing (owed) before it fits at 8k tok/step.

6. Actions, by measured payoff

1. torch.compile — the 2× lever (spike: 273 ms compiled vs 357 eager at the 523M shape). Two attempts died: dynamic=True spent 40+ min in the guard tree; static compile ran 45 min without landing. An OVERNIGHT run is the honest way; the inductor cache persists once it lands.
2. NS steps 5 → 3 — ≈ −130 ms (Muon is the opt’s biggest kernel block); gate: orthogonality probe + a 3-seed A/B at small.
3. clip_grads 107 ms — 2228 add_ calls; one fused pass over a flattened grad list is the burn-era fix repeated at b1 scale (there it took 252 ms → ~5 ms).
4. TF32 as the b1 default — pending only the loss-parity check (tf32_check2.py fixed and rerun; b8·s1024 dies on VRAM before parity matters at that shape — run it at b8·s768).
5. Iteration-level activation checkpointing — the enabler for b16·s512/b4·s2048 (14.5 → ~12 GiB), cost ≈ +20-30% step time.
6. bnb8bit — only with the tables in.