Repository navigation
Backward of batched attention builds a contraction over every pair of sequences and heads (9.4M rows, 23 s of a 66 s GPT step) #217
Description
🌀 Claude: In the backward pass of batched, multi-head attention, ddx stores a contraction step that is a cross product over every pair of sequences and heads. The gradient is right, but the step has (sequences × heads × positions)² rows, where the scores have sequences × heads × positions²/2. In xarray-sql's GPT demo (benchmarks/ml/02_nanogpt.py, xqlsystems/xarray-sql#269), at nanoGPT's suggested CPU size with one layer (4 heads, 128 channels, 64 positions, batch 12), that step stores 9,437,184 rows and takes 23 s of a 65.6 s training step on DuckDB. The two steps that read it take another 24 s. The forward loss takes 2.9 s. Both of ddxdb's backends do it, so it's in ddx's core.
The cause, as far as I can tell. The attention output joins the weights a with v on a.seq = v.seq AND a.h = v.h AND a.j = v.pos, and groups by a.seq, a.i, a.h, v.c. To differentiate with respect to the weights, ddx contracts v with the cotangent of the output. That contraction joins the two on c alone, and groups by both sides' sequence, head and position:
… FROM v INNER JOIN cotangent_2 ON (v.c IS NOT DISTINCT FROM cotangent_2.c)
GROUP BY v.seq, v.h, v.pos, cotangent_2.seq, cotangent_2.pos, cotangent_2.hv.seq equals the cotangent's seq only through a (v.seq = a.seq, and a.seq is the output's seq), and the same holds for h. The contraction is built before a is joined, so it uses only the key that v and the cotangent share directly. A later step joins back on those keys with IS NOT DISTINCT FROM, so the result is right, but it's computed through the cross product. Taking the transitive closure of the join equalities, before choosing a contraction's join keys, would keep seq and h in it. The output would then have sequences × heads × positions² rows, at most twice the score table's.
Repro, with ddxdb 0.4.3, DuckDB 1.5.6 and DataFusion 55.0.0, all from PyPI. It counts each step table's rows on both backends, and checks the gradient against JAX:
"""ddx's backward of batched, multi-head causal attention on DuckDB: step table sizes."""
import sys
import duckdb, numpy as np, pyarrow as pa
import ddxdb.duckdb as dd
from ddxdb import ad
from datafusion import SessionContext
S, H, T, C = 3, 2, 8, 4 # sequences, heads, positions, channels per head
rows = []
orig = dd.DuckDBBackend.materialize
def materialize(self, name, plan):
orig(self, name, plan)
rows.append((name.split("_", 3)[-1], self.con.execute(f"SELECT count(*) FROM temp.main.{dd._ident(name)}").fetchone()[0]))
dd.DuckDBBackend.materialize = materialize
orig_df = ad.DataFusionBackend.materialize
def materialize_df(self, name, plan):
orig_df(self, name, plan)
rows.append(("datafusion " + name.split("_", 3)[-1], self.ctx.sql(f'SELECT count(*) AS n FROM "{name}"').to_pydict()["n"][0]))
ad.DataFusionBackend.materialize = materialize_df
con = duckdb.connect()
idx = np.indices((S, T, H, C)).reshape(4, -1)
for name, phase in (("q", 0.0), ("k", 1.0), ("v", 2.0)):
con.register(name, pa.table({"seq": idx[0], "pos": idx[1], "h": idx[2], "c": idx[3],
"val": np.sin(np.arange(idx.shape[1]) * 0.37 + phase)}))
SCORES = {
"all keys in ON": "FROM q JOIN k ON q.seq = k.seq AND q.h = k.h AND q.c = k.c WHERE q.pos >= k.pos",
}
MASK = sys.argv[1] if len(sys.argv) > 1 else "all keys in ON"
LOSS = f"""WITH scores AS (
SELECT q.seq, q.h, q.pos AS i, k.pos AS j, SUM(q.val * k.val) / 2.0 AS s
{SCORES[MASK]}
GROUP BY q.seq, q.h, q.pos, k.pos),
weights AS (
SELECT seq, h, i, j, exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p FROM scores),
norm AS (SELECT seq, h, i, SUM(p) AS z FROM weights GROUP BY seq, h, i),
heads AS (
SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val) AS val
FROM weights a
JOIN v ON a.seq = v.seq AND a.h = v.h AND a.j = v.pos
JOIN norm n ON a.seq = n.seq AND a.h = n.h AND a.i = n.i
GROUP BY a.seq, a.i, a.h, v.c),
loss AS (SELECT SUM(val * val) / 2 AS l FROM heads) """
got = dd.sql(con, LOSS + "SELECT seq, pos, h, c, grad(loss, val) AS g FROM q ORDER BY seq, pos, h, c").to_arrow_table()
ctx = SessionContext()
for name in "qkv":
ctx.register_record_batches(name, [con.sql(f"SELECT * FROM {name}").to_arrow_table().to_batches()])
got_df = ad.sql(ctx, LOSS + "SELECT seq, pos, h, c, grad(loss, val) AS g FROM q ORDER BY seq, pos, h, c").to_arrow_table()
import jax, jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
def loss(q, k, v): # (seq, pos, h, c)
s = jnp.einsum("sihc,sjhc->shij", q, k) / 2.0
s = jnp.where(jnp.tril(jnp.ones((T, T), bool)), s, -jnp.inf)
p = jax.nn.softmax(s, axis=-1)
return jnp.sum(jnp.einsum("shij,sjhc->sihc", p, v) ** 2) / 2
q, k, v = (np.sin(np.arange(S * T * H * C) * 0.37 + ph).reshape(S, T, H, C) for ph in (0.0, 1.0, 2.0))
want = np.asarray(jax.grad(loss)(q, k, v)).ravel()
print("DataFusion gradient matches JAX:", np.allclose(np.array(got_df["g"]), want, rtol=1e-12, atol=1e-14))
print("DuckDB gradient matches JAX:", np.allclose(np.array(got["g"]), want, rtol=1e-12, atol=1e-14))
print(f"S*H*T*(T+1)/2 = {S*H*T*(T+1)//2} score rows; (S*H*T)^2 = {(S*H*T)**2}; S*H*T*T = {S*H*T*T}")
for name, n in rows:
print(f" {name}: {n}")DataFusion gradient matches JAX: True
DuckDB gradient matches JAX: True
S*H*T*(T+1)/2 = 216 score rows; (S*H*T)^2 = 2304; S*H*T*T = 384
97c82f0_0_cte_1: 0
bd42312_1_saved_0: 216
bd42312_1_saved_1: 48
bd42312_1_saved_2: 192
bd42312_1_saved_3: 1
bd42312_1_value: 1
bd42312_1_cotangent_3: 1
bd42312_1_cotangent_2: 192
bd42312_1_contraction_0: 2304
bd42312_1_window_1: 216
bd42312_1_cotangent_1: 48
bd42312_1_window_2: 216
bd42312_1_cotangent_0: 216
bd42312_1_grad_0_q: 192
datafusion 1227462c_2_cte_1: 0
datafusion 128b83ab_3_saved_0: 216
datafusion 128b83ab_3_saved_1: 48
datafusion 128b83ab_3_saved_2: 192
datafusion 128b83ab_3_saved_3: 1
datafusion 128b83ab_3_value: 1
datafusion 128b83ab_3_cotangent_3: 1
datafusion 128b83ab_3_cotangent_2: 192
datafusion 128b83ab_3_contraction_0: 2304
datafusion 128b83ab_3_window_1: 216
datafusion 128b83ab_3_cotangent_1: 48
datafusion 128b83ab_3_window_2: 216
datafusion 128b83ab_3_cotangent_0: 216
datafusion 128b83ab_3_grad_0_q: 192
contraction_0 has (3 × 2 × 8)² = 2,304 rows on both backends, against 216 score rows. At the GPT's size above, that's 3,072² = 9,437,184 rows, from about 300 million joined rows, where sequences × heads × positions² is 196,608.
Ask: keep every key that the joins make equal, transitively, in a contraction's join and GROUP BY. That way no step is larger than the forward pass's own intermediates.
Activity
alxmrs commentedon Oct 28, 2026
🌀 Claude: It's also the training step's peak memory. At the GPT size above, on DuckDB, the process peaks at 2.0 GB through
saved_11, then at 9.2 GB atcontraction_8, the cross product. The loss alone peaks at 1.2 GB. Measured withresource.getrusageafter each step, on the same machine and versions.