Phase 1, Section 25: StoryLM-0 decoder-wide AdamW

StoryLM-0 train slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses Section 15 NLL ∘ reverseMode and Section 16 AdamW, one copy per trainable weight.

1. What this section commits to

Section 16 locks AdamW on one [4, 8] tensor. Section 15 applies SGD to every decoder weight. Section 24’s host loop still uses SGD. This section applies Section 16 to each trainable weight after the Section 15 projection.

The trainable set is

{token_embedding, wq, wk, wv, wo, lm_head}

g_mask and g_scale are bound and unused. Moments m and v start at host zeros. t lives on the host; each call takes β1 ** t and β2 ** t.

The numeric train instance is the Section 24 corpus, 40 epochs, lr = 0.2, β1 = 0.9, β2 = 0.999, ε = 1e-8, λ = 0.01. After training, Section 23 greedy from [1, 2, 3, 4] must emit the same locked cycle as Section 24.

There is no adamw keyword and no tuple-map combinator. Write one Section 16 composition per tensor.

2. Why not a primitive

Section 16 already is the algebra. A decoder-wide fused update would invent a parameter bundle ABI. Independent (w', m', v') results per tensor are the existing composition.

3. Operands (numeric instance)

Decoder widths match Section 15. AdamW scalars match Section 16 except lr = 0.2 for the train loop (the one-step check uses the first corpus window, t = 1, zero moments, and those same scalars).

Each weight has the same shape as its g, m, and v:

Tensor	Shape
`token_embedding`	`[8, 4]`
`wq` `wk` `wv` `wo`	`[4, 4]`
`lm_head`	`[4, 8]`

4. Operational meaning

(nll, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm)
  = reverseMode storylm0_nll(…)
(emb', m_emb', v_emb') = adamw(emb, g_emb, m_emb, v_emb, …)
(wq',  m_wq',  v_wq')  = adamw(wq,  g_wq,  m_wq,  v_wq,  …)
… same rule for wk, wv, wo, lm_head …

Host train loop (Section 24 windows, t counts every window):

t = 0
for epoch in 0 .. 39:
    for (ids, targets) in corpus:
        t = t + 1
        (nll, w, m, v) = storylm0_adamw_step(ids, targets, w, m, v, β**t, …)

Corpus sum-NLL must fall below 0.5. Greedy eight ids prefixed by the prompt must equal [1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6].

5. Typing (derived)

Each AdamW copy is Section 16 typing at that weight’s shape. The wrapper result is a constructed flat tuple:

(nll, emb', m_emb', v_emb', wq', m_wq', v_wq', wk', m_wk', v_wk',
 wv', m_wv', v_wv', wo', m_wo', v_wo', lm', m_lm', v_lm')

6. Surface instance

storylm0_adamw_step.kina binds shape-specific adamw helpers, storylm0_nll, and step. @main calls step with host placeholders. The numeric check loads step from the .so and also storylm0_rollout.kina for greedy. Hyperparameters live in tests/compiler/fixtures/storylm0/adamw.json; the Section 24 corpus is unchanged.

No new tokens. AdamW sits outside the differentiated function.

7. Tests

storylm0_adamw_step.kina — §6 program; golden accepted by mlir-opt; step result is the §5 tuple.
scripts/validate_decoder_adamw.py compiles that example and storylm0_rollout.kina to .so files, checks the first window’s (w', m', v') against NumPy / PyTorch Section 16, runs the §4 loop, and checks corpus sum-NLL < 0.5 plus the locked greedy cycle.

8. Out of scope

Mean NLL as this loop’s objective (see Section 26).
Differentiating through AdamW.
AMSGrad, fused kernels, or an in-graph step clock.
A real corpus or production checkpoint ABI.
Phase 1 exit criteria.

9. Open issues

Mean NLL. Specified in Section 26. This loop stays sum-NLL.
Host pow .so lowering. Unchanged from Section 16.
Trained 32-token fixture. Eight greedy ids, not a 32-token prose check.
Prefill-T family / in-place write. Unchanged.