rms_norm_backward_f32

void rms_norm_backward_f32(const float * x, const float * gamma, const float * saved_rstd, const float * g, float * dx, float * dgamma, int outer, int N)

source edit

Single-precision RMSNorm backward — produces input and weight gradients.

Reuses the saved per-row rstd to avoid recomputing the row's RMS. dgamma is accumulated; callers reducing over multiple sub-batches must zero it once at the start. dx is written fully (not accumulated).

Math

With $r = \text{rstd}$ and $c = \sum_j \gamma_j g_j x_j$ ,

\qquad \frac{\partial L}{\partial \gamma_i} = \sum_{\text{rows}} g_i x_i r.$$

Parameters

xconst float*

Original forward input of shape (outer, N).

gammaconst float*

Affine scale of shape (N,).

saved_rstdconst float*

Per-row reciprocal RMS from the matching forward call.

gconst float*

Upstream gradient of shape (outer, N).

dxfloat*

Output gradient w.r.t. x; shape (outer, N); overwritten.

dgammafloat*

Output gradient w.r.t. gamma; shape (N,); accumulated.

outer, Nstd::size_t

Layout extents.