S1-28: GLU_BACK — a VJP for every GLU variant#22
Open
dillon-blake wants to merge 3 commits into
Open
Conversation
ggml could differentiate exactly ONE member of the GLU family: split SwiGLU, via a
SILU_BACK composite. Fused SwiGLU tripped `GGML_ASSERT(src1 && "only split swiglu")`,
and REGLU / GEGLU / GEGLU_ERF / GEGLU_QUICK / SWIGLU_OAI hit
`GGML_ABORT("unsupported glu op for backward pass")`. Every Gemma (GEGLU) and gpt-oss
(SWIGLU_OAI) model was untrainable.
New op GGML_OP_GLU_BACK at the enum tail. One kernel, one row pass, whichever way the
caller packed its operands: dst carries d_a's shape, and a split forward views half 0 as
d_a and half 1 as d_b. Split SwiGLU keeps its existing SILU_BACK composite untouched --
it is the FFN of every dense llama, it is covered today, and routing it through the new
op would change nothing mathematically while risking a regression in the one place this
project cannot afford one.
Derivatives are taken against ggml's OWN scalar forwards in vec.h, not against a paper:
if ggml's gelu uses the tanh approximation, the correct derivative is the derivative OF
THAT, and a finite difference will insist on it.
An odd-width fused `a` was the first thing to bite. ggml_glu computes its output width as
a->ne[0]/2 with INTEGER division, so the sweep's ne_a[0]=5 leaves a trailing element the
forward never reads. It gets a zero gradient -- and it has to be WRITTEN zero, not left as
whatever the allocator handed us, or d_a is not a->ne[0] wide and compute_backward's
same-shape assert fires.
THE TICKET'S HEADLINE BLOCKER IS REFUTED BY MEASUREMENT.
It says the GGML_GELU_FP16 lookup table makes GEGLU's MODE_GRAD unpassable, because the
F32 forward reads an F16 table and the FD would differentiate a step function. Tested by
making the F32 GEGLU forward exact: the error got WORSE (0.024 vs 0.0049), and the
table-free variants moved around just as much. The table is not the cause. That saves a
risky change to inference numerics on every Gemma.
MODE_GRAD IS A WIRING CHECK FOR THIS OP, AND THE COMMENT SAYS SO.
mean_abs_asymm divides by (gn + ga), not (|gn| + |ga|), so a near-zero gradient element
sends the ratio to infinity -- and every GLU gradient is a PRODUCT (dx = dy*g*act'(x)), so
near-zero elements are ordinary, not exotic. Measured: FD noise reaches 0.80 while a
genuinely broken kernel measures 0.18-0.60. They OVERLAP. No tolerance separates them, so
none is pretended: the bound asserts the backward builds, schedules, has the right shapes
and does not abort. That much was impossible before this commit.
Correcting the metric to |gn| + |ga| drops the noise floor to 0.022 and catches every
defect with 3.6-12x margin. It is written and measured and it is NOT here, because it also
removes a NaN-based free pass that TANH, SIGMOID and CROSS_ENTROPY_LOSS have relied on
since S1-19 (0/0 = NaN; `NaN > tol` is false; they pass unconditionally). That is ticket
S1-37.
So the NUMERICS get a real oracle: tests/test-glu-back.cpp, a new ctest binary that
differentiates ggml's own scalar forwards in float64. All six variants match to ~1e-7. It
catches what MODE_GRAD cannot -- perturbing GEGLU's tanh argument by 5% is invisible to
MODE_GRAD and fails here instantly. (It also caught a stale library mid-session, which is
how I know it works.)
test_glu now calls ggml_set_param -- it never did, so the FUSED backward was completely
unexercised while `grad -o REGLU` still printed OK. The param is the BASE tensor, never the
view: ggml_set_param asserts op == GGML_OP_NONE.
added 2 commits
July 14, 2026 13:41
Two bugs in MODE_GRAD's error metric, and they compound.
1. THE DENOMINATOR WAS THE SIGNED SUM.
const float asymm = (a[i] - b[i]) / (a[i] + b[i]);
So the ratio goes to infinity whenever the two gradients nearly cancel, and in
particular whenever the TRUE gradient is near zero and both are rounding noise. That
is not exotic: any op whose output is a product has such elements for ordinary inputs
(every GLU: dx = dy * g * act'(x)), and any op summing over a selected subset can
manufacture them (MUL_MAT_ID's expert routing).
Measured on GLU_BACK, with the kernel verified exact against a float64 reference the
whole time: the "noise" reached MAA 1.80 while a genuinely broken kernel measures
0.18-0.60. The noise OVERLAPPED the defects, and because the metric is unbounded, NO
tolerance was safe. S1-28 had to ship its grad cases with a 0.9 bound -- i.e. a test
that asserted almost nothing -- and it STILL flaked.
Now |a| + |b|. That is what a symmetric relative error is, it is bounded in [-1, 1],
and since |a|+|b| >= |a+b| it can only ever make MAA smaller. GLU's noise floor drops
from 1.80 to 0.022, and its bound is now 5e-2: 2.3x above the noise, 3.6-90x below
every one of six injected defects, all six caught.
2. AN EXACTLY-ZERO GRADIENT WAS AN UNCONDITIONAL PASS.
a == b == 0 gives 0/0 = NaN. `NaN > max_maa_err()` is FALSE. The case passes.
tanh(150) and sigmoid(150) are 1.0 to float precision -- derivative exactly zero -- and
test_unary initializes in [-150, 150]. TANH AND SIGMOID HAVE BEEN ON THE VENDOR-BUMP
ALLOWLIST SINCE S1-19 ON THE STRENGTH OF A NaN. So has CROSS_ENTROPY_LOSS, whose
softmax is one-hot at [-100, 100]. Their gradients had never once been compared.
With the metric fixed they are checked for the first time and fail: MAA 0.14, 0.37 and
0.60. The kernels are fine -- a grad_eps of 15 against a range of 150 measures nothing
a derivative would recognise. Conditioned (range +-3, eps 0.02 for the saturating
unaries; range +-3 for CE) the residual drops to 5.5e-4, 1.1e-2 and 3.8e-3, and each
gets a measured bound. The EVAL sweep keeps its wide range: it is hunting NaNs in the
tails, which is the opposite requirement.
Also pinned, both measured, both previously flaking about 1 run in 15 against the default
1e-4: ADD_ID (identity VJP, noise 1.2e-4 -> 1e-3) and CONCAT (routing VJP, noise 2e-4 ->
1e-3).
Full 16-op grad sweep: 0/15 runs failing.
Found by training one. Neither could have been found any other way, and that is the
point of the end-to-end test.
1. DIV'S BACKWARD DID NOT HANDLE A BROADCAST src1. MUL's always has.
d/d(b) of (a/b) is -a/b^2, which has a's shape. When b is broadcast against a, that
gradient must be REDUCED back onto b's shape with repeat_back -- exactly as MUL does
six lines above it. Without that, ggml_build_backward_expand aborts on its own
same-shape assert.
build_moe_ffn normalizes the top-k router weights with
weights = div(weights[n_used, n_tok], sum_rows(weights)[1, n_tok])
so EVERY MoE model with norm_w -- which is every Mixtral -- hits it the moment a
gradient reaches the router. MUL_MAT_ID having a backward was necessary and nowhere
near sufficient.
2. ggml_get_rows_back REQUIRED A MATRIX GRAD AND A VECTOR INDEX TENSOR.
GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && ...)
while ggml_get_rows itself has always been fully general:
out[i0, i10, i11, i12] = a[i0, b[i10,i11,i12], i11, i12]
So any get_rows with a 3D source and a 2D index tensor had a forward and no backward.
build_moe_ffn gathers each token's top-k router probabilities with exactly that:
weights = get_rows(probs[1, n_expert, n_tok], selected_experts[n_used, n_tok])
Generalized to the real semantics; the old 2-D loop is the i11 == i12 == 0 slice of the
new one. Its flat indexing of src0 was also only correct by accident -- it ignored the
structure entirely, which is why test_get_rows_back could build a grad of [n, r, b]
against a [r/2, b] index tensor and still pass. Those shapes are now consistent.
WHY NO TEST COULD SEE EITHER OF THESE.
test_bin_bcast gated MODE_GRAD on
// The backward pass supports broadcasting only for GGML_ADD:
const bool grad_supported = op == ggml_add && ggml_are_same_shape(a, b[0]) && ...
Two consequences. MUL and DIV were never grad-tested AT ALL, at any shape. And no
BROADCASTING case was ever grad-tested, for any op -- the same-shape clause excluded them.
So DIV's broadcast backward was structurally invisible.
The class was already tuned for it: grad_eps, grad_precise and max_maa_err all branch on
`op == ggml_div` / `op == ggml_mul`. All of that was dead code behind the gate.
Now: ADD 59 grad cases, MUL 59 (previously ZERO), DIV 57, GET_ROWS 111, GET_ROWS_BACK 55.
DIV's bound is 5e-3, measured -- its gradient is second order in b, so its finite
difference is intrinsically noisier than ADD's or MUL's.
Full 21-op grad sweep: 0/12 runs failing. Forwards unregressed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #21 (S1-26).
ggml could differentiate exactly one member of the GLU family: split SwiGLU, via a
SILU_BACKcomposite. Fused SwiGLU tripped
GGML_ASSERT(src1 && "only split swiglu"), and REGLU / GEGLU /GEGLU_ERF / GEGLU_QUICK / SWIGLU_OAI hit
GGML_ABORT("unsupported glu op for backward pass").Every Gemma (GEGLU) and gpt-oss (SWIGLU_OAI) model was untrainable.
New op
GGML_OP_GLU_BACKat the enum tail. One kernel, one row pass, whichever way the callerpacked its operands. Split SwiGLU keeps its existing composite untouched — it is the FFN of
every dense llama and the one path this project cannot afford to regress.
Derivatives are taken against ggml's own scalar forwards in
vec.h, not against a paper: ifggml's gelu uses the tanh approximation, the correct derivative is the derivative of that.
The ticket's headline blocker is refuted by measurement
The ticket says the
GGML_GELU_FP16lookup table makes GEGLU's MODE_GRAD unpassable — the F32forward reads an F16 table, so the FD would differentiate a step function.
Tested it. Made the F32 GEGLU forward exact: the error got worse (0.024 vs 0.0049), and the
table-free variants moved just as much. The table is not the cause. That saves a risky change to
inference numerics on every Gemma.
MODE_GRAD is a wiring check for this op, and the code says so
mean_abs_asymmdivides by(gn + ga), not(|gn| + |ga|), so a near-zero gradient element sendsthat ratio to infinity — and every GLU gradient is a product (
dx = dy·g·act'(x)), so near-zeroelements are ordinary. Measured:
act'(x)fromdxact(x)fromdgx(1-s)termThe noise overlaps the defects. No tolerance separates them, so none is pretended — the bound
asserts the backward builds, schedules, has the right shapes and does not abort. That much was
impossible before this commit.
Correcting the metric to
|gn| + |ga|drops the noise floor to 0.022 and catches every defectwith 3.6–12× margin. It is written and measured and deliberately not in this PR: it also removes
a NaN-based free pass that
TANH,SIGMOIDandCROSS_ENTROPY_LOSShave relied on since S1-19(
0/0 = NaN;NaN > tolis false; they pass unconditionally). That is ticket S1-37.So the numerics get a real oracle
tests/test-glu-back.cpp— a new ctest binary that differentiates ggml's own scalar forwards infloat64. All six variants match to ~1e-7.
It catches what MODE_GRAD cannot: perturbing GEGLU's tanh argument by 5% is invisible to MODE_GRAD
and fails here instantly. (It also caught a stale library mid-session, which is how I know it works.)
Also
test_glunever calledggml_set_param, so the fused backward was completely unexercised whilegrad -o REGLUstill printedOK. It does now — on the base tensor, never the view(
ggml_set_paramassertsop == GGML_OP_NONE).An odd-width fused
awas the first thing to bite:ggml_gluhalves the width with integerdivision, so
ne_a[0]=5leaves a trailing element the forward never reads. It gets a zero gradient— and it must be written zero, or
d_aisn'ta->ne[0]wide andcompute_backward's same-shapeassert fires.