Memory and Precision

The Memory Anatomy of a Step

Params 4P + grads 4P + Adam states 8P = 16P bytes before activations. Activations scale with batch×depth×width and dominate; peak lands early in backward.

Measuring: Count vs. Plan

torch.cuda.memory._record_memory_history(max_entries=100000)
for _ in range(3):
    opt.zero_grad(set_to_none=True)
    net(X).sum().backward()
    opt.step()
snapshot = torch.cuda.memory._snapshot()
torch.cuda.memory._record_memory_history(enabled=None)

# Reconstruct cumulative allocated bytes from the allocation trace.
cum, series = 0, []
for ev in snapshot['device_traces'][0]:
    if ev['action'] == 'alloc':
        cum += ev['size']
    elif ev['action'] == 'free_completed':
        cum -= ev['size']
    series.append(cum / 1e6)
d2l.plot(list(range(len(series))), [series],
         'allocation event', 'allocated (MB)')

The sawtooth is the anatomy: grow in forward, shrink in backward.

def loss_fn(params, X):
    h = X
    for W in params:
        h = jax.nn.relu(h @ W)
    return h.sum()

key = jax.random.PRNGKey(0)
params = [jax.random.normal(k, (2048, 2048)) * 0.02
          for k in jax.random.split(key, 6)]
X = jax.random.normal(key, (256, 2048))

compiled = jax.jit(jax.grad(loss_fn)).lower(params, X).compile()
a = compiled.memory_analysis()
print(f'compiler-planned temp memory: {a.temp_size_in_bytes / 1e6:.0f} MB')
print(f'argument + output: '
      f'{(a.argument_size_in_bytes + a.output_size_in_bytes) / 1e6:.0f} MB')

compiler-planned temp memory: 15 MB
argument + output: 203 MB

PyTorch counts at runtime; XLA plans at compile time.

Mixed Precision, Fairly Measured

def net_fn(params, X, dtype):
    h = X.reshape(X.shape[0], -1).astype(dtype)
    for W in params[:-1]:
        h = jax.nn.relu(h @ W.astype(dtype))
    return h @ params[-1].astype(dtype)

key = jax.random.PRNGKey(0)
sizes = [(64 * 64, 8192), (8192, 8192), (8192, 10)]
params = [jax.random.normal(k, s) * 0.02
          for k, s in zip(jax.random.split(key, 3), sizes)]
X = jax.random.normal(key, (2048, 64 * 64))

def loss_fn(params, dtype):
    return (net_fn(params, X, dtype) ** 2).sum()

grad_fn = jax.jit(jax.grad(loss_fn), static_argnums=1)
with jax.default_matmul_precision('highest'):  # Unfair baseline: tf32 off
    print(d2l.Benchmark(lambda: grad_fn(params, jnp.float32),
                        desc='fp32 (tf32 off)'))
# The default on this card is tf32-class compute -- the fair baseline
print(d2l.Benchmark(lambda: grad_fn(params, jnp.float32), desc='tf32'))
print(d2l.Benchmark(lambda: grad_fn(params, jnp.bfloat16), desc='bf16'))

fp32 (tf32 off): 26.41 ms/call
tf32: 15.99 ms/call
bf16: 8.70 ms/call

fp32→tf32 is a config (opposite defaults: PyTorch off, JAX on); tf32→bf16 adds ~1.5–2× more. No GradScaler for bf16 — it shares fp32’s exponent range. Only fp16 underflows and needs loss scaling.

Activation Checkpointing

Store a few activations; recompute the rest in backward.

def block(W1, W2, x):
    return x + jax.nn.gelu(x @ W1) @ W2

key = jax.random.PRNGKey(1)
ks = jax.random.split(key, 32)             # one key per weight matrix
Ws = [(jax.random.normal(ks[2 * i], (1024, 1024)) * 0.02,
       jax.random.normal(ks[2 * i + 1], (1024, 1024)) * 0.02)
      for i in range(16)]
X = jax.random.normal(key, (16384, 1024))

# jax.checkpoint (a.k.a. remat) recomputes the block in the backward pass.
ckpt_block = jax.checkpoint(block)
def forward(Ws, X, blk):
    h = X
    for W1, W2 in Ws:
        h = blk(W1, W2, h)
    return (h ** 2).sum()

for name, blk in [('store all', block), ('checkpointed', ckpt_block)]:
    compiled = jax.jit(jax.grad(forward), static_argnums=2).lower(
        Ws, X, blk).compile()
    mb = compiled.memory_analysis().temp_size_in_bytes / 1e6
    t = d2l.Benchmark(lambda: compiled(Ws, X)).time
    print(f'{name}: compiler temp {mb:.0f} MB, fwd+bwd {1000 * t:.1f} ms')

store all: compiler temp 3230 MB, fwd+bwd 50.5 ms
checkpointed: compiler temp 1216 MB, fwd+bwd 58.7 ms

Peak memory cut roughly in half or better (~2 GB of activations released) for ~15–20% more time. Same trade as Mamba’s recompute-in-kernel (§12), one level up. It buys memory, not speed — use it when memory binds.

Gradient Accumulation

B_{\textrm{global}} = B_{\textrm{micro}} \times k: split the batch, sum the gradients, step once.

W = jax.random.normal(jax.random.PRNGKey(0), (1024, 1))
X = jax.random.normal(jax.random.PRNGKey(1), (128, 1024))
y = jax.random.normal(jax.random.PRNGKey(2), (128, 1))

def mse(W, X, y):
    return ((X @ W - y) ** 2).mean()

full = jax.grad(mse)(W, X, y)
acc = sum(jax.grad(mse)(W, X[i * 32:(i + 1) * 32], y[i * 32:(i + 1) * 32])
          for i in range(4)) / 4
print(f'max gradient difference: {jnp.abs(full - acc).max():.2e}')

max gradient difference: 3.10e-06

Gradients match a full-batch step to noise (scale by 1/k for mean-losses). Big effective batch at micro-batch memory — ≈ same wall-clock, never faster.

The Ladder So Far

slow (bandwidth/overhead) → compile (usually free)
slow (compute) or tight → bf16 (~1.5–2×, half the bytes)
doesn’t fit → checkpoint (large memory cut, modest time)
want a bigger batch → accumulate (≈ same wall-clock)
still too slow / too big → more devices (next)