es-ude / es-ude/OnDeviceTraining
LayerNorm float backward: route dgamma/dbeta through the identity-kernel funnel (packed LN grads)
- Dominant language
- C
- Stars
- 1
- Forks
- 3
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 8
Description
# LayerNorm float backward: route dgamma/dbeta through the identity-kernel funnel (packed LN grads)
**Status:** Draft for Leo confirmation before filing at PR time (PR3 Task 6).
**Context:** `layerNormBackwardFloat` (`src/layer/LayerNorm.c`) is the FLOAT32-arithmetic
backward path for LayerNorm. Unlike `layerNormBackwardSymInt32` — whose dgamma/dbeta
increments are written through `executeOp` with `cfg->weightGradAccMode` /
`cfg->biasGradAccMode` (the identity-kernel funnel, `executeOpIdentityKernel`,
LayerNorm.c around the two `executeOp(&(opSpec_t){.kernel = executeOpIdentityKernel, ...`
calls after the grad-increment accumulation loop) — `layerNormBackwardFloat` raw-casts
`cfg->gamma->grad->data` / `cfg->beta->grad->data` straight to `float*` and does
`dgamma[j] += ...` / `dbeta[j] += ...` directly on the buffer. It never goes through
`executeOp`, so it has no accumulate-mode awareness and, more importantly, no dtype
awareness: it assumes the grad tensor is a flat `float[]` unconditionally.
PR3 Task 6 added a fail-fast dtype guard in `layerNormBackward`'s `ARITH_FLOAT32` case
(the block right after the existing forwardInput/loss/gamma->param FLOAT32 check): if
`cfg->gamma->grad->quantization->type` or `cfg->beta->grad->quantization->type` is not
`FLOAT32`, it `PRINT_ERROR`s (naming the raw-cast reason and this follow-up) and
`exit(1)`s instead of corrupting a packed grad buffer. That guard is a stopgap, not a
fix: it makes packed (SYM/ASYM) grad storage for LayerNorm's FLOAT32 backward path an
explicit hard error rather than silent memory corruption, but it does not make packed LN
grads possible.
**Precedent:** the SYM_INT32 backward path already shows the target shape. It builds a
transient FLOAT32 scratch tensor for the per-call increment (`dgammaT`/`dbetaT` via
`setTensorValues` over a stack VLA), then calls
`executeOp(&(opSpec_t){.kernel = executeOpIdentityKernel, .inputs = {&dgammaT}, .nInputs = 1,
.arithmetic = {.type = ARITH_FLOAT32, .roundingMode = HALF_AWAY}, .mode = cfg->weightGradAccMode},
cfg->gamma->grad)` (and the beta twin with `cfg->biasGradAccMode`). The funnel's
`accumulateOut` epilogue (`ExecuteOp.c`) already has FLOAT32/SYM_INT32/SYM/ASYM
accumulate arms (PR3 §4) — routing the float backward's dgamma/dbeta increments through
the same identity-kernel call, instead of the raw `+=` on a `float*`, would let
`cfg->gamma->grad` / `cfg->beta->grad` legitimately be packed SYM/ASYM tensors, closing
the gap this task's guard only detects.
**Relationship to #261 / PR3:** #261 established that grads must never be stored
SYM_INT32 (compute-only), and PR3 (packed sub-byte grad storage,
`docs/superpowers/specs/2026-07-03-pr3-packed-grad-storage-design.md` §7) explicitly
scoped this LayerNorm float-backward funnel migration out as a named follow-up — the
spec's exact words: "Follow-up issue (filed at PR time): route LN float dgamma/dbeta
through the identity-kernel funnel like its SYM_INT32 path, so packed LN grads become
possible."
**Task (target: post-PR3 follow-up):**
- [ ] Refactor `layerNormBackwardFloat`'s dgamma/dbeta accumulation to (a) accumulate the
per-call increments into local FLOAT32 scratch (mirroring the SYM_INT32 path's
`dgammaInc`/`dbetaInc` VLAs), then (b) write them into `cfg->gamma->grad` /
`cfg->beta->grad` via `executeOp` + `executeOpIdentityKernel`, honoring
`cfg->weightGradAccMode` / `cfg->biasGradAccMode` — same call shape as the
SYM_INT32 backward already uses.
- dx (`propLoss`) is unaffected — it stays a direct `float*` write (FLOAT32-arithmetic
backward output is not eligible for packed grad storage; only gamma/beta grads are).
- [ ] Remove (or narrow) the Task 6 dtype guard once the funnel route makes non-FLOAT32
grads a supported path rather than a hard error — decide whether any residual
guard is still needed (e.g. still rejecting BOOL/INT32).
- [ ] Extend `UnitTestLayerNorm.c` with a positive test: FLOAT32-backward math + packed
SYM gamma/beta grad storage, asserting the grads land correctly through the
funnel's accumulate arm (mirrors `testFactoryFullSymProfileTrainsSymGrads`, which
exercises the same thing on the SYM_INT32 backward path).
- [ ] Bit-parity: existing FLOAT32-grad LN tests (`testBackwardFloat*`, `testGold*Backward*`)
must stay verbatim green — the refactor changes only the write path for grads
whose storage is FLOAT32 already (the funnel's FLOAT32 arm is a direct write,
byte-identical to today's raw `+=`), not gamma/beta grad *values*.
**Recon source:** PR3 spec §7 (`docs/superpowers/specs/2026-07-03-pr3-packed-grad-storage-design.md`);
guard added in `src/layer/LayerNorm.c` `layerNormBackward` (PR3 Task 6, this issue's
originating change).
Contributor guide
Research direction
Read src/layer/LayerNorm.c alongside the SYM_INT32 backward path and the identity-kernel accumulate logic in ExecuteOp.c. Run the existing LayerNorm backward tests in UnitTestLayerNorm.c first; done means FLOAT32 backward tests remain unchanged and a new packed SYM gamma/beta gradient test passes without the current dtype guard rejecting supported storage.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100