nn.LayerNorm two-pass variance produces ~27% more HLO reduce ops under vmap + grad

[CatchemAL]

eqx.nn.LayerNorm uses two-pass variance (jnp.mean + jnp.var). Under vmap + jax.grad, this produces ~27% more HLO reduce ops and runs 1.45x slower than an identical model using jax.nn.standardize (one-pass variance: E[x²] - E[x]²).

Benchmark

8-block residual MLP, DIM=64, BATCH=256, SEQ=32, double vmap (batch × seq):

Variant	Time	HLO lines	Reduces	Broadcasts
(a) eqx.nn.LayerNorm	49.7 ms	1385	132	342
(b) LayerNormNew (drop-in)	34.2 ms	1011	104	238

Full script here to show timings.

Algorithm

eqx.nn.LayerNorm.__call__ (in _normalisation.py) computes:

mean = jnp.mean(x, keepdims=True)
variance = jnp.var(x, keepdims=True)        # internally: mean((x - mean)²)
variance = jnp.maximum(0.0, variance)
inv = jax.lax.rsqrt(variance + self.eps)
out = (x - mean) * inv

jnp.var is two-pass, jax.nn.standardize uses one-pass variance:

mean = jnp.mean(x, axis, keepdims=True)
variance = jnp.mean(jnp.square(x), axis, keepdims=True) - jnp.square(mean)

Suggested fix

In _normalisation.py, replace the two-pass variance:

mean = jnp.mean(x, keepdims=True)
variance = jnp.var(x, keepdims=True)
variance = jnp.maximum(0.0, variance)
inv = jax.lax.rsqrt(variance + self.eps)
out = (x - mean) * inv

with jax.nn.standardize:

out = jax.nn.standardize(x, axis=range(len(x.shape)), epsilon=self.eps)

This is a one-line change. jax.nn.standardize already handles the clipping (jnp.clip(variance, 0)), so jnp.maximum is no longer needed.

Confirmation of Correctness

Here is an suite of 50 unit tests confirming identical outputs - just faster.

Versions

• equinox 0.13.4
• jax 0.9.0
• CPU (no GPU)

I would be happy to prepare a PR if useful.

[patrick-kidger]

Thank you for catching this one too! Yup, very happy to take a PR on this, thank you :)

[CatchemAL]

Thanks. Will be with you in the next couple of days.

Hey, just a small heads-up: measuring with the CPU backend of JAX is sort of flaky, as it's much less optimized than the other backends. Also comparing pre-optimized StableHLO is often difficult, as XLA will rewrite/fuse away a lot of stuff. I wrote a small benchmark script that seems to indicate that (at least on my machine ;)) this new version is in fact not faster.

These are the timings I measured on JAX nightly with a 5060TI:

Shape           | Mode       | Two-Pass (ms)   | One-Pass (ms)   | Speedup
---------------------------------------------------------------------------
(64, 64)        | Fwd        | 0.0625          | 0.0544          | 1.15x
(64, 64)        | Fwd+Bwd    | 0.0564          | 0.0545          | 1.04x
(128, 128)      | Fwd        | 0.0445          | 0.0537          | 0.83x
(128, 128)      | Fwd+Bwd    | 0.0524          | 0.0552          | 0.95x
(256, 256)      | Fwd        | 0.0445          | 0.0547          | 0.81x
(256, 256)      | Fwd+Bwd    | 0.0540          | 0.0645          | 0.84x
(512, 512)      | Fwd        | 0.0538          | 0.0520          | 1.03x
(512, 512)      | Fwd+Bwd    | 0.0541          | 0.0654          | 0.83x
(1024, 1024)    | Fwd        | 0.0518          | 0.0559          | 0.93x
(1024, 1024)    | Fwd+Bwd    | 0.0537          | 0.0620          | 0.86x
(2048, 2048)    | Fwd        | 0.0848          | 0.0948          | 0.89x
(2048, 2048)    | Fwd+Bwd    | 0.0900          | 0.1021          | 0.88x
(4096, 4096)    | Fwd        | 0.4177          | 0.5768          | 0.72x
(4096, 4096)    | Fwd+Bwd    | 0.4179          | 0.7337          | 0.57x
(8192, 4096)    | Fwd        | 0.7846          | 1.1027          | 0.71x
(8192, 4096)    | Fwd+Bwd    | 0.7856          | 1.4121          | 0.56x
(8192, 8192)    | Fwd        | 1.5117          | 2.1444          | 0.70x
(8192, 8192)    | Fwd+Bwd    | 1.5152          | 2.7879          | 0.54x
---------------------------------------------------------------------------
Average Fwd Speedup:     0.87x
Average Fwd+Bwd Speedup: 0.78x

