LayerNorm

LayerNorm has two learned vectors and a hand-written backward pass. This page covers both directions, then derives the backward formula from the forward definition.

LayerNorm forward

What is it?

LayerNormImplemented

Basic operations

Normalises each row of its input on its own, so that the row has a mean near 0 and a variance near 1 across its d features. It then scales and shifts each feature with two learned vectors.

Why does the model need it?

Each block gives attention and the feed-forward layer a normalised copy of its input, and the head reads a normalised copy too. The statistics are computed per row, so the code needs no batch axis and keeps no running averages.

Note

In general, normalising inputs keeps their scale steady, so later layers receive numbers in a similar range whatever earlier layers produce, which tends to make training more stable. This is a general explanation. Nothing in this repository measures that effect.

Mathematics

Each row is handled on its own, over its last axis.

Mean of one row, over the last axis #

\mu_t = \frac{1}{d} \sum_{i} X_{ti}

The mean of the squares divides by d, not by d - 1, so this is the biased variance.

Biased variance of one row. It divides by d #

\sigma_t^2 = \frac{1}{d} \sum_{i} \left(X_{ti} - \mu_t\right)^2

Inverse standard deviation #

s_t = \frac{1}{\sqrt{\sigma_t^2 + \epsilon}}

The normalised row #

\hat{X}_{ti} = \left(X_{ti} - \mu_t\right) s_t

Scale and shift, one pair per feature #

Y_{ti} = \gamma_i \hat{X}_{ti} + \beta_i

The function returns the output and a cache, (Xhat, s). The backward pass needs those two arrays, and nothing else from the forward pass.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X$	Input, one row per token	`X`	(T, d)
$\gamma$	Learned scale for each feature. A parameter	`gamma`	(d,)
$\beta$	Learned shift for each feature. A parameter	`beta`	(d,)
$\epsilon$	Small number that keeps s finite when a row has no spread	`epsilon`	scalar
$\mu$	Mean of each row. keepdims keeps it broadcastable	`mu`	(T, 1)
$\sigma^2$	Variance of each row	`var`	(T, 1)
$s$	Inverse standard deviation of each row. Kept in the cache	`s`	(T, 1)
$\hat{X}$	Normalised input. Kept in the cache	`Xhat`	(T, d)
$Y$	Output. The function returns it without naming it	not named in the code	(T, d)

Tiny numerical example

Edit the row, gamma or beta. Setting all four values of X to the same number shows what epsilon is for.

One row through LayerNormIllustrative ExampleHand-picked numbers from test_layer_norm_worked_example that you can edit. These are not learned values.

X
	i0	i1	i2	i3
t0

gamma
	i0	i1	i2	i3
gamma

beta
	i0	i1	i2	i3
beta

Mean: 3

Variance: 3.5

s: 0.5345

Xhat, shape (1, 4)
-1.069	-0.5345	0	1.6036

Output, gamma * Xhat + beta, shape (1, 4)
-1.069	-0.5345	0	4.2071

Tensor shapes

The shape does not change. gamma and beta have shape (d,) and broadcast over the rows.

Before attentiontransformer.py:133

