Three backward-pass gaps blocking standard ViT training: LayerNorm, GELU, conv_2d/CONT

[OriPekelman]

Three independent backward-pass gaps that block training a vanilla ViT-Tiny (and likely any vision transformer) on ggml's auto-backward path. Found while building a Ruby/Spinel ML framework on top of ggml; happy to provide reproducers / discuss vendoring patches.

1. GGML_OP_NORM (LayerNorm) — no backward

ggml_build_backward_expand aborts at ggml.c:6874:

ggml_compute_backward: unsupported ggml op for backward pass: NORM

GGML_OP_RMS_NORM has a registered backward (ggml_rms_norm_back); LayerNorm doesn't. This forces every transformer that wants ggml training to swap LayerNorm → RMSNorm. For LLMs (Llama family) that's natural since they already use RMSNorm; for ViT variants and GPT-2-shape models, it's a real conversion cost (with downstream effects on weight-loader compat).

2. GGML_UNARY_OP_GELU — no backward

ggml_build_backward_expand aborts at ggml.c:6843 with the same shape:

ggml_compute_backward: unsupported unary op for backward pass: GELU

The unary backward dispatch switch (ggml.c:6780-6845) handles ABS, SGN, NEG, STEP, RELU, SILU, EXP, EXPM1, SOFTPLUS — but not GELU (or GELU_ERF, GELU_QUICK). Same effect: ViT MLPs (standard GELU FFN) can't backprop and must swap to SiLU.

A ggml-side ggml_gelu_back doesn't appear to exist; the GELU derivative is well-known (0.5 * (1 + tanh(sqrt(2/π)*(x + 0.044715*x³))) + ...) and could be added either as a new op or as a synthesized backward chain (it's a small composition of existing ops).

3. ggml_conv_2d backward asserts contiguous gradient

ggml.c:6649: GGML_ASSERT(ggml_is_contiguous(grad)) failed

Source of the issue: ggml_conv_2d's implementation (ggml.c:4585-4604) ends in ggml_cont(ggml_permute(...)). The auto-backward for GGML_OP_CONT (ggml.c:6645-6655) asserts the incoming gradient is contiguous. In a transformer-shaped graph the gradient flowing back from the next layer into the conv output isn't always contiguous.

Workarounds we've identified:

• Implement patch_embed as a flat linear matmul(W_patch, flat_patches) instead of conv_2d. Mathematically equivalent to a stride=patch / no-overlap / no-padding conv. Trivial — but loses generality (overlapping conv, non-square strides, etc.).
• Vendor a conv2d implementation that ends differently, OR make the auto-bw of OP_CONT materialise a contiguous grad before the assert.

Why this matters

Together these three gaps mean ggml cannot train a standard ViT through its native auto-backward. Each individually has a low-effort vendoring path; together they would unblock vision-transformer training. Happy to write the fixes if there's interest in the approach.

Repro environment

• Linux aarch64 (NVIDIA GB10)
• ggml HEAD as vendored in our repo (May 2026)
• Auto-backward via ggml_build_backward_expand, no manual gradient construction
• Single-image batch (N=1)

Source where the three workarounds are pinned in code:

• lib/vit_tiny_forward_ffi.rb — header comment + per-op notes
• Repository: OriPekelman/toy

[OriPekelman]

Update: gaps 1 & 2 implemented + validated (CPU reference)

We needed these for GPT-2 train-from-scratch (LayerNorm + GELU FFN), so I implemented the first two as dedicated fused backward ops, modeled on the existing ggml_rms_norm_back / ggml_silu_back siblings. Both gradient-check clean.

Gap 1 — GGML_OP_NORM → new GGML_OP_NORM_BACK / ggml_norm_back(ctx, x, dy, eps).
Key simplification: ggml_norm is the pure normalize y=(x-μ)/√(var+ε) with no affine, so only the normalize backward is needed — GPT-2/ViT's γ/β are separate mul/add that already have backward. The kernel is rms_norm_back plus the mean-centering term:

dx_i = rstd * ( dy_i − mean(dy) − xhat_i · mean(dy·xhat) ),  xhat=(x−μ)·rstd

Autograd: case GGML_OP_NORM → ggml_norm_back(ctx, src0, grad, eps) (eps read from op_params, like RMS_NORM).

Gap 2 — GGML_UNARY_OP_GELU → new GGML_OP_GELU_BACK / ggml_gelu_back(ctx, dy, x).
Elementwise dy · gelu'(x) with the tanh-approx derivative:

g = √(2/π)·x·(1+a·x²);  t=tanh(g);  g' = √(2/π)·(1+3a·x²)
gelu'(x) = 0.5(1+t) + 0.5·x·(1−t²)·g'      (a = 0.044715)

We went dedicated kernel rather than synthesized subgraph — the synthesized route needs scalar-add (ggml_add1 deprecated) + per-row broadcast for NORM's mean, which is fiddly and error-prone; the dedicated kernels copy the silu_back/rms_norm_back plumbing and have directly-verifiable math.

Gotcha worth flagging for an upstream version: the CPU gelu forward uses an f16 lookup table (GGML_GELU_FP16). The analytic backward is the derivative of the exact tanh gelu, not of the f16-table forward — fine for training, but a finite-difference check must difference the exact function or f16 buckets dominate the result.

Validation: autograd grad on the real ggml_backend_sched_graph_compute (CPU) path vs an independent pure-C central finite difference (numpy-allclose, atol 1e-3 / rtol 2e-2) — both PASS at the f32 FD rounding floor (gelu max_abs 5e-5, norm 2e-4), on aarch64. (Side note: while doing this I re-confirmed #1491 — rms_norm_back is still wrong on the sched path, while the new norm_back is correct on the identical path.)

Gap 3 (conv_2d/CONT contiguous-grad assert) is ViT-only and we haven't tackled it yet.

If there's interest I'm happy to open a PR for the two backward ops (the diff touches ggml.h/ggml.c op tables + autograd dispatch, ggml-cpu/{ops.cpp,vec.h}, and the alloc/meta switches).