and running your script:

batch=256, seq=32, dim=64, blocks=8
  (a) eqx.nn.LayerNorm            1.0 ms
  (b) LayerNormNew                1.1 ms

Code to reproduce

import jax
import jax.numpy as jnp
import time

def two_pass_layernorm(x, eps=1e-5):
    mean = jnp.mean(x, keepdims=True)
    variance = jnp.var(x, keepdims=True)
    variance = jnp.maximum(0.0, variance)
    inv = jax.lax.rsqrt(variance + eps)
    return (x - mean) * inv

def one_pass_layernorm(x, eps=1e-5):
    return jax.nn.standardize(x, axis=tuple(range(x.ndim)), epsilon=eps)

def loss_two(x): return jnp.sum(two_pass_layernorm(x))
def loss_one(x): return jnp.sum(one_pass_layernorm(x))

fwd_two = jax.jit(jax.vmap(two_pass_layernorm))
fwd_one = jax.jit(jax.vmap(one_pass_layernorm))
bwd_two = jax.jit(jax.vmap(jax.grad(loss_two)))
bwd_one = jax.jit(jax.vmap(jax.grad(loss_one)))

def bench(fn, x, n=1000):
    fn(x).block_until_ready()
    start = time.perf_counter()
    for _ in range(n):
        fn(x).block_until_ready()
    return (time.perf_counter() - start) / n * 1000

shapes = [
    (64, 64), (128, 128), (256, 256),
    (512, 512), (1024, 1024), (2048, 2048),
    (4096, 4096), (8192, 4096), (8192, 8192)
]
key = jax.random.PRNGKey(0)

print(f"{'Shape':<15} | {'Mode':<10} | {'Two-Pass (ms)':<15} | {'One-Pass (ms)':<15} | {'Speedup'}")
print("-" * 75)

fwd_speedups = []
bwd_speedups = []

for shape in shapes:
    x = jax.random.normal(key, shape)
    
    t_fwd_two = bench(fwd_two, x)
    t_fwd_one = bench(fwd_one, x)
    speedup_fwd = t_fwd_two / t_fwd_one
    fwd_speedups.append(speedup_fwd)
    
    print(f"{str(shape):<15} | {'Fwd':<10} | {t_fwd_two:<15.4f} | {t_fwd_one:<15.4f} | {speedup_fwd:.2f}x")
    
    t_bwd_two = bench(bwd_two, x)
    t_bwd_one = bench(bwd_one, x)
    speedup_bwd = t_bwd_two / t_bwd_one
    bwd_speedups.append(speedup_bwd)
    
    print(f"{str(shape):<15} | {'Fwd+Bwd':<10} | {t_bwd_two:<15.4f} | {t_bwd_one:<15.4f} | {speedup_bwd:.2f}x")

print("-" * 75)
print(f"Average Fwd Speedup:     {sum(fwd_speedups) / len(fwd_speedups):.2f}x")
print(f"Average Fwd+Bwd Speedup: {sum(bwd_speedups) / len(bwd_speedups):.2f}x")

[CatchemAL]

Wow nice find @ZagButNoZig and quite a surprising result! Ok, I'll do some GPU investigations on both this and #1195 and see what I find. It'll take me a couple of days. But indeed, if this is harmful to GPU then we should simply close this (and maybe #1195) as "Won't do". I'll revert shortly.

But what a pity if CPU performance runs counter to GPU.