X(T, d)(16, 32)
Xn(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Before the feed-forward layertransformer.py:136

X1(T, d)(16, 32)
Xn2(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

At the end of the model, before the headtransformer.py:165

X(T, d)(16, 32)
Xf(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Actual implementation

ops.pylayer_norm_forwardGitHub

42def layer_norm_forward(X, gamma, beta, epsilon = 1e-5):43    mu = X.mean(axis=-1, keepdims=True)44    var = ((X - mu) ** 2).mean(axis=-1, keepdims=True)45    s = 1.0 / np.sqrt(var + epsilon)46    Xhat = (X - mu) * s47    return gamma * Xhat + beta, (Xhat, s)

Code and math connection

Equation	Role	Lines	Code
Mean of one row, over the last axis	Computes it	ops.py:43	`mu = X.mean(axis=-1, keepdims=True)`
Biased variance of one row. It divides by d	Computes it	ops.py:44	`var = ((X - mu) ** 2).mean(axis=-1, keepdims=True)`
Inverse standard deviation	Computes it	ops.py:45	`s = 1.0 / np.sqrt(var + epsilon)`
The normalised row	Computes it	ops.py:46	`Xhat = (X - mu) * s`
Scale and shift, one pair per feature	Computes it	ops.py:47	`return gamma * Xhat + beta, (Xhat, s)`
Scale and shift, one pair per feature	Calls it	transformer.py:133	`Xn, c_ln1 = layer_norm_forward(X, params[p + "ln1_g"], params[p + "ln1_b"], cfg.ln_eps) # make a normalized copy for attention to read`
Scale and shift, one pair per feature	Calls it	transformer.py:136	`Xn2, c_ln2 = layer_norm_forward(X1, params[p + "ln2_g"], params[p + "ln2_b"], cfg.ln_eps) # normalized copy again, for the feed-forward`
Scale and shift, one pair per feature	Calls it	transformer.py:165	`Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)`
Scale and shift, one pair per feature	Tests it	tests/test_ops.py:73	`y, cache = layer_norm_forward(x, gamma, beta)`

Where is it used?

Source

ops.py, lines 42-47, layer_norm_forward

PDF reference

Section 4.3, Eq. 30, 31, p. 7

Used by

Called by: transformer.block_forward, transformer.forward

Tested by: tests.test_ops.test_layer_norm_worked_example

Calls it: transformer.py, line 133 (block_forward), for Scale and shift, one pair per feature
Calls it: transformer.py, line 136 (block_forward), for Scale and shift, one pair per feature
Calls it: transformer.py, line 165 (forward), for Scale and shift, one pair per feature
Tests it: tests/test_ops.py, line 73 (test_layer_norm_worked_example), for Scale and shift, one pair per feature

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Mean of one row, over the last axis	Computes it	`ops.py`	`layer_norm_forward`	43	GitHub: ops.py, line 43
Biased variance of one row. It divides by d	Computes it	`ops.py`	`layer_norm_forward`	44	GitHub: ops.py, line 44
Inverse standard deviation	Computes it	`ops.py`	`layer_norm_forward`	45	GitHub: ops.py, line 45
The normalised row	Computes it	`ops.py`	`layer_norm_forward`	46	GitHub: ops.py, line 46
Scale and shift, one pair per feature	Computes it	`ops.py`	`layer_norm_forward`	47	GitHub: ops.py, line 47

LayerNorm backward

What is it?

LayerNorm, backwardImplemented

Basic operations

Given dY, the gradient on LayerNorm's output, and the cache from the forward pass, this computes three gradients. They are dX for the input, dgamma for the scale and dbeta for the shift.

Why does the model need it?

dgamma and dbeta are what the optimizer needs to update this layer. dX carries the gradient back through the normalisation to whatever produced the input.

Mathematics

Gradient on the normalised row #

d\hat{X}_{ti} = dY_{ti}\, \gamma_i

Gradient on the input #

dX_{tj} = s_t \left( d\hat{X}_{tj} - \operatorname{mean}_i\!\left(d\hat{X}_{ti}\right) - \hat{X}_{tj}\, \operatorname{mean}_i\!\left(d\hat{X}_{ti} \hat{X}_{ti}\right) \right)

Gradient on the scale, summed over rows #

d\gamma_i = \sum_{t} dY_{ti}\, \hat{X}_{ti}

Gradient on the shift, summed over rows #

d\beta_i = \sum_{t} dY_{ti}

The next section derives these, one step at a time.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$dY$	Gradient of the loss on the output Y. The caller supplies it	`dY`	(T, d)
$\hat{X},\, s$	The cache from the forward pass	`Xhat, s`	(T, d) and (T, 1)
$\gamma$	The scale parameter	`gamma`	(d,)
$d\hat{X}$	Gradient on the normalised row	`dXhat`	(T, d)
$dX$	Gradient on the input, passed to the previous layer	`dX`	(T, d)
$d\gamma$	Gradient on the scale	`dgamma`	(d,)
$d\beta$	Gradient on the shift	`dbeta`	(d,)
$\operatorname{mean}_i$	Mean over the last axis, row by row	`.mean(axis=-1, keepdims=True)`	(T, 1)

Tiny numerical example

A worked backwardIllustrative ExampleHand-picked numbers from test_layer_norm_worked_example. dY is chosen by hand, not computed from a loss.

Step	Code	Value
Gradient on the output	`dY`	`[[1.0000, 0.0000, 0.0000, 1.0000]]`
Cache	`Xhat`	`[[-1.0690, -0.5345, 0.0000, 1.6036]]`
Cache	`s`	`[[0.5345]]`
Gradient on the normalised row	`dXhat = dY * gamma`	`[[1.0000, 0.0000, 0.0000, 2.0000]]`
First mean	`dXhat.mean(axis=-1, keepdims=True)`	`[[0.7500]]`
Second mean	`(dXhat * Xhat).mean(axis=-1, keepdims=True)`	`[[0.5345]]`
Gradient on the input	`dX`	`[[0.4391, -0.2482, -0.4009, 0.2100]]`
Gradient on the scale	`dgamma = (dY * Xhat).sum(axis=0)`	`[-1.0690, 0.0000, 0.0000, 1.6036]`
Gradient on the shift	`dbeta = dY.sum(axis=0)`	`[1.0000, 0.0000, 0.0000, 1.0000]`

One entry: dX[0][0] = 0.5345 × (1 - 0.75 - (-1.069) × 0.5345) = 0.4391

Tensor shapes

dX has the shape of X. dgamma and dbeta have the shape of gamma and beta, because the sum over rows removes the first axis.

The final LayerNormtransformer.py:181

dXf(T, d)(16, 32)
dX(T, d)(16, 32)
dg(d,)(32,)
db(d,)(32,)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

LayerNorm before attention, inside a blocktransformer.py:149

dXn(T, d)(16, 32)
dX_branch(T, d)(16, 32)
dg(d,)(32,)
db(d,)(32,)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Actual implementation

ops.pylayer_norm_backwardGitHub

49def layer_norm_backward(dY, cache, gamma):50    Xhat, s = cache51    dXhat = dY * gamma52    dX = s * (dXhat - dXhat.mean(axis=-1, keepdims=True)53                - Xhat * (dXhat * Xhat).mean(axis=-1, keepdims=True))54    55    dgamma = (dY * Xhat).sum(axis=0)56    dbeta = dY.sum(axis=0)57    return dX, dgamma, dbeta

Code and math connection

Equation	Role	Lines	Code
Gradient on the normalised row	Computes it	ops.py:51	`dXhat = dY * gamma`
Gradient on the input	Computes it	ops.py:52	`dX = s * (dXhat - dXhat.mean(axis=-1, keepdims=True)`
Gradient on the input	Computes it	ops.py:53	`- Xhat * (dXhat * Xhat).mean(axis=-1, keepdims=True))`
Gradient on the input	Calls it	transformer.py:144	`dX1_branch, dg, db = layer_norm_backward(dXn2, c["ln2"], params[p + "ln2_g"])`
Gradient on the input	Calls it	transformer.py:149	`dX_branch, dg, db = layer_norm_backward(dXn, c["ln1"], params[p + "ln1_g"])`
Gradient on the input	Calls it	transformer.py:181	`dX, dg, db = layer_norm_backward(dXf, cache["lnf"], params["lnf_g"])`
Gradient on the input	Tests it	tests/test_ops.py:75	`dx, dgamma, dbeta = layer_norm_backward(np.array([[1., 0., 0., 1.]]), cache, gamma)`
Gradient on the scale, summed over rows	Computes it	ops.py:55	`dgamma = (dY * Xhat).sum(axis=0)`
Gradient on the shift, summed over rows	Computes it	ops.py:56	`dbeta = dY.sum(axis=0)`

Where is it used?

tests/test_gradcheck.py also reaches these gradients through the whole model: test_every_tensor_passes compares every parameter except the key biases with central differences, which includes gamma and beta.

Source

ops.py, lines 49-57, layer_norm_backward

PDF reference

Section 4.3, Eq. 32, 34, p. 7

Used by

Called by: transformer.block_backward, transformer.loss_and_grads

Tested by: tests.test_ops.test_layer_norm_worked_example

Calls it: transformer.py, line 144 (block_backward), for Gradient on the input
Calls it: transformer.py, line 149 (block_backward), for Gradient on the input
Calls it: transformer.py, line 181 (loss_and_grads), for Gradient on the input
Tests it: tests/test_ops.py, line 75 (test_layer_norm_worked_example), for Gradient on the input

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Gradient on the normalised row	Computes it	`ops.py`	`layer_norm_backward`	51	GitHub: ops.py, line 51
Gradient on the input	Computes it	`ops.py`	`layer_norm_backward`	52	GitHub: ops.py, line 52
Gradient on the input	Computes it	`ops.py`	`layer_norm_backward`	53	GitHub: ops.py, line 53
Gradient on the scale, summed over rows	Computes it	`ops.py`	`layer_norm_backward`	55	GitHub: ops.py, line 55
Gradient on the shift, summed over rows	Computes it	`ops.py`	`layer_norm_backward`	56	GitHub: ops.py, line 56

How this reaches a parameter update. The caller adds dgamma and dbeta to grads, as at transformer.py:183 and transformer.py:145. batch_loss_and_grads averages them over the batch (train.py:61). clip_grad_norm may scale them (train.py:62), and opt.step uses them (train.py:66). Both are one-dimensional, so they get no weight decay. dX is not stored. It goes on to the earlier layers, through the residual connections in block_backward (transformer.py:152).

Deriving the backward formula

What is it?

Deriving the backward formulaDerived from implementation

Basic operations

The steps that lead from the forward definition to the coded dX. The code holds only the result, in ops.layer_norm_backward. These steps are worked out here from ops.layer_norm_forward.

Why does the model need it?

The model does not run these steps. They show why the formula has its three terms, and that the result is exact, not an approximation, even with epsilon inside s.

Mathematics

Start from the loss and go back one step. Y depends on the normalised row through gamma.

Step 1: from the output to the normalised row #

\frac{\partial L}{\partial \hat{X}_{ti}} = \frac{\partial L}{\partial Y_{ti}}\, \frac{\partial Y_{ti}}{\partial \hat{X}_{ti}} = dY_{ti}\, \gamma_i = d\hat{X}_{ti}

Next, work out how the normalised row depends on the input. It depends on mu and s, and s depends on the variance. Every entry of the row changes the mean by the same amount.

Step 2: the derivative of the mean #

\frac{\partial \mu_t}{\partial X_{tj}} = \frac{1}{d}

The sum of X - mu over a row is 0, so the term from the mean drops out.

Step 3: the derivative of the variance, where the mu term cancels #

\frac{\partial \sigma_t^2}{\partial X_{tj}} = \frac{2}{d} \sum_{i} \left(X_{ti} - \mu_t\right)\left(\delta_{ij} - \frac{1}{d}\right) = \frac{2}{d}\left(X_{tj} - \mu_t\right)

Epsilon is a constant, so it stays inside the power. That power of var + epsilon is s cubed, exactly.

Step 4: the derivative of s, with epsilon absorbed into s #

\frac{\partial s_t}{\partial X_{tj}} = -\frac{1}{2}\left(\sigma_t^2 + \epsilon\right)^{-3/2} \frac{2}{d}\left(X_{tj} - \mu_t\right) = -\frac{1}{d}\, s_t^3 \left(X_{tj} - \mu_t\right)

Apply the product rule to (X - mu) * s. In the second term, (X - mu) * s appears twice, and each is a Xhat.

Step 5: the product rule on the normalised row #

\frac{\partial \hat{X}_{ti}}{\partial X_{tj}} = s_t\left(\delta_{ij} - \frac{1}{d}\right) + \left(X_{ti} - \mu_t\right)\frac{\partial s_t}{\partial X_{tj}} = s_t\left(\delta_{ij} - \frac{1}{d} - \frac{1}{d}\,\hat{X}_{ti}\hat{X}_{tj}\right)

Multiply by dXhat and sum over the features of the row. Each 1/d and its sum is one of the two means in the code, which gives the formula above and its code lines.

Step 6: the chain rule over all features of the row #

dX_{tj} = \sum_{i} d\hat{X}_{ti}\, \frac{\partial \hat{X}_{ti}}{\partial X_{tj}} = s_t \left( d\hat{X}_{tj} - \frac{1}{d}\sum_{i} d\hat{X}_{ti} - \hat{X}_{tj}\, \frac{1}{d}\sum_{i} d\hat{X}_{ti}\hat{X}_{ti} \right)

gamma and beta enter directly. Each is shared by every row, which is why their gradients are sums over rows.

The parameters, each of which is shared by every row #

\frac{\partial Y_{ti}}{\partial \gamma_i} = \hat{X}_{ti},\quad \frac{\partial Y_{ti}}{\partial \beta_i} = 1

What does each symbol mean?

Symbol	Meaning	In code	Shape
$L$	The scalar loss	`loss`	scalar
$\delta_{ij}$	1 when i equals j, otherwise 0. It records that X_ti depends only on itself	not named in the code	scalar
$i,\, j$	Feature indexes within one row. The rows do not affect each other	not named in the code	scalar
$d\hat{X}$	Gradient on the normalised row, from step 1	`dXhat`	(T, d)
$s$	Inverse standard deviation, with epsilon inside it	`s`	(T, 1)

Tiny numerical example

The derived matrix against the codeIllustrative ExampleThe same hand-picked row as the backward example. The matrix is step 5, multiplied through in step 6. The last result is stored from running the real layer_norm_backward.

Step 5, entry (i, j), shape (4, 4)
0.2482	-0.21	-0.1336	0.0954
-0.21	0.3627	-0.1336	-0.0191
-0.1336	-0.1336	0.4009	-0.1336
0.0954	-0.0191	-0.1336	0.0573

Central differences of the forward, shape (4, 4)
0.2482	-0.21	-0.1336	0.0954
-0.21	0.3627	-0.1336	-0.0191
-0.1336	-0.1336	0.4009	-0.1336
0.0954	-0.0191	-0.1336	0.0573

Step 6, as one row, shape (1, 4)
0.4391	-0.2482	-0.4009	0.21

Real layer_norm_backward, shape (1, 4)
0.4391	-0.2482	-0.4009	0.21

Checking the whole backward numericallyIllustrative ExampleA test-time idea, not part of the pipeline: gradcheck.numerical_grad is only called by tests. It nudges each number up and down by h = 1e-5 and measures how (Y * dY).sum(), whose gradient on Y is exactly dY, changes. Tests compare the central differences here with that real function.

Gradient	From the backward formula	From central differences	Largest gap
`dX`	`[0.4391, -0.2482, -0.4009, 0.2100]`	`[0.4391, -0.2482, -0.4009, 0.2100]`	`4.7e-11`
`dgamma`	`[-1.0690, 0.0000, 0.0000, 1.6036]`	`[-1.0690, 0.0000, 0.0000, 1.6036]`	`1.4e-11`
`dbeta`	`[1.0000, 0.0000, 0.0000, 1.0000]`	`[1.0000, 0.0000, 0.0000, 1.0000]`	`3.8e-11`

Tensor shapes

Each row has its own matrix of derivatives. The code never builds itops.py:52

dXhat row(d,)(32,)
dX/dXhat matrix(d, d)(32, 32)
dX row(d,)(32,)

d = emb_dim = 32.

Actual implementation

The code has no function for these steps. Its result is ops.layer_norm_backward, shown in the previous section, and the forward lines it comes from are in ops.layer_norm_forward.

Code and math connection

Equation	Role	Lines	Code
Step 1: from the output to the normalised row	Computes it	ops.py:51	`dXhat = dY * gamma`
Step 2: the derivative of the mean	Computes it	ops.py:43	`mu = X.mean(axis=-1, keepdims=True)`
Step 3: the derivative of the variance, where the mu term cancels	Computes it	ops.py:44	`var = ((X - mu) ** 2).mean(axis=-1, keepdims=True)`
Step 4: the derivative of s, with epsilon absorbed into s	Computes it	ops.py:45	`s = 1.0 / np.sqrt(var + epsilon)`
Step 5: the product rule on the normalised row	Computes it	ops.py:46	`Xhat = (X - mu) * s`
Step 6: the chain rule over all features of the row	Computes it	ops.py:52	`dX = s * (dXhat - dXhat.mean(axis=-1, keepdims=True)`
Step 6: the chain rule over all features of the row	Computes it	ops.py:53	`- Xhat * (dXhat * Xhat).mean(axis=-1, keepdims=True))`
The parameters, each of which is shared by every row	Computes it	ops.py:55	`dgamma = (dY * Xhat).sum(axis=0)`
The parameters, each of which is shared by every row	Computes it	ops.py:56	`dbeta = dY.sum(axis=0)`

Where is it used?

Nowhere at run time. It explains a line that the pipeline uses.

Source

ops.py, lines 42-47, layer_norm_forward

ops.py, lines 49-57, layer_norm_backward

PDF reference

Not stated in the PDF

Used by

Called by: transformer.block_forward, transformer.forward, transformer.block_backward, transformer.loss_and_grads

Tested by: tests.test_ops.test_layer_norm_worked_example, tests.test_ops.test_layer_norm_worked_example

Status

Derived from implementation

Implemented in

Equation	Role	File	Function	Lines	GitHub
Step 1: from the output to the normalised row	Computes it	`ops.py`	`layer_norm_backward`	51	GitHub: ops.py, line 51
Step 2: the derivative of the mean	Computes it	`ops.py`	`layer_norm_forward`	43	GitHub: ops.py, line 43
Step 3: the derivative of the variance, where the mu term cancels	Computes it	`ops.py`	`layer_norm_forward`	44	GitHub: ops.py, line 44
Step 4: the derivative of s, with epsilon absorbed into s	Computes it	`ops.py`	`layer_norm_forward`	45	GitHub: ops.py, line 45
Step 5: the product rule on the normalised row	Computes it	`ops.py`	`layer_norm_forward`	46	GitHub: ops.py, line 46
Step 6: the chain rule over all features of the row	Computes it	`ops.py`	`layer_norm_backward`	52	GitHub: ops.py, line 52
Step 6: the chain rule over all features of the row	Computes it	`ops.py`	`layer_norm_backward`	53	GitHub: ops.py, line 53
The parameters, each of which is shared by every row	Computes it	`ops.py`	`layer_norm_backward`	55	GitHub: ops.py, line 55
The parameters, each of which is shared by every row	Computes it	`ops.py`	`layer_norm_backward`	56	GitHub: ops.py, line 56

Parameters

Each block has two LayerNorms, and the model has one more at the end. Each has its own pair of vectors, and init_params sets gamma to ones and beta to zeros. So at the start of training, the output of each LayerNorm is exactly Xhat.

Name	Shape	With config	Scalars	Role
`b{l}.ln1_g`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 scale, before attention.Starts as ones. No decay.
`b{l}.ln1_b`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 shift, before attention.Starts as zeros. No decay.
`b{l}.ln2_g`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 scale, before feed-forward.Starts as ones. No decay.
`b{l}.ln2_b`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 shift, before feed-forward.Starts as zeros. No decay.
`lnf_g`transformer.py:39	`(d,)`	`(32,)`	`d`32	Final LayerNorm scale.Starts as ones. No decay.
`lnf_b`transformer.py:39	`(d,)`	`(32,)`	`d`32	Final LayerNorm shift.Starts as zeros. No decay.

Where LayerNorm is applied

The blocks are pre-LN: each branch reads a normalised copy, and the residual adds its result to the value before normalisation. In block_forward, the first LayerNorm runs before attention:

transformer.pyblock_forwardGitHub

133    Xn, c_ln1 = layer_norm_forward(X, params[p + "ln1_g"], params[p + "ln1_b"], cfg.ln_eps) # make a normalized copy for attention to read

The second runs before the feed-forward layer:

transformer.pyblock_forwardGitHub

136    Xn2, c_ln2 = layer_norm_forward(X1, params[p + "ln2_g"], params[p + "ln2_b"], cfg.ln_eps) # normalized copy again, for the feed-forward

After the last block, forward applies the final LayerNorm, and the head reads its output:

transformer.pyforwardGitHub

165    Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)

Which epsilon is passed. All three calls pass cfg.ln_eps. train.main builds ModelConfig without that argument (train.py:48-51), so the value comes from the dataclass default:

transformer.pyModelConfigGitHub

15    ln_eps: float = 1e-5

config.py also has an ln_eps (config.py:8), but nothing reads it, so editing it changes nothing. Today it has the same value, 1e-5. The default on the function itself (ops.py:42) applies only where a caller leaves the argument out, as test_layer_norm_worked_example does.