Train a Transformer on Silicon: #3 The Memory Is the Chip

Part of the Train a Transformer on Silicon series — start at #0. This one pays off the cliffhanger from #2.

In #2 I built a decoder layer that trains itself on-chip, and the die came out 85% SRAM. The compute — GEMM, LayerNorm, softmax, backprop through all of it — fit in 0.117 mm². The memory did not fit in anything reasonable. This post is what I did about it.

TL;DR: I gave back a third of the die without touching one logic gate. A liveness analysis on v0’s 56 buffers, plus rescheduling each weight’s SGD update to fire early, shrank the memory map from 145,952 words to 90,896 (−38%) and the die from 6.99 mm² to 4.55 mm² (−35%). Logic area, cell count, power, and timing are all within noise of v0 — exactly what a memory-only optimization should produce. And the liveness pass found a buffer that was written but never read: 1,024 words of pure waste nobody would have caught by eye.

The problem: a memory map that never frees anything

v0’s memory map is the simplest correct thing: every intermediate the forward pass produces gets its own buffer so the backward pass can read it later; nothing is reused, nothing freed. That’s 56 buffers, 145,952 words × 32 bits ≈ 4.67 Mbit — the reason the die is mostly memory. Everything in this post is memory-subsystem sizing: how many words the training step actually needs, and when.

But “later” is only a handful of ops, after which the buffer is dead — its words could hold something else. Most of those 56 buffers are alive for a small slice of the 50-op program and dead for the rest. If I can prove which slices, non-overlapping buffers can share physical words. That’s the whole idea of H01 — a classic register-allocation problem wearing an SRAM costume.

The mechanism: liveness → linear-scan → early-SGD reschedule

Three steps, all in the Python generator that emits the memory map, so the packing is decided before any RTL is touched.

1. Liveness analysis. For each buffer I compute its live interval [\text{first write}, \text{last read}] across the 50-op schedule, with per-unit read/write-port semantics: a unit that reads and writes the same slot in one op can’t have that slot aliased to another live buffer, or it clobbers itself. The pass also reports anything suspicious:

v0 flat map: TOTAL=145952 words, AW=19, NPC=50
transients: 44; written-but-never-read: ['dxs']

v0 flat map: TOTAL=145952 words, AW=19, NPC=50
transients: 44; written-but-never-read: ['dxs']

dxs — 1,024 words the forward pass computes and nothing ever reads. A dead store. It costs memory and it’s a code smell (why is that even being written?), and liveness flags it for free.

2. Linear-scan, first-fit allocation. Sort the live intervals and drop each buffer into the first free address range large enough to hold it, releasing ranges as intervals end. Two buffers that are never alive at the same time land on the same words.

3. Early-SGD reschedule. This is the multiplier. In v0 the ten SGD weight updates run at the very end (ops 40–49), so every gradient stays alive from the moment it’s computed all the way to step end. Instead, I move each weight’s update to fire immediately after that weight’s last read in the step — the gradient’s live range ends right there:

SGD(W2) moved: op 45 -> after body-op 17
SGD(W1) moved: op 44 -> after body-op 20
SGD(Wo) moved: op 43 -> after body-op 24
SGD(Wq) moved: op 40 -> after body-op 34
...

SGD(W2) moved: op 45 -> after body-op 17
SGD(W1) moved: op 44 -> after body-op 20
SGD(Wo) moved: op 43 -> after body-op 24
SGD(Wq) moved: op 40 -> after body-op 34
...

Pack-only gets me to 110,352 words (−24%). Pack plus the reschedule gets to 90,896 (−38%), because collapsing the gradient live ranges drops the peak of simultaneously-live memory from 105,872 words to 85,536.

4. Legality by construction. The generator asserts that no op ever reads a weight or its gradient after that weight’s update point; the build fails if it does. The reschedule is legal because a check says so, not because I stared at it.

The picture tells the story: 94,496 words of live-once activations, gradients, and scratch fold into a single 39,440-word region everything time-shares — a 2.4× compression. What doesn’t move is the persistent floor: the input, target, and ten weight tensors, read at step start and written at step end, impossible to recycle. That’s 51,456 words, and no scheduler crosses it.

Why it works: most buffers are dead most of the time

Plot the working set — the words actually alive — at each op and the argument becomes visual. The curve rides far below the flat map’s 145,952-word ceiling for essentially the whole step. That entire gap is memory v0 kept resident and H01 proves is dead. The reschedule (the bold curve) also keeps the peak lower and earlier than pack-only, because gradients stop living to step end.

One honest wrinkle: the allocator lands at 90,896 words even though the peak live set is only 85,536 — first-fit leaves fragmentation. A smarter allocator could recover some of that ~5,000-word gap, but the win already changed the address width, so I shipped it.

