es-ude / es-ude/OnDeviceTraining

LayerNorm float backward: route dgamma/dbeta through the identity-kernel funnel (packed LN grads)

Open
#275 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.