Verification: does the packed netlist still learn?

Same protocol as always — golden model first, then the RTL has to match it to the last bit. Re-running v0’s checks against the H01 netlist:

• 1-step training, all 20 tensors (10 gradients + 10 updated weights) vs golden: 0 LSB — exact.
• Multi-step loss trajectory: identical to golden over 120 steps — step for step, from the opening loss of 1.20951708 onward. Packing changed no numbers.
• SGD-update legality: assert-checked, passes.

The war story: packing broke my gradient check

Here’s the part that bit me. In v0, verifying gradients was easy: run a step, read all ten gradient tensors out of the RAM at the end, compare to golden. In H01 that read returns garbage — and correctly so. The whole point of packing is that a gradient’s slot gets recycled the instant its SGD update fires. By end-of-step, those addresses hold some later buffer. My trusted verification method was reading dead memory and comparing noise to golden.

The design was fine; the harness was reading at the wrong time. The fix is a mid-step hierarchical RAM snapshot: a read fires at the exact issue cycle of each SGD op, capturing that gradient the instant before its slot is reused. Same bit-exact guarantee as v0’s end-of-step check — relocated in time. (Second time in this series the harness lied while the design was innocent; in #2 it was stale testbench plumbing. When the numbers look wrong, suspect the measurement before the design.)

Physical results: logic untouched, die shrunk

RTL trusted, I pushed H01 through the same flow — Genus synthesis, Innovus P&R, educational 45 nm PDK, SRAM as a blackboxed macro. H01 against its v0 parent — the PPA scorecard (power, performance, area, side by side):

Everything below the die row is flat — the evidence that only memory moved: the packing rewrote the address map, not the datapath, so not one logic gate changed. (The tiny cell-count delta is a shared defensive-RTL workaround, not the packing — see caveats.) The die shrink tracks the memory shrink, quantized through the macro library: AW19 → AW18, plus fixed logic and channel that don’t scale.

The floorplan is the same story at a smaller scale: the SRAM macro that was 85% of v0’s die is now ~82% of a much smaller one, the datapath still packed into its thin west-edge channel. The datapath — unchanged — is now 2.6% of the die instead of under 2%; it didn’t grow, the die shrank around it. And the block ranking is exactly as it was in #2: LayerNorm backward is still the biggest logic block, 30.7% of logic and 4× the GEMM engine. The cost of strict nanoGPT fidelity didn’t go anywhere; the memory just stopped hiding it.

How low can this go? Not much, by allocation alone. The persistent state is 51,456 words and the peak working set is 85,536 — both above what an AW17 macro holds — so AW18 is the floor. Getting under it isn’t a packing problem; it needs weight streaming or gradient recompute, a different design.

Honest caveats

• Tiny model, educational PDK. D = 64, one layer, one head, T = 16, generic 45 nm — a mechanism demonstration, not a competitive accelerator, and not signoff-grade.
• The SRAM area is modeled, not compiled. The die numbers use a hand-built LEF macro at 40.82 µm²/word. As a cross-check, CACTI 7 at 45 nm puts the data arrays 1.5–2× smaller, so my model is conservative — a real SRAM compiler would give a smaller, foundry-specific macro. Die claims stay labeled modeled-not-compiled.
• Power is vectorless and logic-only. The macro is a blackbox to the tools, so on a die that’s 82% SRAM the reported 1.39 mW is not the real total.
• Fixed-netlist context. Both v0 and H01 carry a defensive-RTL fix (gemm_engine_v3) that works around a Genus mapping defect I hit mid-project; it adds ~340 cells and leaves die, power, and timing essentially unchanged. H01 re-synthesized, re-routed, and re-passed DRC (0 violations, fresh session).
• The floor is real. ~51,456 persistent words is a hard limit for allocation; only recompute or streaming goes lower.

The lesson

Profile before you optimize, then optimize where the area actually is. v0’s logic was under 2% of the die — tuning it would have been effort on a rounding error. Reclaiming memory, with nothing more exotic than liveness analysis and a schedule tweak, cut the die by a third with the datapath untouched. And liveness pays a bonus hand-inspection can’t: it surfaced a dead store (dxs, 1,024 words written and never read) sitting in the map the whole time.

Next up — #4: Giving It a Vocabulary. So far this thing trains on a copy-shift task with no tokens — it’s a decoder with nothing to decode. To make it a real language model it needs an embedding table and an output projection over a vocabulary, which lands a large new persistent tensor right on top of the 51,456-word floor liveness just told me it can’t touch. The floor is about to get taller. That’s the next post.

The full roadmap lives in Part 0 — the hub for the series.

If you build hardware, do ML, or work in EDA — follow along. And if you catch a mistake, tell me. For a series like this, that’s the best thing that can happen.