From 5d2b0d5fba398321fe69e2f312c25ff0519c8953 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 15:05:54 -0700 Subject: [PATCH 01/11] Add an ml extra for gradients in SQL, and three ML demos The `ml` extra installs ddxdb. With it, every Context registers ddx_stop_gradient, and Context.sql computes grad(), vjp() and jvp() calls in a select list. xarray_sql.ad.sql_all shares one backward pass across statements that differentiate the same loss. benchmarks/ml writes einfold's three target demos as SQL: causal attention in both of einfold's spellings, a 1-layer GPT trained with SGD through ddx, and a GraphCast-shaped GNN over WeatherBench 2 ERA5. Each checks its results against the same model in JAX and times the same SQL on DataFusion and DuckDB, reporting plan sizes alongside. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- README.md | 24 ++ benchmarks/ml/01_attention.py | 243 ++++++++++++ benchmarks/ml/02_nanogpt.py | 395 +++++++++++++++++++ benchmarks/ml/03_graphcast.py | 718 ++++++++++++++++++++++++++++++++++ benchmarks/ml/README.md | 169 ++++++++ benchmarks/ml/_bench.py | 149 +++++++ benchmarks/ml/_hosts.py | 111 ++++++ pyproject.toml | 8 +- tests/test_ad.py | 99 +++++ xarray_sql/ad.py | 131 +++++++ xarray_sql/sql.py | 28 +- 11 files changed, 2073 insertions(+), 2 deletions(-) create mode 100644 benchmarks/ml/01_attention.py create mode 100644 benchmarks/ml/02_nanogpt.py create mode 100644 benchmarks/ml/03_graphcast.py create mode 100644 benchmarks/ml/README.md create mode 100644 benchmarks/ml/_bench.py create mode 100644 benchmarks/ml/_hosts.py create mode 100644 tests/test_ad.py create mode 100644 xarray_sql/ad.py diff --git a/README.md b/README.md index 42fe05a7..12f3ecba 100644 --- a/README.md +++ b/README.md @@ -237,6 +237,30 @@ Some highlights: * Dense tensor operations are slower in SQL than Xarray. It seems like the vanilla HashJoin implementation in standard DB engines are not as performant as array ops (they don't take advantage of contiguous array layouts), though this gives us some idea for how it could [one day be possible](https://github.com/xqlsystems/einfold). * Sparse tensor operations are just as fast (or faster!) than Xarray today. Specifically, SQL is an excellent way to perform regridding and reprojecting between different spatial scales (see also: we include a `proj` udf via the `xarray-sql[geo]` extra). +## Gradients in SQL + +With the `ml` extra (`pip install "xarray-sql[ml]"`), a `Context` takes +gradients of queries with [ddx](https://github.com/xqlsystems/ddx). In a +select list, `grad(loss, column)` is, at each row, the derivative of the +number the CTE `loss` computes with respect to that row's `column`. A +gradient-descent step on a parameter table is then one statement: + +```python +ctx.sql(""" +WITH pred AS ( + SELECT x.sample, SUM(x.value * w.val) AS y + FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample), +loss AS ( + SELECT AVG(power(pred.y - t.value, 2)) AS l + FROM pred JOIN t ON pred.sample = t.sample) +SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w +""") +``` + +[`benchmarks/ml`](benchmarks/ml) trains a small GPT and runs attention and +a GraphCast-shaped weather model this way, checked against JAX and timed on +several engines. + ## Why does this work? Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py new file mode 100644 index 00000000..d6b3c932 --- /dev/null +++ b/benchmarks/ml/01_attention.py @@ -0,0 +1,243 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[ml,duckdb]", +# "jax", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Demo 1: attention in SQL, the query einfold should turn into FlashAttention. + +Attention, the core of a transformer, is ``softmax(Q·Kᵀ/√d)·V``. In the +XQL model, ``Q``, ``K`` and ``V`` are tables of ``(t, d, x)`` (token, +feature, value), and attention is two folds over joins: an einsum for the +scores, and a weighted average of ``V`` with weights ``exp(score)``. + +The SQL below is the naive plan, which stores every one of the N×N +scores. FlashAttention never does: it streams tiles of keys past tiles of +queries, with an online softmax. einfold's first target is to find that +plan from this SQL, in both of its spellings, with no attention-specific +rule (einfold's ``docs/demos.md``, Demo 1). This script is the baseline +that plan must match and beat: + +* **forward**: each spelling on each host, against a JAX reference; +* **gradient**: ddx differentiates the same SQL with respect to Q, K + and V, against ``jax.grad``; +* **masks are absent rows**: an infinity in V reaches only the queries + allowed to see it, as SQL's semantics require of any faster plan; +* **a control that is not attention**: the Nadaraya–Watson estimator, a + Gaussian-kernel weighted average used in regridding. The rules that + fuse attention must fuse it too. + +Run ``uv run benchmarks/ml/01_attention.py --help`` for the options. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import xarray as xr +from _bench import Report, assert_close, cli, hosts + +import xarray_sql as xql + +jax.config.update("jax_enable_x64", True) + +D = 64 # features per head; the scores are scaled by 1/√64 = 1/8. + +# The scores, softmax and weighted average, in einfold's window spelling. +# The causal mask is a join predicate: a masked (i, j) pair is a row that +# does not exist. The max only keeps exp() from overflowing and cancels +# out of the ratio, so it carries no gradient (ddx_stop_gradient). +ATTENTION_WINDOW = """ +scores AS ( + SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s + FROM q JOIN k ON q.d = k.d + WHERE q.t >= k.t + GROUP BY q.t, k.t), +weights AS ( + SELECT i, j, exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i))) AS w + FROM scores), +norm AS ( + SELECT i, SUM(w) AS z FROM weights GROUP BY i), +attention AS ( + SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x + FROM weights w + JOIN v ON w.j = v.t + JOIN norm n ON w.i = n.i + GROUP BY w.i, v.d) +""" + +# The same attention as ddx spells it: the row max as a GROUP BY, joined +# back to the scores. +ATTENTION_GROUP_BY = """ +scores AS ( + SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s + FROM q JOIN k ON q.d = k.d + WHERE q.t >= k.t + GROUP BY q.t, k.t), +row_max AS ( + SELECT i, ddx_stop_gradient(MAX(s)) AS m FROM scores GROUP BY i), +weights AS ( + SELECT s.i, s.j, exp(s.s - r.m) AS w + FROM scores s JOIN row_max r ON s.i = r.i), +norm AS ( + SELECT i, SUM(w) AS z FROM weights GROUP BY i), +attention AS ( + SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x + FROM weights w + JOIN v ON w.j = v.t + JOIN norm n ON w.i = n.i + GROUP BY w.i, v.d) +""" + +SPELLINGS = {"window": ATTENTION_WINDOW, "group by": ATTENTION_GROUP_BY} + +# A scalar of the output to differentiate: half its squared norm. +LOSS = "loss AS (SELECT SUM(x * x) / 2 AS l FROM attention)" + +# Nadaraya–Watson: at each target point, the average of the sources' y +# weighted by a Gaussian kernel of their distance (bandwidth h = 0.5). +# Structurally attention: a weighted average with weights exp(score). +NADARAYA_WATSON = """ +kernel AS ( + SELECT t.p, s.q, exp(-(power(t.x - s.x, 2) + power(t.y - s.y, 2)) / 0.25) AS w + FROM target t CROSS JOIN source s) +SELECT k.p, SUM(k.w * s.v) / SUM(k.w) AS v +FROM kernel k JOIN source s ON k.q = s.q +GROUP BY k.p +ORDER BY k.p +""" + + +def attention(q, k, v): + """The reference: causal softmax attention, as JAX writes it.""" + n = q.shape[0] + scores = q @ k.T / np.sqrt(D) + causal = jnp.tril(jnp.ones((n, n), dtype=bool)) + weights = jax.nn.softmax(jnp.where(causal, scores, -jnp.inf), axis=-1) + return weights @ v + + +def loss(q, k, v): + return jnp.sum(attention(q, k, v) ** 2) / 2 + + +def table(values: np.ndarray) -> xr.Dataset: + """A (t, d) matrix as the XQL table ``(t, d, x)``.""" + t, d = values.shape + return xr.Dataset( + {"x": (("t", "d"), values)}, + coords={"t": np.arange(t), "d": np.arange(d)}, + ) + + +def matrix(rows, name: str = "x") -> np.ndarray: + """The rows of a ``(t, d, …)`` result as a (t, d) matrix.""" + return xql.to_dataset(rows, dims=["t", "d"])[name].values + + +def main(): + parser = cli(__doc__.splitlines()[0], sizes=[64, 256], size_name="tokens") + args = parser.parse_args() + report = Report("attention", "tokens", args.reps) + rng = np.random.default_rng(0) + + for n in args.sizes: + q, k, v = (rng.normal(size=(n, D)) for _ in range(3)) + want = attention(q, k, v) + want_grads = jax.grad(loss, argnums=(0, 1, 2))(q, k, v) + + for host in hosts(args.hosts): + for name, data in zip("qkv", (q, k, v)): + host.load(name, table(data)) + + for spelling, ctes in SPELLINGS.items(): + forward = f"WITH {ctes} SELECT * FROM attention ORDER BY t, d" + rows = report.time( + f"forward ({spelling})", host, n, host.query, forward + ) + report.check(assert_close("forward", matrix(rows), want)) + report.note("operators", host.operators(forward)) + + step = f"gradient ({spelling})" + if not host.differentiates: + report.skip(step, host, n, "no ddx backend") + continue + # One statement per input, sharing one backward pass. + gradients = [ + f"WITH {ctes}, {LOSS} " + f"SELECT t, d, grad(loss, x) AS x FROM {name} ORDER BY t, d" + for name in "qkv" + ] + grads = report.time(step, host, n, host.query_all, gradients) + errors = [ + assert_close(f"d/d{name}", matrix(got), ref) + for name, got, ref in zip("qkv", grads, want_grads) + ] + report.check(max(errors)) + + check_masks_are_absent_rows(args.hosts) + check_nadaraya_watson(args.hosts, report) + report.finish(args.csv) + + +def check_masks_are_absent_rows(host_names): + """An infinity in V at token j reaches the queries at or after j, and + no others: a masked pair is skipped, not multiplied by a zero weight, + which would turn 0·∞ into NaN. (xarray-sql reads NaN as missing data, + a NULL that SUM skips, so an infinity is the value that tells.)""" + n, j = 16, 9 + rng = np.random.default_rng(1) + q, k, v = (rng.normal(size=(n, D)) for _ in range(3)) + v[j, 0] = np.inf + print(f"\n masks: +inf in V at token {j}, feature 0") + for host in hosts(host_names): + for name, data in zip("qkv", (q, k, v)): + host.load(name, table(data)) + for spelling, ctes in SPELLINGS.items(): + out = matrix( + host.query(f"WITH {ctes} SELECT * FROM attention ORDER BY t, d") + ) + assert np.isfinite(out[:j]).all(), (host.name, spelling) + assert np.isposinf(out[j:, 0]).all(), (host.name, spelling) + assert np.isfinite(out[:, 1:]).all(), (host.name, spelling) + print(f" {spelling:<9} {host.name:<11} ok") + + +def check_nadaraya_watson(host_names, report): + """The control: a kernel-weighted average that is not attention.""" + rng = np.random.default_rng(2) + targets, sources = rng.uniform(size=(64, 2)), rng.uniform(size=(256, 2)) + y = np.sin(4 * sources[:, 0]) + np.cos(3 * sources[:, 1]) + dist2 = ((targets[:, None, :] - sources[None, :, :]) ** 2).sum(-1) + w = np.exp(-dist2 / 0.25) + want = (w @ y) / w.sum(axis=1) + + target = xr.Dataset( + {"x": ("p", targets[:, 0]), "y": ("p", targets[:, 1])}, + coords={"p": np.arange(len(targets))}, + ) + source = xr.Dataset( + {"x": ("q", sources[:, 0]), "y": ("q", sources[:, 1]), "v": ("q", y)}, + coords={"q": np.arange(len(sources))}, + ) + for host in hosts(host_names): + host.load("target", target) + host.load("source", source) + rows = report.time( + f"nadaraya-watson ({len(sources)} sources)", + host, + len(targets), + host.query, + f"WITH {NADARAYA_WATSON}", + ) + got = xql.to_dataset(rows, dims=["p"])["v"].values + report.check(assert_close("nadaraya-watson", got, want)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py new file mode 100644 index 00000000..62df9ccb --- /dev/null +++ b/benchmarks/ml/02_nanogpt.py @@ -0,0 +1,395 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[ml,duckdb]", +# "jax", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Demo 2: a GPT in SQL, trained by gradient descent through ddx. + +nanoGPT is Andrej Karpathy's small, readable GPT-2. This is its model as +one SQL query: the cross-entropy loss of a batch of characters from +Shakespeare, computed by a transformer whose parameters are tables. + +================= ============================================ +Model piece In SQL +================= ============================================ +Embeddings a join of the batch with ``wte`` and ``wpe`` +Layer norm per-position ``AVG``s, joined back +Attention Demo 1's query, per head, causal +MLP two einsums around ReLU², ``power(greatest(x, 0), 2)`` +Output layer an einsum with ``wte`` again (tied, as in GPT-2) +Loss log-sum-exp per position, minus the target's logit +Gradients ddx: ``grad(loss, w)``, one statement per table +Optimizer step ``w - lr * grad(loss, w)``, the new table +================= ============================================ + +Each training step updates every parameter table with its own +statement; the statements share one backward pass. The same model in JAX, +from the same initial parameters and the same batches, is the reference: +its loss at every step and its gradients must match the SQL's. + +This is Part A of einfold's Demo 2 (``docs/demos.md``). Part B, the four +tricks of the NanoGPT speedrun's sparsity record as a small diff to this +SQL, builds on it. Simplifications from nanoGPT: no dropout, no biases, +and layer norms without a learned gain. + +Run ``uv run benchmarks/ml/02_nanogpt.py --help`` for the options. +""" + +from __future__ import annotations + +import pathlib +import urllib.request + +import jax +import jax.numpy as jnp +import numpy as np +import xarray as xr +from _bench import Report, assert_close, cli, hosts + +import xarray_sql as xql + +jax.config.update("jax_enable_x64", True) + +# A tiny character-level GPT. +LAYERS = 1 # transformer blocks +EMBED = 16 # e: channels of the residual stream +HEADS = 2 # h: attention heads +HEAD = EMBED // HEADS # c: channels per head +HIDDEN = 4 * EMBED # f: channels of the MLP +CONTEXT = 16 # pos: characters per sequence +LEARNING_RATE = 0.5 + +# Each parameter table's dimensions; its values are the column w. +PARAMS = { + "wte": ("tok", "e"), + "wpe": ("pos", "e"), + "wq": ("layer", "e", "h", "c"), + "wk": ("layer", "e", "h", "c"), + "wv": ("layer", "e", "h", "c"), + "wo": ("layer", "h", "c", "e"), + "w_fc": ("layer", "e", "f"), + "w_proj": ("layer", "f", "e"), +} + +# -------------------------------------------------------------------------- +# The model, in SQL. Activations are tables of (seq, pos, e, val): one row +# per sequence in the batch, position, and channel. +# -------------------------------------------------------------------------- + + +def layer_norm(x: str, out: str) -> str: + """Normalize each position of ``x`` to zero mean and unit variance.""" + return f""" +{out}_mean AS ( + SELECT seq, pos, AVG(val) AS mean FROM {x} GROUP BY seq, pos), +{out}_var AS ( + SELECT x.seq, x.pos, AVG(power(x.val - m.mean, 2)) AS var + FROM {x} x JOIN {out}_mean m ON x.seq = m.seq AND x.pos = m.pos + GROUP BY x.seq, x.pos), +{out} AS ( + SELECT x.seq, x.pos, x.e, (x.val - m.mean) / sqrt(v.var + 1e-5) AS val + FROM {x} x + JOIN {out}_mean m ON x.seq = m.seq AND x.pos = m.pos + JOIN {out}_var v ON x.seq = v.seq AND x.pos = v.pos),""" + + +def block(layer: int, x: str, out: str) -> str: + """One transformer block: ``out = x + attention(ln(x))``, then + ``out = out + mlp(ln(out))``.""" + l = f"l{layer}" + projections = "".join( + f""" +{l}_{name} AS ( + SELECT x.seq, x.pos, p.h, p.c, SUM(x.val * p.w) AS val + FROM {l}_ln1 x JOIN w{name} p ON x.e = p.e + WHERE p.layer = {layer} + GROUP BY x.seq, x.pos, p.h, p.c),""" + for name in ("q", "k", "v") + ) + return f"""{layer_norm(x, f"{l}_ln1")}{projections} +{l}_scores AS ( + SELECT q.seq, q.h, q.pos AS i, k.pos AS j, + SUM(q.val * k.val) / {np.sqrt(HEAD)} AS s + FROM {l}_q q JOIN {l}_k k ON q.seq = k.seq AND q.h = k.h AND q.c = k.c + WHERE q.pos >= k.pos + GROUP BY q.seq, q.h, q.pos, k.pos), +{l}_weights AS ( + SELECT seq, h, i, j, + exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY seq, h, i))) AS p + FROM {l}_scores), +{l}_norm AS ( + SELECT seq, h, i, SUM(p) AS z FROM {l}_weights GROUP BY seq, h, i), +{l}_heads AS ( + SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val) AS val + FROM {l}_weights a + JOIN {l}_v v ON a.seq = v.seq AND a.h = v.h AND a.j = v.pos + JOIN {l}_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), +{l}_attn AS ( + SELECT a.seq, a.pos, p.e, SUM(a.val * p.w) AS val + FROM {l}_heads a JOIN wo p ON a.h = p.h AND a.c = p.c + WHERE p.layer = {layer} + GROUP BY a.seq, a.pos, p.e), +{l}_mid AS ( + SELECT x.seq, x.pos, x.e, x.val + a.val AS val + FROM {x} x JOIN {l}_attn a ON x.seq = a.seq AND x.pos = a.pos AND x.e = a.e), +{layer_norm(f"{l}_mid", f"{l}_ln2")} +{l}_hidden AS ( + SELECT x.seq, x.pos, p.f, power(greatest(SUM(x.val * p.w), 0.0), 2) AS val + FROM {l}_ln2 x JOIN w_fc p ON x.e = p.e + WHERE p.layer = {layer} + GROUP BY x.seq, x.pos, p.f), +{l}_mlp AS ( + SELECT x.seq, x.pos, p.e, SUM(x.val * p.w) AS val + FROM {l}_hidden x JOIN w_proj p ON x.f = p.f + WHERE p.layer = {layer} + GROUP BY x.seq, x.pos, p.e), +{out} AS ( + SELECT x.seq, x.pos, x.e, x.val + m.val AS val + FROM {l}_mid x JOIN {l}_mlp m ON x.seq = m.seq AND x.pos = m.pos AND x.e = m.e),""" + + +def model() -> str: + """The CTEs of the whole model, ending in ``loss``: the mean + cross-entropy of predicting each next character of the batch.""" + blocks = "".join( + block(layer, f"x{layer}", f"x{layer + 1}") for layer in range(LAYERS) + ) + return f""" +x0 AS ( + SELECT b.seq, b.pos, te.e, te.w + pe.w AS val + FROM batch b + JOIN wte te ON b.tok = te.tok + JOIN wpe pe ON b.pos = pe.pos AND te.e = pe.e), +{blocks} +{layer_norm(f"x{LAYERS}", "final")} +logits AS ( + SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val + FROM final x JOIN wte te ON x.e = te.e + GROUP BY x.seq, x.pos, te.tok), +shifted AS ( + SELECT seq, pos, tok, val, + ddx_stop_gradient(MAX(val) OVER (PARTITION BY seq, pos)) AS m + FROM logits), +log_sum_exp AS ( + SELECT seq, pos, MAX(m) + ln(SUM(exp(val - m))) AS lse + FROM shifted GROUP BY seq, pos), +loss AS ( + SELECT AVG(s.lse - l.val) AS l + FROM log_sum_exp s + JOIN batch b ON s.seq = b.seq AND s.pos = b.pos + JOIN logits l ON l.seq = b.seq AND l.pos = b.pos AND l.tok = b.next)""" + + +def step_statements(lr: float) -> list[str]: + """One SGD step: each parameter table's next version, as a statement.""" + return [ + f"WITH {model()} SELECT {', '.join(dims)}, " + f"w - {lr} * grad(loss, w) AS w FROM {table} " + f"ORDER BY {', '.join(dims)}" + for table, dims in PARAMS.items() + ] + + +def loss_statement() -> str: + return f"WITH {model()} SELECT l FROM loss" + + +# -------------------------------------------------------------------------- +# The same model in JAX: the reference. +# -------------------------------------------------------------------------- + + +def jax_layer_norm(x): + mean = x.mean(-1, keepdims=True) + var = ((x - mean) ** 2).mean(-1, keepdims=True) + return (x - mean) / jnp.sqrt(var + 1e-5) + + +def jax_loss(p, tok, nxt): + x = p["wte"][tok] + p["wpe"][None] + causal = jnp.tril(jnp.ones((CONTEXT, CONTEXT), dtype=bool)) + for layer in range(LAYERS): + h = jax_layer_norm(x) + q, k, v = ( + jnp.einsum("bte,ehc->bhtc", h, p[w][layer]) + for w in ("wq", "wk", "wv") + ) + s = jnp.einsum("bhic,bhjc->bhij", q, k) / np.sqrt(HEAD) + a = jax.nn.softmax(jnp.where(causal, s, -jnp.inf), axis=-1) + heads = jnp.einsum("bhij,bhjc->bhic", a, v) + x = x + jnp.einsum("bhtc,hce->bte", heads, p["wo"][layer]) + h = jax_layer_norm(x) + hidden = jnp.maximum(h @ p["w_fc"][layer], 0.0) ** 2 + x = x + hidden @ p["w_proj"][layer] + logits = jax_layer_norm(x) @ p["wte"].T + target = jnp.take_along_axis(logits, nxt[..., None], -1)[..., 0] + return jnp.mean(jax.nn.logsumexp(logits, -1) - target) + + +# -------------------------------------------------------------------------- +# Data and parameters, as xarray Datasets. +# -------------------------------------------------------------------------- + +SHAKESPEARE = ( + "https://raw.githubusercontent.com/karpathy/char-rnn/master/" + "data/tinyshakespeare/input.txt" +) + + +def shakespeare() -> str: + """Tiny Shakespeare, cached in ~/.cache; a sonnet offline.""" + cache = ( + pathlib.Path.home() / ".cache" / "xarray-sql" / "tinyshakespeare.txt" + ) + if not cache.exists(): + cache.parent.mkdir(parents=True, exist_ok=True) + try: + with urllib.request.urlopen(SHAKESPEARE, timeout=10) as r: + cache.write_bytes(r.read()) + except OSError: + return SONNET_18 * 4 + return cache.read_text() + + +def batches(text: str, size: int, rng): + """Endless batches of (seq, pos) characters and the characters after + them, as one Dataset: the table ``batch(seq, pos, tok, next)``.""" + vocab = sorted(set(text)) + ids = np.array([vocab.index(c) for c in text]) + while True: + starts = rng.integers(0, len(ids) - CONTEXT - 1, size) + tok = np.stack([ids[s : s + CONTEXT] for s in starts]) + nxt = np.stack([ids[s + 1 : s + CONTEXT + 1] for s in starts]) + yield xr.Dataset( + {"tok": (("seq", "pos"), tok), "next": (("seq", "pos"), nxt)}, + coords={"seq": np.arange(size), "pos": np.arange(CONTEXT)}, + ) + + +def init_params(vocab: int, rng) -> dict[str, np.ndarray]: + """GPT-2's initialization: normal, standard deviation 0.02.""" + sizes = { + "tok": vocab, + "pos": CONTEXT, + "layer": LAYERS, + "e": EMBED, + "h": HEADS, + "c": HEAD, + "f": HIDDEN, + } + return { + name: rng.normal(0, 0.02, [sizes[d] for d in dims]) + for name, dims in PARAMS.items() + } + + +def as_table(name: str, values: np.ndarray) -> xr.Dataset: + dims = PARAMS[name] + return xr.Dataset( + {"w": (dims, values)}, + coords={d: np.arange(n) for d, n in zip(dims, values.shape)}, + ) + + +def as_array(name: str, rows) -> np.ndarray: + return xql.to_dataset(rows, dims=list(PARAMS[name]))["w"].values + + +def main(): + parser = cli(__doc__.splitlines()[0], sizes=[4], size_name="batch") + parser.add_argument("--steps", type=int, default=5, help="SGD steps") + args = parser.parse_args() + report = Report("nanogpt", "batch", args.reps) + text = shakespeare() + vocab = len(set(text)) + lines = len(loss_statement().splitlines()) + print( + f" {LAYERS}-layer GPT, vocab {vocab}: the loss is {lines} lines of SQL" + ) + + for size in args.sizes: + init = init_params(vocab, np.random.default_rng(0)) + data = list( + zip( + range(args.steps), batches(text, size, np.random.default_rng(1)) + ) + ) + + for host in hosts(args.hosts): + reference = dict(init) + for name, values in init.items(): + host.load(name, as_table(name, values)) + + for step, batch in data: + host.load("batch", batch) + tok, nxt = batch["tok"].values, batch["next"].values + + rows = report.time( + f"loss, step {step}", + host, + size, + host.query, + loss_statement(), + ) + want = jax_loss(reference, tok, nxt) + report.check( + assert_close("loss", rows["l"][0].as_py(), want, rtol=1e-9) + ) + report.note("loss", round(float(want), 6)) + report.note("operators", host.operators(loss_statement())) + + if not host.differentiates: + report.skip( + f"train, step {step}", host, size, "no ddx backend" + ) + continue + updated = report.time( + f"train, step {step}", + host, + size, + host.query_all, + step_statements(LEARNING_RATE), + ) + grads = jax.grad(jax_loss)(reference, tok, nxt) + errors = [] + for name, rows in zip(PARAMS, updated): + reference[name] = ( + reference[name] - LEARNING_RATE * grads[name] + ) + trained = as_array(name, rows) + errors.append( + assert_close( + f"{name}, step {step}", trained, reference[name] + ) + ) + host.load(name, as_table(name, trained)) + report.check(max(errors)) + + report.finish(args.csv) + + +SONNET_18 = """Shall I compare thee to a summer's day? +Thou art more lovely and more temperate: +Rough winds do shake the darling buds of May, +And summer's lease hath all too short a date; +Sometime too hot the eye of heaven shines, +And often is his gold complexion dimm'd; +And every fair from fair sometime declines, +By chance or nature's changing course untrimm'd; +But thy eternal summer shall not fade, +Nor lose possession of that fair thou ow'st; +Nor shall death brag thou wander'st in his shade, +When in eternal lines to time thou grow'st: +So long as men can breathe or eyes can see, +So long lives this, and this gives life to thee. +""" + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py new file mode 100644 index 00000000..3ca61684 --- /dev/null +++ b/benchmarks/ml/03_graphcast.py @@ -0,0 +1,718 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[ml,duckdb]", +# "jax", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Demo 3: a GraphCast-shaped weather model in SQL, over ERA5 in Zarr. + +GraphCast (Lam et al., 2023) forecasts the atmosphere 6 hours ahead from +the two previous states. It is a graph neural network in three stages: + +* an **encoder** maps the latitude–longitude grid onto a "multi-mesh", + a refined icosahedron that keeps the edges of every refinement level; +* a **processor** passes messages along the mesh's edges, round after + round; +* a **decoder** maps the mesh back onto the grid, as the 6-hour change. + +Here it is in miniature, as SQL over ERA5 read through xarray-sql from +WeatherBench 2's Zarr store (or a synthetic stand-in, offline). Every +graph is a table of edges ``(edge, sender, receiver)``, every state a +table of ``(id, f, val)``, and every stage the same relational steps: + +================= ============================================ +GNN step In SQL +================= ============================================ +Gather join the nodes' states to the edges, by sender + and by receiver +Concatenate ``UNION ALL``, offsetting each part's features +MLP einsums with ``params`` around swish, ``x / (1 + exp(-x))`` +Aggregate ``SUM`` of the edges' messages, ``GROUP BY`` receiver +================= ============================================ + +The geometry (the mesh, and which grid cells and mesh nodes are joined) +is computed once with numpy; the model, its inputs and its forecast are +SQL. The script checks the forecast, and ddx's gradient of the forecast +error, against the same model in JAX. Then it measures what einfold's +Demo 3 is about (``docs/demos.md``): how much less a regional forecast +needs than a global one. It reports the receptive field: how many mesh +nodes and grid cells each stage needs for the region. + +Simplifications from GraphCast: random weights, a 64×32 grid and six +variables, a mesh refined 2 times rather than 6, 2 processor rounds +rather than 16, no biases or layer norms, edge features of 3-D +displacement and length, decoder edges from each cell's 3 nearest mesh +nodes, and edge states that are embedded once rather than updated. + +Run ``uv run benchmarks/ml/03_graphcast.py --help`` for the options. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pandas as pd +import xarray as xr +from ddxdb import DdxError +from _bench import Report, assert_close, cli, hosts + +jax.config.update("jax_enable_x64", True) + +ERA5 = ( + "gs://weatherbench2/datasets/era5/" + "1959-2023_01_10-6h-64x32_equiangular_conservative.zarr" +) +ERA5_TABLES = { + ("time", "longitude", "latitude"): "surface", + ("time", "level", "longitude", "latitude"): "atmosphere", +} +# The forecast's variables, as (table, variable, pressure level). +VARIABLES = [ + ("surface", "2m_temperature", None), + ("surface", "mean_sea_level_pressure", None), + ("surface", "10m_u_component_of_wind", None), + ("surface", "10m_v_component_of_wind", None), + ("atmosphere", "geopotential", 500), + ("atmosphere", "temperature", 850), +] +# The two input states and the state to forecast, 6 hours apart. +TIMES = pd.to_datetime(["2020-01-01T00", "2020-01-01T06", "2020-01-01T12"]) + +HIDDEN = 16 # channels of every node and edge state +INPUTS = 2 * len(VARIABLES) + 3 # two states, cos(lat), sin(lon), cos(lon) +EDGE_FEATURES = 4 # receiver minus sender in 3-D, and the edge's length +# A region to forecast: the contiguous United States, roughly. +BOUNDS = {"latitude": (25, 50), "longitude": (235, 295)} +REGION = " AND ".join( + f"c.{coord} BETWEEN {low} AND {high}" + for coord, (low, high) in BOUNDS.items() +) + +# -------------------------------------------------------------------------- +# The model, in SQL. +# -------------------------------------------------------------------------- + + +def mlp(out: str, x: str, net: str) -> str: + """A two-layer MLP over the features of each ``id`` of ``x``.""" + return f""" +{out}_hidden AS ( + SELECT x.id, p.g AS f, SUM(x.val * p.w) AS val + FROM {x} x JOIN params p ON p.net = '{net}.0' AND x.f = p.f + GROUP BY x.id, p.g), +{out} AS ( + SELECT h.id, p.g AS f, SUM(h.val / (1 + exp(-h.val)) * p.w) AS val + FROM {out}_hidden h JOIN params p ON p.net = '{net}.1' AND h.f = p.f + GROUP BY h.id, p.g),""" + + +def message_passing(out, net, edges, edge_h, senders, receivers) -> str: + """One round of an interaction network on the graph ``edges``: each + edge's message is an MLP of the edge, its sender and its receiver; + each receiver adds an MLP of itself and its messages' sum to itself.""" + return f""" +{out}_edge_in AS ( + SELECT id, f, val FROM {edge_h} + UNION ALL + SELECT e.edge AS id, s.f + {HIDDEN} AS f, s.val + FROM {edges} e JOIN {senders} s ON e.sender = s.id + UNION ALL + SELECT e.edge AS id, r.f + {2 * HIDDEN} AS f, r.val + FROM {edges} e JOIN {receivers} r ON e.receiver = r.id), +{mlp(f"{out}_message", f"{out}_edge_in", f"{net}.edge")} +{out}_sum AS ( + SELECT e.receiver AS id, m.f, SUM(m.val) AS val + FROM {out}_message m JOIN {edges} e ON m.id = e.edge + GROUP BY e.receiver, m.f), +{out}_node_in AS ( + SELECT id, f, val FROM {receivers} + UNION ALL + SELECT id, f + {HIDDEN} AS f, val FROM {out}_sum), +{mlp(f"{out}_update", f"{out}_node_in", f"{net}.node")} +{out} AS ( + SELECT r.id, r.f, r.val + u.val AS val + FROM {receivers} r JOIN {out}_update u ON r.id = u.id AND r.f = u.f),""" + + +def era5_state(time, first_feature: int) -> list[str]: + """Each variable at ``time``, normalized, as (cell, feature, value) + rows. Only these variables, levels and times are read from ERA5.""" + parts = [] + for var, (table, name, level) in enumerate(VARIABLES): + on_level = f" AND x.level = {level}" if level else "" + parts.append( + f""" SELECT c.cell AS id, {first_feature + var} AS f, + (x."{name}" - n.mean) / n.std AS val + FROM era5.{table} x + JOIN cells c ON x.latitude = c.latitude AND x.longitude = c.longitude + JOIN scale n ON n.var = {var} + WHERE x.time = TIMESTAMP '{time}'{on_level}""" + ) + return parts + + +def position(table: str, key: str, first_feature: int) -> list[str]: + """Where each row of ``table`` is on the sphere, as 3 features.""" + features = [ + "cos(radians(latitude))", + "sin(radians(longitude))", + "cos(radians(longitude))", + ] + return [ + f" SELECT {key} AS id, {first_feature + i} AS f, {feature} AS val " + f"FROM {table}" + for i, feature in enumerate(features) + ] + + +def union(name: str, parts: list[str]) -> str: + return f"\n{name} AS (\n" + "\n UNION ALL\n".join(parts) + ")," + + +def edge_features(graph: str) -> str: + return f"(SELECT edge AS id, f, val FROM {graph}.features)" + + +def model(rounds: int) -> str: + """The CTEs of the whole model, ending in ``forecast``: the change of + each normalized variable over the next 6 hours, at each cell.""" + n = len(VARIABLES) + processor = "".join( + message_passing( + f"mesh{r + 1}", + f"processor.{r}", + "multimesh.edges", + "multimesh_h", + f"mesh{r}", + f"mesh{r}", + ) + for r in range(rounds) + ) + return f"""{ + union( + "grid_in", + era5_state(TIMES[0], 0) + + era5_state(TIMES[1], n) + + position("cells", "cell", 2 * n), + ) + } +{union("mesh_in", position("mesh", "node", 0))} +{mlp("grid_h", "grid_in", "embed.grid")} +{mlp("mesh_h", "mesh_in", "embed.mesh")} +{mlp("g2m_h", edge_features("g2m"), "embed.g2m")} +{mlp("multimesh_h", edge_features("multimesh"), "embed.multimesh")} +{mlp("m2g_h", edge_features("m2g"), "embed.m2g")} +{message_passing("mesh0", "encoder", "g2m.edges", "g2m_h", "grid_h", "mesh_h")} +{processor} +{ + message_passing( + "grid_out", + "decoder", + "m2g.edges", + "m2g_h", + f"mesh{rounds}", + "grid_h", + ) + } +forecast AS ( + SELECT g.id, p.g AS f, SUM(g.val * p.w) AS val + FROM grid_out g JOIN params p ON p.net = 'output' AND g.f = p.f + GROUP BY g.id, p.g)""" + + +def forecast_statement(rounds: int, where: str = "TRUE") -> str: + """The forecast change of each variable at each cell in ``where``.""" + return f"""WITH {model(rounds)} +SELECT c.latitude, c.longitude, f.f AS var, f.val AS change +FROM forecast f JOIN cells c ON f.id = c.cell +WHERE {where} +ORDER BY c.latitude, c.longitude, f.f""" + + +def gradient_statement(rounds: int) -> str: + """Each weight's gradient of the forecast's mean squared error + against what ERA5 recorded 6 hours later.""" + n = len(VARIABLES) + return f"""WITH {model(rounds)}, +{union("later", era5_state(TIMES[2], 0))} +error AS ( + SELECT AVG(power(f.val - (y.val - x.val), 2)) AS l + FROM forecast f + JOIN later y ON f.id = y.id AND f.f = y.f + JOIN grid_in x ON f.id = x.id AND f.f + {n} = x.f) +SELECT weight, grad(error, w) AS w FROM params ORDER BY weight""" + + +def receptive_field(host, rounds: int) -> list[tuple[str, int]]: + """How many cells and mesh nodes each stage needs for the region, + found by walking the edges backward from it: semi-joins in SQL.""" + + def count(sql): + return host.query(f"SELECT COUNT(*) AS n FROM ({sql}) t")["n"][ + 0 + ].as_py() + + needed = f"SELECT cell AS id FROM cells c WHERE {REGION}" + stages = [("decoder: grid cells forecast", count(needed))] + needed = ( + f"SELECT DISTINCT sender AS id FROM m2g.edges " + f"WHERE receiver IN ({needed})" + ) + stages.append((f"processor round {rounds}: mesh nodes", count(needed))) + for r in reversed(range(rounds)): + needed = ( + f"SELECT DISTINCT sender AS id FROM multimesh.edges " + f"WHERE receiver IN ({needed}) UNION SELECT id FROM ({needed}) n" + ) + label = f"processor round {r}" if r else "encoder" + stages.append((f"{label}: mesh nodes", count(needed))) + needed = ( + f"SELECT DISTINCT sender AS id FROM g2m.edges " + f"WHERE receiver IN ({needed})" + ) + stages.append(("encoder: grid cells read", count(needed))) + return stages + + +# -------------------------------------------------------------------------- +# The same model in JAX: the reference. +# -------------------------------------------------------------------------- + + +def jax_mlp(p, net, x): + h = x @ p[f"{net}.0"] + return (h * jax.nn.sigmoid(h)) @ p[f"{net}.1"] + + +def jax_message_passing(p, net, edges, edge_h, senders, receivers): + sender, receiver = edges + edge_in = jnp.concatenate( + [edge_h, senders[sender], receivers[receiver]], axis=1 + ) + message = jax_mlp(p, f"{net}.edge", edge_in) + total = jax.ops.segment_sum(message, receiver, len(receivers)) + node_in = jnp.concatenate([receivers, total], axis=1) + return receivers + jax_mlp(p, f"{net}.node", node_in) + + +def jax_forecast(p, graph, grid_in, rounds): + grid_h = jax_mlp(p, "embed.grid", grid_in) + mesh_h = jax_mlp(p, "embed.mesh", graph["mesh_in"]) + edge_h = { + g: jax_mlp(p, f"embed.{g}", graph[g]["features"]) + for g in ("g2m", "multimesh", "m2g") + } + mesh = jax_message_passing( + p, "encoder", graph["g2m"]["edges"], edge_h["g2m"], grid_h, mesh_h + ) + for r in range(rounds): + mesh = jax_message_passing( + p, + f"processor.{r}", + graph["multimesh"]["edges"], + edge_h["multimesh"], + mesh, + mesh, + ) + grid_out = jax_message_passing( + p, "decoder", graph["m2g"]["edges"], edge_h["m2g"], mesh, grid_h + ) + return grid_out @ p["output"] + + +def jax_error(p, graph, grid_in, target, rounds): + return jnp.mean((jax_forecast(p, graph, grid_in, rounds) - target) ** 2) + + +# -------------------------------------------------------------------------- +# Geometry, with numpy: the multi-mesh and the edges between it and the grid. +# -------------------------------------------------------------------------- + + +def icosahedron(): + t = (1 + 5**0.5) / 2 + vertices = [ + (-1, t, 0), + (1, t, 0), + (-1, -t, 0), + (1, -t, 0), + (0, -1, t), + (0, 1, t), + (0, -1, -t), + (0, 1, -t), + (t, 0, -1), + (t, 0, 1), + (-t, 0, -1), + (-t, 0, 1), + ] + faces = [ + (0, 11, 5), + (0, 5, 1), + (0, 1, 7), + (0, 7, 10), + (0, 10, 11), + (1, 5, 9), + (5, 11, 4), + (11, 10, 2), + (10, 7, 6), + (7, 1, 8), + (3, 9, 4), + (3, 4, 2), + (3, 2, 6), + (3, 6, 8), + (3, 8, 9), + (4, 9, 5), + (2, 4, 11), + (6, 2, 10), + (8, 6, 7), + (9, 8, 1), + ] + vertices = [np.array(v) / np.linalg.norm(v) for v in vertices] + return vertices, faces + + +def refine(vertices: list, faces: list) -> list: + """Split each triangle into four, adding its edges' midpoints, on the + sphere, to ``vertices``.""" + midpoints = {} + + def midpoint(a, b): + key = (min(a, b), max(a, b)) + if key not in midpoints: + m = vertices[a] + vertices[b] + vertices.append(m / np.linalg.norm(m)) + midpoints[key] = len(vertices) - 1 + return midpoints[key] + + finer = [] + for a, b, c in faces: + ab, bc, ca = midpoint(a, b), midpoint(b, c), midpoint(c, a) + finer += [(a, ab, ca), (b, bc, ab), (c, ca, bc), (ab, bc, ca)] + return finer + + +def multimesh(refinements: int): + """The vertices of an icosahedron refined ``refinements`` times, the + edges of every level of refinement in both directions, and the length + of the finest level's longest edge.""" + vertices, faces = icosahedron() + levels = [faces] + for _ in range(refinements): + levels.append(refine(vertices, levels[-1])) + edges = { + (u, v) + for level in levels + for face in level + for u, v in zip(face, face[1:] + face[:1]) + } + edges |= {(v, u) for u, v in edges} + vertices = np.array(vertices) + finest = max( + np.linalg.norm(vertices[u] - vertices[v]) + for face in levels[-1] + for u, v in zip(face, face[1:] + face[:1]) + ) + return vertices, np.array(sorted(edges)).T, finest + + +def xyz(latitude, longitude): + lat, lon = np.radians(latitude), np.radians(longitude) + return np.stack( + [np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], -1 + ) + + +def graph_table(senders, receivers, points_from, points_to) -> xr.Dataset: + """A graph as a Dataset: the table ``edges(edge, sender, receiver)`` + and the table ``features(edge, f, val)``.""" + offset = points_to[receivers] - points_from[senders] + features = np.concatenate( + [offset, np.linalg.norm(offset, axis=1, keepdims=True)], axis=1 + ) + return xr.Dataset( + { + "sender": ("edge", senders), + "receiver": ("edge", receivers), + "val": (("edge", "f"), features), + }, + coords={"edge": np.arange(len(senders)), "f": np.arange(EDGE_FEATURES)}, + ) + + +GRAPH_TABLES = {("edge",): "edges", ("edge", "f"): "features"} + + +def geometry(era5: xr.Dataset, refinements: int) -> dict[str, xr.Dataset]: + """The tables of the model's fixed structure, as Datasets.""" + lat, lon = xr.broadcast(era5.latitude, era5.longitude) + lat, lon = lat.values.ravel(), lon.values.ravel() + cell_xyz = xyz(lat, lon) + node_xyz, (senders, receivers), finest = multimesh(refinements) + node_lat = np.degrees(np.arcsin(node_xyz[:, 2])) + node_lon = np.degrees(np.arctan2(node_xyz[:, 1], node_xyz[:, 0])) % 360 + + # Grid to mesh: each cell to the mesh nodes within 0.6 of the finest + # mesh's longest edge, GraphCast's radius. + distance = np.linalg.norm(cell_xyz[:, None] - node_xyz[None], axis=-1) + g2m_cell, g2m_node = np.nonzero(distance < 0.6 * finest) + # Mesh to grid: each cell from its 3 nearest mesh nodes. + nearest = np.argsort(distance, axis=1)[:, :3] + m2g_cell = np.repeat(np.arange(len(lat)), 3) + m2g_node = nearest.ravel() + + return { + "cells": xr.Dataset( + {"latitude": ("cell", lat), "longitude": ("cell", lon)}, + coords={"cell": np.arange(len(lat))}, + ), + "mesh": xr.Dataset( + {"latitude": ("node", node_lat), "longitude": ("node", node_lon)}, + coords={"node": np.arange(len(node_xyz))}, + ), + "g2m": graph_table(g2m_cell, g2m_node, cell_xyz, node_xyz), + "multimesh": graph_table(senders, receivers, node_xyz, node_xyz), + "m2g": graph_table(m2g_node, m2g_cell, node_xyz, cell_xyz), + } + + +# -------------------------------------------------------------------------- +# Data and parameters. +# -------------------------------------------------------------------------- + + +def open_era5() -> xr.Dataset: + """WeatherBench 2's ERA5, lazily; a synthetic stand-in offline.""" + try: + return xr.open_zarr( + ERA5, chunks=None, storage_options={"token": "anon"} + ) + except Exception as e: # noqa: BLE001 - any failure to reach the store + print(f" ERA5 unavailable ({type(e).__name__}); using synthetic data") + lon = np.arange(64) * 5.625 + lat = -87.1875 + np.arange(32) * 5.625 + rng = np.random.default_rng(3) + surface = rng.normal(size=(len(TIMES), 64, 32)) + atmosphere = rng.normal(size=(len(TIMES), 2, 64, 32)) + data = { + name: ( + ("time", "longitude", "latitude"), + surface + i, + ) + for i, (table, name, _) in enumerate(VARIABLES) + if table == "surface" + } + data["geopotential"] = ( + ("time", "level", "longitude", "latitude"), + atmosphere, + ) + data["temperature"] = ( + ("time", "level", "longitude", "latitude"), + -atmosphere, + ) + return xr.Dataset( + data, + coords={ + "time": TIMES, + "level": [500, 850], + "longitude": lon, + "latitude": lat, + }, + ) + + +def era5_arrays(era5: xr.Dataset, time) -> np.ndarray: + """Each variable at ``time`` as a column, one row per cell, read with + xarray: the reference's input, independent of the SQL's.""" + columns = [] + for _, name, level in VARIABLES: + field = era5[name].sel(time=time) + if level: + field = field.sel(level=level) + columns.append(field.transpose("latitude", "longitude").values.ravel()) + return np.stack(columns, axis=1) + + +def init_params(rounds: int, rng) -> dict[str, np.ndarray]: + """Each network's weight matrices, by name, with fan-in scaling.""" + shapes = {"output": (HIDDEN, len(VARIABLES))} + + def two_layer(net, inputs): + shapes[f"{net}.0"] = (inputs, HIDDEN) + shapes[f"{net}.1"] = (HIDDEN, HIDDEN) + + two_layer("embed.grid", INPUTS) + two_layer("embed.mesh", 3) + for graph in ("g2m", "multimesh", "m2g"): + two_layer(f"embed.{graph}", EDGE_FEATURES) + for net in ["encoder", "decoder"] + [ + f"processor.{r}" for r in range(rounds) + ]: + two_layer(f"{net}.edge", 3 * HIDDEN) + two_layer(f"{net}.node", 2 * HIDDEN) + return { + net: rng.normal(0, 1 / np.sqrt(shape[0]), shape) + for net, shape in shapes.items() + } + + +def params_table(params: dict[str, np.ndarray]) -> xr.Dataset: + """Every weight as a row of ``params(weight, net, f, g, w)``.""" + rows = [ + (net, f, g, w) + for net, matrix in params.items() + for (f, g), w in np.ndenumerate(matrix) + ] + net, f, g, w = map(np.array, zip(*rows)) + return xr.Dataset( + { + "net": ("weight", net), + "f": ("weight", f), + "g": ("weight", g), + "w": ("weight", w), + }, + coords={"weight": np.arange(len(rows))}, + ) + + +def main(): + parser = cli(__doc__.splitlines()[0], sizes=[2], size_name="refinements") + parser.add_argument( + "--rounds", type=int, default=2, help="processor rounds" + ) + args = parser.parse_args() + report = Report("graphcast", "refinements", args.reps) + rounds = args.rounds + era5 = open_era5() + + # The normalization: each variable's mean and spread at the inputs. + now, before, later = (era5_arrays(era5, t) for t in TIMES[[1, 0, 2]]) + mean = np.concatenate([before, now]).mean(axis=0) + std = np.concatenate([before, now]).std(axis=0) + scale = xr.Dataset( + {"mean": ("var", mean), "std": ("var", std)}, + coords={"var": np.arange(len(VARIABLES))}, + ) + params = init_params(rounds, np.random.default_rng(0)) + + for refinements in args.sizes: + tables = geometry(era5, refinements) + cells = tables["cells"] + print( + f" mesh refined {refinements}×: {tables['mesh'].sizes['node']} " + f"nodes, {tables['multimesh'].sizes['edge']} mesh edges, " + f"{tables['g2m'].sizes['edge']} grid-to-mesh, " + f"{tables['m2g'].sizes['edge']} mesh-to-grid; " + f"{len(params_table(params).weight)} weights" + ) + + # The reference, in JAX, from arrays read with xarray. + graph = { + name: { + "edges": ( + tables[name].sender.values, + tables[name].receiver.values, + ), + "features": tables[name].val.values, + } + for name in ("g2m", "multimesh", "m2g") + } + graph["mesh_in"] = np.stack( + [ + np.cos(np.radians(tables["mesh"].latitude.values)), + np.sin(np.radians(tables["mesh"].longitude.values)), + np.cos(np.radians(tables["mesh"].longitude.values)), + ], + axis=1, + ) + position = np.stack( + [ + np.cos(np.radians(cells.latitude.values)), + np.sin(np.radians(cells.longitude.values)), + np.cos(np.radians(cells.longitude.values)), + ], + axis=1, + ) + grid_in = np.concatenate( + [(before - mean) / std, (now - mean) / std, position], axis=1 + ) + target = (later - now) / std + want = jax_forecast(params, graph, grid_in, rounds) + + want_grad = jax.grad(jax_error)(params, graph, grid_in, target, rounds) + want_grad = np.concatenate([want_grad[net].ravel() for net in params]) + + for host in hosts(args.hosts): + host.load("era5", era5, table_names=ERA5_TABLES) + for name, table in tables.items(): + options = ( + {"table_names": GRAPH_TABLES} + if "edge" in table.dims + else {} + ) + host.load(name, table, **options) + host.load("scale", scale) + host.load("params", params_table(params)) + + statement = forecast_statement(rounds) + rows = report.time( + "forecast, global", host, refinements, host.query, statement + ) + got = rows["change"].to_numpy().reshape(len(cells.cell), -1) + # The rows are ordered by latitude then longitude, as cells are. + report.check(assert_close("forecast", got, want, rtol=1e-9)) + report.note("operators", host.operators(statement)) + + statement = forecast_statement(rounds, REGION) + rows = report.time( + "forecast, regional", host, refinements, host.query, statement + ) + in_region = np.logical_and.reduce( + [ + (cells[c] >= lo) & (cells[c] <= hi) + for c, (lo, hi) in BOUNDS.items() + ] + ) + regional = rows["change"].to_numpy().reshape(-1, len(VARIABLES)) + report.check( + assert_close("regional", regional, want[in_region]) + ) + report.note("operators", host.operators(statement)) + + step = "gradient of the error" + if not host.differentiates: + report.skip(step, host, refinements, "no ddx backend") + continue + try: + rows = report.time( + step, + host, + refinements, + host.query, + gradient_statement(rounds), + ) + except DdxError as e: + # Recorded rather than raised: a gap in ddx, which this + # row tracks until ddx closes it. + reason = str(e).removeprefix("not supported by ddx-ad yet: ") + report.skip(step, host, refinements, f"ddx: {reason}") + continue + report.check( + assert_close("gradient", rows["w"].to_numpy(), want_grad) + ) + + print("\n receptive field of the regional forecast:") + for stage, n in receptive_field(host, rounds): + print(f" {stage:<34} {n:>6}") + + report.finish(args.csv) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md new file mode 100644 index 00000000..053a31b0 --- /dev/null +++ b/benchmarks/ml/README.md @@ -0,0 +1,169 @@ +# Machine learning in SQL + +These scripts write three machine-learning models as SQL over +[xarray-sql](../../README.md), take their gradients with +[ddx](https://github.com/xqlsystems/ddx), and time the same SQL on several +engines. The models are the three target demos of +[einfold](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md): + +| # | Script | Model | What it measures | +|---|---|---|---| +| 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, in einfold's window and `GROUP BY` spellings | the naive plan FlashAttention must beat; gradients with respect to Q, K and V; masked pairs as absent rows; a Nadaraya–Watson control | +| 02 | [`02_nanogpt.py`](02_nanogpt.py) | a 1-layer character-level GPT on Tiny Shakespeare | the loss on each host; SGD training steps, every parameter table checked after every step | +| 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature GraphCast (encoder, processor, decoder on an icosahedral multi-mesh) over ERA5 in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | + +The SQL favors legibility over speed. Each model reads like its definition, +and the slow parts are left slow on purpose: they are what ddx, xarray-sql +and einfold need to make fast. + +Every result is checked before it is timed into the table. Each script +computes the same model in JAX from the same inputs and parameters, and a +mismatch fails the run. + +## Running + +```shell +uv run benchmarks/ml/01_attention.py +uv run benchmarks/ml/02_nanogpt.py --steps 5 +uv run benchmarks/ml/03_graphcast.py --rounds 1 +``` + +Each script carries [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/) +and runs against the xarray-sql in this checkout, with the `ml` extra. The +options are the same in all three: + +```shell +--hosts datafusion duckdb # the engines to run on (default: all) +--sizes 64 256 # problem sizes: tokens, batch size, or mesh refinements +--reps 3 # timed repetitions; the median is reported +--csv results.csv # also write the results table +``` + +Demo 2 adds `--steps`, the number of training steps. Demo 3 adds +`--rounds`, the number of processor rounds, and `--lazy`, which registers +all of ERA5 (65 years, 60 variables) rather than the three time steps the +forecast reads. + +The output ends with one table: a row per step, host and size, with its +median time, its error against JAX, and the number of operators in the +host's physical plan. + +## The `ml` extra + +```shell +pip install "xarray-sql[ml]" +``` + +The extra installs [ddxdb](https://pypi.org/project/ddxdb/). With it, an +`xarray_sql.Context` understands `grad`, `vjp` and `jvp` in a select list, +and `ddx_stop_gradient` as a function. A gradient-descent step on a +parameter table is one statement: + +```sql +WITH pred AS ( + SELECT x.sample, SUM(x.value * w.val) AS y + FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample), +loss AS ( + SELECT AVG(power(pred.y - t.value, 2)) AS l + FROM pred JOIN t ON pred.sample = t.sample) +SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w +``` + +`xarray_sql.ad.sql_all(ctx, statements)` runs several such statements for +one backward pass, which is how Demo 2 updates its eight parameter tables. + +## Hosts + +einfold calls the engine that runs a plan its *host*. +[`_hosts.py`](_hosts.py) defines each one in a few lines: how it registers +an `xr.Dataset` (through xarray-sql), runs SQL, and counts the operators +in its plan. + +| Host | Forward pass | Gradients | +|---|---|---| +| `datafusion` | `xarray_sql.Context` | ddx | +| `duckdb` | DuckDB, through `xql.register` | not yet: ddx runs its programs on DataFusion only | + +A host that cannot differentiate still runs every forward query, and its +gradient rows say `no ddx backend`. To add a host, such as DataFusion with +einfold's optimizer or DuckDB with a GPU extension, write one class in +`_hosts.py` and add it to `HOSTS`. + +## Findings + +Measured on 2026-10-17 on an Intel Core i7-8700 (6 cores, 12 threads) +with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. + +**ddx's gradients agree with JAX.** The gradients of attention with +respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the +largest value at every size. After each of five SGD steps, all eight of +the GPT's parameter tables match JAX's to 1e-15, and the loss falls from +4.18 to 3.85 (each step on a different batch). + +**DataFusion computes a CTE once for every place that reads it.** A +model reads its activations more than once: a layer norm reads its input +three times, a residual connection twice, a message-passing round reads +the mesh as sender, receiver and residual. DataFusion inlines each read, +so the plan grows with the model's depth; DuckDB computes a CTE that is +read more than once a single time. + +| Query | DataFusion operators | DuckDB operators | DataFusion | DuckDB | +|---|---:|---:|---:|---:| +| GPT loss, 1 layer, batch 4 | 12,166 | 120 | 8.4 s | 0.055 s | +| GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | +| GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | + +A GPT training step, forward and backward through ddx, takes 2.7 s on +DataFusion: less than its forward query alone, because ddx stores the +result of each step of its program as a table. Computing shared +subexpressions once is the first thing xarray-sql's DataFusion path needs +for these models, and einfold should plan them as graphs rather than +trees. + +**Attention stores every score.** Causal attention's forward pass at 512 +tokens takes 0.68 s on DataFusion in the window spelling and 1.3 s in the +`GROUP BY` spelling; DuckDB takes 0.52 s and 0.26 s. Both grow with the +square of the tokens. The gradients take 3 to 4 times the forward pass. + +**A regional forecast costs what a global one does.** For the contiguous +United States, the forecast covers 55 of 2,048 grid cells. Walking the +graph back from them, the decoder needs 15 mesh nodes, the processor's +two rounds need 36 and then 82, and the encoder reads 1,103 grid cells +(54%); with one round, 576 (28%). Neither host uses any of this: the +regional forecast takes as long as the global one on both. Pushing the +region's filter back through the graph is einfold's Demo 3. + +**What ddx cannot differentiate yet.** Each of these raised a typed +error naming the construct, never a wrong number: + +- A window `MAX(s) OVER (...)` used as a value. Attention and the + cross-entropy subtract the row maximum before `exp`, and that shift + cancels out of the result, so the demos wrap it in `ddx_stop_gradient`, + which is exact. ddx could treat this as a rule. +- A `GROUP BY` key that depends on a differentiated column. einfold's + window spelling groups by the normalizer `z`; the demos divide inside + the sum instead, `SUM(w / z * v)`, which einfold also lists. +- `UNION ALL` between a parameter table and the loss. GraphCast + concatenates features with `UNION ALL`, so its gradient row records the + error. It passes once ddx differentiates a union. +- Any engine but DataFusion: ddx's DuckDB support is planned. + +**xarray-sql reads NaN as missing.** A NaN in a Dataset becomes SQL +`NULL`, which `SUM` skips, so Demo 1 checks masks with an infinity: a +plan that multiplies a masked pair's zero weight by it gives NaN where +SQL gives a number. + +## Next + +- Part B of the NanoGPT demo: the four tricks of the speedrun's sparsity + record, as a small diff to `02_nanogpt.py`, with Adam's state as + versioned tables. +- Peak memory per host, from a fresh process per measurement. +- An einfold host, once einfold produces plans. +- With `03_graphcast.py --lazy`, the bytes each forecast reads from ERA5. + +[`_bench.py`](_bench.py) holds the shared command line, timing, checks and +results table. + +[`_bench.py`](_bench.py) holds the shared command line, timing, checks and +results table. diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py new file mode 100644 index 00000000..815db316 --- /dev/null +++ b/benchmarks/ml/_bench.py @@ -0,0 +1,149 @@ +"""Timing, checking and reporting for the ML demos. + +Each demo is a script with the same command line:: + + uv run benchmarks/ml/01_attention.py --hosts datafusion duckdb \\ + --sizes 64 256 --reps 3 --csv attention.csv + +It runs every step on every chosen host and size, checks each result +against a JAX reference, and ends with one table: a row per step, host +and size, with the median time and whether the result matched. Checks +come first: a wrong answer fails the run, however fast it was. +""" + +from __future__ import annotations + +import argparse +import csv +import statistics +import sys +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field + +import numpy as np +from _hosts import HOSTS, Host + + +def cli(description: str, sizes: Sequence[int], size_name: str = "n"): + """The shared command line, for a demo to extend and parse.""" + parser = argparse.ArgumentParser(description=description) + parser.add_argument( + "--hosts", + nargs="+", + default=list(HOSTS), + choices=list(HOSTS), + help="engines to run on (default: all)", + ) + parser.add_argument( + "--sizes", + nargs="+", + type=int, + default=list(sizes), + help=f"problem sizes, as {size_name} (default: %(default)s)", + ) + parser.add_argument( + "--reps", + type=int, + default=1, + help="timed repetitions per step; the median is reported", + ) + parser.add_argument("--csv", help="also write the results table here") + return parser + + +def hosts(names: Sequence[str]) -> list[Host]: + """A fresh instance of each named host.""" + return [HOSTS[name]() for name in names] + + +def assert_close(name: str, got, want, *, rtol=1e-9, atol=1e-9) -> float: + """Assert two arrays match, NaN for NaN. Returns the largest error, + relative to the largest value.""" + got, want = np.asarray(got), np.asarray(want) + np.testing.assert_allclose( + got, want, rtol=rtol, atol=atol, equal_nan=True, err_msg=name + ) + finite = np.isfinite(want) + scale = np.max(np.abs(want[finite]), initial=0.0) or 1.0 + return float(np.max(np.abs(got - want)[finite], initial=0.0) / scale) + + +@dataclass +class Report: + """The demo's results: one row per step, host and size.""" + + demo: str + size_name: str + reps: int = 1 + rows: list[dict] = field(default_factory=list) + + def time(self, step: str, host: Host, size: int, run: Callable, *args): + """Run ``run(*args)`` ``reps`` times, record the median, and return + the last result.""" + times = [] + for _ in range(self.reps): + start = time.perf_counter() + result = run(*args) + times.append(time.perf_counter() - start) + self.rows.append( + { + "demo": self.demo, + "step": step, + "host": host.name, + self.size_name: size, + "seconds": statistics.median(times), + "check": "", + } + ) + print( + f" {step:<28} {host.name:<11} {self.size_name}={size:<6} " + f"{statistics.median(times):8.3f}s", + flush=True, + ) + return result + + def check(self, error: float) -> None: + """Mark the last row's result as checked, with its error (from + :func:`assert_close`).""" + self.rows[-1]["check"] = f"ok ({error:.0e})" + + def skip(self, step: str, host: Host, size: int, why: str) -> None: + """Record a step the host cannot run, and why.""" + self.rows.append( + { + "demo": self.demo, + "step": step, + "host": host.name, + self.size_name: size, + "seconds": None, + "check": why, + } + ) + + def note(self, key: str, value) -> None: + """Attach a measurement besides time to the last row.""" + self.rows[-1][key] = value + + def finish(self, csv_path: str | None = None) -> None: + """Print the table, and write it as CSV if asked.""" + columns = list(dict.fromkeys(k for row in self.rows for k in row)) + print(f"\n| {' | '.join(columns)} |") + print(f"|{'|'.join('---' for _ in columns)}|") + for row in self.rows: + cells = [_cell(row.get(c)) for c in columns] + print(f"| {' | '.join(cells)} |") + if csv_path: + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=columns) + writer.writeheader() + writer.writerows(self.rows) + print(f"\nwrote {csv_path}", file=sys.stderr) + + +def _cell(value) -> str: + if value is None: + return "n/a" + if isinstance(value, float): + return f"{value:.3f}" + return str(value) diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py new file mode 100644 index 00000000..9e41bbbd --- /dev/null +++ b/benchmarks/ml/_hosts.py @@ -0,0 +1,111 @@ +"""The SQL engines ("hosts") the ML demos run on. + +einfold calls the engine that runs a plan its *host*. Every demo writes +its model once, as SQL, and runs that same text on each host, so a +difference in time is a difference between engines, never between +queries. A host does four things: + +* ``load(name, data, **options)`` registers an ``xr.Dataset`` as the + table ``name`` through xarray-sql, replacing any table of that name + (``options`` such as ``table_names`` go to xarray-sql); +* ``query(sql)`` runs a statement and returns its rows as Arrow; +* ``query_all(sqls)`` runs several statements, and, on a host that + differentiates, shares the backward pass between those that take + ``grad`` of the same loss; +* ``operators(sql)`` counts the operators in the physical plan it would + run, which shows, for one, whether a CTE read twice is computed twice. + +Forward queries run everywhere. Gradients (``grad``, ``vjp``, ``jvp``) +need ddx, which today runs its programs on DataFusion only; a host that +cannot differentiate says so with ``differentiates = False``, and the +demos report its gradient steps as not run rather than skip it silently. + +On every host, ``ddx_stop_gradient(x)`` is the identity. It is how a +model marks a value whose derivative is zero (the max subtracted before a +softmax's ``exp``), and it changes nothing in a forward pass. + +Adding a host (einfold's optimizer on DataFusion, DuckDB with a GPU +extension, or DuckDB running ddx's programs once ddx supports it) is one +class here and one entry in ``HOSTS``. +""" + +from __future__ import annotations + +import json +from collections.abc import Sequence + +import duckdb +import pyarrow as pa +import xarray as xr + +import xarray_sql as xql + + +class Host: + """An engine that runs the demos' SQL. Subclasses fill in the rest.""" + + name = "" + differentiates = False + + def load(self, name: str, data: xr.Dataset, **options) -> None: + raise NotImplementedError + + def query(self, sql: str) -> pa.Table: + raise NotImplementedError + + def query_all(self, sqls: Sequence[str]) -> list[pa.Table]: + return [self.query(sql) for sql in sqls] + + def operators(self, sql: str) -> int: + raise NotImplementedError + + +class DataFusion(Host): + """xarray-sql's ``Context``, with ddx's gradients (the ``ml`` extra).""" + + name = "datafusion" + differentiates = True + + def __init__(self): + self.ctx = xql.Context() + + def load(self, name, data, **options): + self.ctx.deregister_table(name) + self.ctx.from_dataset(name, data, **options) + + def query(self, sql): + return self.ctx.sql(sql).to_arrow_table() + + def query_all(self, sqls): + return [df.to_arrow_table() for df in xql.ad.sql_all(self.ctx, sqls)] + + def operators(self, sql): + plan = self.ctx.sql(sql).execution_plan().display_indent() + return len(plan.splitlines()) + + +class DuckDB(Host): + """DuckDB, reading the same Datasets through xarray-sql's adapter.""" + + name = "duckdb" + + def __init__(self): + self.con = duckdb.connect() + self.con.sql("CREATE MACRO ddx_stop_gradient(x) AS x") + + def load(self, name, data, **options): + xql.register(self.con, name, data, **options) + + def query(self, sql): + return self.con.sql(sql).fetch_arrow_table() + + def operators(self, sql): + (_, plan), *_ = self.con.sql(f"EXPLAIN (FORMAT JSON) {sql}").fetchall() + return sum(_count(node) for node in json.loads(plan)) + + +def _count(node: dict) -> int: + return 1 + sum(_count(child) for child in node.get("children", [])) + + +HOSTS = {host.name: host for host in (DataFusion, DuckDB)} diff --git a/pyproject.toml b/pyproject.toml index 8d80eae7..f9384fb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,11 +58,17 @@ polars = [ geo = [ "pyproj", ] +# Gradients in SQL: grad(), vjp() and jvp() from ddx. ddxdb is pre-1.0 +# and breaks its API between minor versions. It reads and writes the +# Substrait of DataFusion 54, so bump it together with that pin. +ml = [ + "ddxdb>=0.3.0,<0.4; sys_platform != 'emscripten'", +] test = [ "adbc-driver-flightsql>=1.12", "adbc-driver-sqlite>=1.12", "cftime", - "xarray-sql[adbc,dask,duckdb,polars,geo]", + "xarray-sql[adbc,dask,duckdb,polars,geo,ml]", "pytest", "xarray[io]", "gcsfs", diff --git a/tests/test_ad.py b/tests/test_ad.py new file mode 100644 index 00000000..b395a247 --- /dev/null +++ b/tests/test_ad.py @@ -0,0 +1,99 @@ +"""Tests for the ddx gradient extension (`xarray_sql.ad`).""" + +import numpy as np +import pytest +import xarray as xr + +pytest.importorskip("ddxdb") + +from ddxdb import DdxError + +from xarray_sql import Context, ad + +# Least squares: pred = x @ w, loss = sum((pred - y)^2). +LOSS = """ +pred AS ( + SELECT x.sample, SUM(x.x * w.w) AS p + FROM data x JOIN weights w ON x.feature = w.feature + GROUP BY x.sample), +loss AS ( + SELECT SUM(power(pred.p - y.y, 2)) AS l + FROM pred JOIN targets y ON pred.sample = y.sample) +""" + + +@pytest.fixture +def problem(): + rng = np.random.default_rng(0) + x = rng.normal(size=(6, 3)) + w = rng.normal(size=3) + y = rng.normal(size=6) + sample, feature = np.arange(6), np.arange(3) + ctx = Context() + ctx.from_dataset( + "data", + xr.Dataset( + {"x": (("sample", "feature"), x)}, + coords={"sample": sample, "feature": feature}, + ), + ) + ctx.from_dataset( + "weights", + xr.Dataset({"w": ("feature", w)}, coords={"feature": feature}), + ) + ctx.from_dataset( + "targets", xr.Dataset({"y": ("sample", y)}, coords={"sample": sample}) + ) + return ctx, x, w, y + + +def test_grad_matches_closed_form(problem): + ctx, x, w, y = problem + got = ctx.sql( + f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature" + ).to_dataset() + np.testing.assert_allclose(got["w"].values, 2 * x.T @ (x @ w - y)) + + +def test_gradient_step_reads_like_sql(problem): + ctx, x, w, y = problem + got = ctx.sql( + f"WITH {LOSS} " + "SELECT feature, w - 0.01 * grad(loss, w) AS w FROM weights " + "ORDER BY feature" + ).to_pandas() + np.testing.assert_allclose(got["w"], w - 0.01 * 2 * x.T @ (x @ w - y)) + + +def test_sql_all_shares_one_loss(problem): + ctx, x, w, y = problem + grad = 2 * x.T @ (x @ w - y) + select = "SELECT feature, {} AS w FROM weights ORDER BY feature" + gradient, step = ad.sql_all( + ctx, + [ + f"WITH {LOSS} " + select.format("grad(loss, w)"), + f"WITH {LOSS} " + select.format("w - 0.01 * grad(loss, w)"), + ], + ) + np.testing.assert_allclose(gradient.to_dataset()["w"].values, grad) + np.testing.assert_allclose(step.to_dataset()["w"].values, w - 0.01 * grad) + + +def test_stop_gradient_blocks_the_derivative(problem): + ctx, _, w, _ = problem + got = ctx.sql( + "WITH loss AS (SELECT SUM(w * ddx_stop_gradient(w)) AS l " + "FROM weights) " + "SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature" + ).to_dataset() + np.testing.assert_allclose(got["w"].values, w) + + +def test_unsupported_expression_raises(problem): + ctx, *_ = problem + with pytest.raises(DdxError): + ctx.sql( + "WITH loss AS (SELECT SUM(w) AS l FROM weights) " + "SELECT feature, grad(loss, feature) AS g FROM weights" + ).collect() diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py new file mode 100644 index 00000000..b8a9c5fd --- /dev/null +++ b/xarray_sql/ad.py @@ -0,0 +1,131 @@ +"""Gradients of SQL queries: the optional ml extension, built on ddx. + +[ddx](https://github.com/xqlsystems/ddx) differentiates a whole SQL query +the way JAX differentiates a function. With it installed +(``pip install xarray-sql[ml]``), [xarray_sql.Context][]'s ``.sql()`` +understands three calls in a select list: + +* ``grad(loss, column)``: at each row, the derivative of the one number + the CTE ``loss`` computes with respect to that row's ``column`` + (``table.column`` when the ``SELECT`` reads several tables). +* ``vjp(f, column, cotangent)``: the same for any CTE ``f``, pulled back + along ``cotangent``, a table or CTE shaped like ``f``'s output. +* ``jvp(e, table.column, tangent)``: the tangent of the expression ``e`` + along ``tangent``, pushed forward. + +A gradient-descent step on a parameter table is then one statement:: + + WITH pred AS ( + SELECT x.sample, SUM(x.value * w.val) AS y + FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample), + loss AS ( + SELECT AVG(power(pred.y - t.value, 2)) AS l + FROM pred JOIN t ON pred.sample = t.sample) + SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w + +The query being differentiated is ordinary SQL, over tables registered +with ``from_dataset`` or any other way. ddx plans the backward pass as +Substrait, DataFusion runs it, and the statement reads what it computed. + +The one function ddx gives a meaning to is ``ddx_stop_gradient(x)``, the +identity whose derivative is zero (JAX's ``lax.stop_gradient``). It is +registered on every ``Context``. Its typical use is the stable softmax, +whose shift by the row maximum cancels and so carries no gradient:: + + exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i))) + +A statement that calls none of ``grad``, ``vjp`` or ``jvp`` is planned +exactly as it would be without ddx. An expression ddx cannot +differentiate raises a ``ddxdb.DdxError`` subclass naming it, never a +wrong number. + +[sql_all][xarray_sql.ad.sql_all] runs several statements that take +``grad`` of the same loss for the price of one backward pass, which is +how a training step updates each of a model's parameter tables. For the +programs underneath (``grad``, ``vjp``, ``jvp`` and ``run``, to build a +backward pass once and run it every step), use ``ddxdb.ad`` directly; +they take a ``Context`` like any other ``SessionContext``. +""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +from datafusion import SessionContext +from ddxdb import ad as _ddx + +__all__ = ["STOP_GRADIENT", "calls_ad", "register", "sql", "sql_all"] + +STOP_GRADIENT = _ddx.STOP_GRADIENT +"""The SQL name of the stop-gradient function, ``ddx_stop_gradient``.""" + +_AD_CALL = re.compile(r"\b(grad|vjp|jvp)\s*\(", re.IGNORECASE) + + +def register(ctx: SessionContext) -> None: + """Register ``ddx_stop_gradient`` on ``ctx``. + + Works on any DataFusion ``SessionContext`` (``Context`` registers it + automatically when ddxdb is installed). ``grad``, ``vjp`` and ``jvp`` + are not UDFs: [sql][xarray_sql.ad.sql] rewrites them before DataFusion + plans the statement. + + Args: + ctx: The DataFusion session context to register the UDF on. + """ + _ddx.register_stop_gradient(ctx) + + +def calls_ad(statement: str) -> bool: + """Whether ``statement`` may call ``grad``, ``vjp`` or ``jvp``. + + A cheap textual check, so statements without a call never reach + ddx's parser. A false positive (the word inside a string literal, + say) is harmless: ddx plans a statement with no call as it is. + """ + return _AD_CALL.search(statement) is not None + + +def sql(ctx: SessionContext, statement: str): + """Run ``statement``, computing its ``grad``, ``vjp`` and ``jvp`` calls. + + Returns a DataFusion DataFrame; ``Context.sql`` wraps it as an + [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame]. + + Args: + ctx: Any DataFusion ``SessionContext``. + statement: A SQL statement. + """ + return _ddx.sql(ctx, statement) + + +def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list: + """Run several statements, sharing the programs their calls need. + + Statements that take ``grad`` of the same loss share one backward + pass, so a training step can update each parameter table with its + own statement and pay for the gradient once:: + + steps = xql.ad.sql_all(ctx, [ + f"WITH loss AS ({LOSS}) " + "SELECT i, val - 0.1 * grad(loss, val) AS val FROM w1", + f"WITH loss AS ({LOSS}) " + "SELECT j, val - 0.1 * grad(loss, val) AS val FROM w2", + ]) + + On a [Context][xarray_sql.Context], each result is an + [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame]. + + Args: + ctx: Any DataFusion ``SessionContext``. + statements: SQL statements. + + Returns: + One DataFrame per statement, in order. + """ + if isinstance(statements, str): + raise TypeError("statements is a list of SQL statements") + frames = _ddx.sql_all(ctx, list(statements)) + wrap = getattr(ctx, "_wrap", None) + return [wrap(f) for f in frames] if wrap else frames diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index b5f3a9c8..31dd1e85 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -25,6 +25,14 @@ except ImportError: # pragma: no cover - depends on the environment _proj = None +_ad: ModuleType | None +try: # ddxdb is an optional dependency (`pip install xarray-sql[ml]`). + from . import ad + + _ad = ad +except ImportError: # pragma: no cover - depends on the environment + _ad = None + class Context(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s. @@ -45,6 +53,10 @@ def __init__(self, *args, **kwargs): # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform. if _proj is not None: _proj.register(self) + # With ddxdb installed, .sql() also takes gradients: grad(), vjp() + # and jvp() in a select list, and ddx_stop_gradient(). + if _ad is not None: + _ad.register(self) def from_dataset( self, @@ -214,6 +226,10 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: ``.to_dataset(dimension_columns=[...])`` for round-tripping the result back to an ``xr.Dataset``. + With the ml extension installed (``pip install xarray-sql[ml]``), + the query may also call ``grad``, ``vjp`` and ``jvp`` (see + [xarray_sql.ad][]). + Args: query: A SQL query string. *args: Forwarded to ``SessionContext.sql``. @@ -222,7 +238,17 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: Returns: An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame. """ - inner = super().sql(query, *args, **kwargs) + if _ad is not None and _ad.calls_ad(query): + if args or kwargs: + raise TypeError( + "a query that calls grad, vjp or jvp takes no other " + "arguments" + ) + return self._wrap(_ad.sql(self, query)) + return self._wrap(super().sql(query, *args, **kwargs)) + + def _wrap(self, inner) -> XarrayDataFrame: + """Wrap a DataFusion DataFrame so it can round-trip to xarray.""" return XarrayDataFrame(inner, templates=self._registered_datasets) From f3df43fe00a8abff11fd09265923093dc52dae85 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 19:12:09 -0700 Subject: [PATCH 02/11] Address review: complete SQL statements, nanoGPT's model, introductions - Attention: each spelling's forward pass and gradient is now one complete SQL statement. The gradient with respect to Q, K and V is a single query (grad of three columns of a join), as fast as three statements sharing a backward pass. - nanoGPT: match model.py. Layer norms have learned gains (ln_1, ln_2, ln_f), the MLP uses GELU (its tanh approximation; DataFusion has no erf), and the residual projections use GPT-2's scaled initialization. The hyperparameters name their GPTConfig counterparts, and the docstring lists what still differs: size, optimizer, dropout. - TODO(xqlsystems/ddx#142) on each window MAX wrapped in ddx_stop_gradient. - Introduce every project, paper and term before using it, in the demos, their README and xarray_sql.ad; record the rule, and the TODO exception to the no-issue-numbers rule, in AGENTS.md. - Drop the README's "Gradients in SQL" section until the feature is ready to announce, and the stale --lazy flag from benchmarks/ml/README.md. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- AGENTS.md | 11 ++- README.md | 24 ----- benchmarks/ml/01_attention.py | 164 ++++++++++++++++++++++++---------- benchmarks/ml/02_nanogpt.py | 153 +++++++++++++++++++++---------- benchmarks/ml/03_graphcast.py | 22 +++-- benchmarks/ml/README.md | 50 ++++++----- benchmarks/ml/_bench.py | 6 +- benchmarks/ml/_hosts.py | 7 +- xarray_sql/ad.py | 8 +- 9 files changed, 291 insertions(+), 154 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fc419a15..dfb638b5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,7 +8,16 @@ summarizes recurring maintainer review feedback so changes land clean. - Keep docstrings and comments self-contained. Do **not** put GitHub issue or PR numbers in docstrings or code comments; a reader should not need the issue tracker to understand the code. Issue references belong in the commit message - and PR description (e.g. `Closes #189`), not in the source. + and PR description (e.g. `Closes #189`), not in the source. The one exception + is a workaround's `TODO`, which names the issue that would remove it: + `TODO(#NN): `, or `TODO(owner/repo#NN): ` for an + issue in another repository. The comment still explains the workaround. +- Introduce every reference before relying on it. The first time a document + (a docstring, comment, README or docs page) mentions a project, paper, + algorithm, dataset or term of art, such as einfold, FlashAttention or a + section of another project's design, say in a few words what it is, and link + it where there is a link. Don't assume the reader has the context of the + conversation that produced the code. - Do not reference the review conversation, chat, or "the reporter" in comments. Describe the behavior, not how it came up. diff --git a/README.md b/README.md index 12f3ecba..42fe05a7 100644 --- a/README.md +++ b/README.md @@ -237,30 +237,6 @@ Some highlights: * Dense tensor operations are slower in SQL than Xarray. It seems like the vanilla HashJoin implementation in standard DB engines are not as performant as array ops (they don't take advantage of contiguous array layouts), though this gives us some idea for how it could [one day be possible](https://github.com/xqlsystems/einfold). * Sparse tensor operations are just as fast (or faster!) than Xarray today. Specifically, SQL is an excellent way to perform regridding and reprojecting between different spatial scales (see also: we include a `proj` udf via the `xarray-sql[geo]` extra). -## Gradients in SQL - -With the `ml` extra (`pip install "xarray-sql[ml]"`), a `Context` takes -gradients of queries with [ddx](https://github.com/xqlsystems/ddx). In a -select list, `grad(loss, column)` is, at each row, the derivative of the -number the CTE `loss` computes with respect to that row's `column`. A -gradient-descent step on a parameter table is then one statement: - -```python -ctx.sql(""" -WITH pred AS ( - SELECT x.sample, SUM(x.value * w.val) AS y - FROM x JOIN w ON x.feature = w.feature GROUP BY x.sample), -loss AS ( - SELECT AVG(power(pred.y - t.value, 2)) AS l - FROM pred JOIN t ON pred.sample = t.sample) -SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w -""") -``` - -[`benchmarks/ml`](benchmarks/ml) trains a small GPT and runs attention and -a GraphCast-shaped weather model this way, checked against JAX and timed on -several engines. - ## Why does this work? Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index d6b3c932..f3d1b539 100644 --- a/benchmarks/ml/01_attention.py +++ b/benchmarks/ml/01_attention.py @@ -8,28 +8,39 @@ # [tool.uv.sources] # xarray-sql = { path = "../../", editable = true } # /// -"""Demo 1: attention in SQL, the query einfold should turn into FlashAttention. - -Attention, the core of a transformer, is ``softmax(Q·Kᵀ/√d)·V``. In the -XQL model, ``Q``, ``K`` and ``V`` are tables of ``(t, d, x)`` (token, -feature, value), and attention is two folds over joins: an einsum for the -scores, and a weighted average of ``V`` with weights ``exp(score)``. - -The SQL below is the naive plan, which stores every one of the N×N -scores. FlashAttention never does: it streams tiles of keys past tiles of -queries, with an online softmax. einfold's first target is to find that -plan from this SQL, in both of its spellings, with no attention-specific -rule (einfold's ``docs/demos.md``, Demo 1). This script is the baseline -that plan must match and beat: - -* **forward**: each spelling on each host, against a JAX reference; -* **gradient**: ddx differentiates the same SQL with respect to Q, K - and V, against ``jax.grad``; +"""Demo 1: causal attention in SQL, on several engines, with its gradient. + +Attention, the core of a transformer, is ``softmax(Q·Kᵀ/√d)·V``: each +token's output is an average of the value vectors ``V``, weighted by how +well the token's query in ``Q`` matches each earlier token's key in ``K``. +Here ``Q``, ``K`` and ``V`` are tables of ``(t, d, x)`` (token, feature, +value), registered from xarray Datasets, and attention is two sums over +joins: the scores, a sum of ``Q`` times ``K`` over features, and the +average of ``V`` weighted by ``exp(score)``. + +This SQL is the naive plan: it stores every one of the N×N scores. +FlashAttention (Dao et al., 2022, https://arxiv.org/abs/2205.14135) +computes the same result without storing them, by streaming blocks of +keys past blocks of queries and keeping a running softmax. einfold +(https://github.com/xqlsystems/einfold), a query optimizer that rewrites +SQL plans into array algorithms like this one, aims to find that plan +from this SQL, in both of the spellings below, with no rule written for +attention. It is the first of einfold's three target demos +(``docs/demos.md`` in its repository). This script is the baseline that +plan must match and then beat: + +* **forward**: each spelling on each *host* (einfold's word for an engine + that runs a plan; here DataFusion and DuckDB), checked against the same + attention in JAX, Google's library for differentiable array programs; +* **gradient**: the gradient of a loss with respect to every value of Q, + K and V, as one SQL statement. ddx (https://github.com/xqlsystems/ddx), + which differentiates SQL queries, computes it, and ``jax.grad`` checks + it; * **masks are absent rows**: an infinity in V reaches only the queries allowed to see it, as SQL's semantics require of any faster plan; * **a control that is not attention**: the Nadaraya–Watson estimator, a - Gaussian-kernel weighted average used in regridding. The rules that - fuse attention must fuse it too. + Gaussian-kernel weighted average used in regridding. Any rule that + speeds up attention must speed up this query too. Run ``uv run benchmarks/ml/01_attention.py --help`` for the options. """ @@ -48,12 +59,14 @@ D = 64 # features per head; the scores are scaled by 1/√64 = 1/8. -# The scores, softmax and weighted average, in einfold's window spelling. -# The causal mask is a join predicate: a masked (i, j) pair is a row that -# does not exist. The max only keeps exp() from overflowing and cancels -# out of the ratio, so it carries no gradient (ddx_stop_gradient). +# Attention in einfold's window spelling. The causal mask is a join +# predicate: a masked (i, j) pair is a row that does not exist. Subtracting +# each row's max only keeps exp() from overflowing: it cancels out of the +# ratio w / z, so it carries no gradient, which ddx_stop_gradient says. +# TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX +# once ddx can differentiate a window aggregate; it refuses one today. ATTENTION_WINDOW = """ -scores AS ( +WITH scores AS ( SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s FROM q JOIN k ON q.d = k.d WHERE q.t >= k.t @@ -69,12 +82,45 @@ JOIN v ON w.j = v.t JOIN norm n ON w.i = n.i GROUP BY w.i, v.d) +SELECT t, d, x FROM attention ORDER BY t, d """ -# The same attention as ddx spells it: the row max as a GROUP BY, joined -# back to the scores. +# The gradient of half the output's squared norm, with respect to every +# value of Q, K and V: the same query, plus a loss and a different SELECT. +ATTENTION_WINDOW_GRADIENT = """ +WITH scores AS ( + SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s + FROM q JOIN k ON q.d = k.d + WHERE q.t >= k.t + GROUP BY q.t, k.t), +weights AS ( + SELECT i, j, exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i))) AS w + FROM scores), +norm AS ( + SELECT i, SUM(w) AS z FROM weights GROUP BY i), +attention AS ( + SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x + FROM weights w + JOIN v ON w.j = v.t + JOIN norm n ON w.i = n.i + GROUP BY w.i, v.d), +loss AS ( + SELECT SUM(x * x) / 2 AS l FROM attention) +SELECT q.t, q.d, + grad(loss, q.x) AS dq, + grad(loss, k.x) AS dk, + grad(loss, v.x) AS dv +FROM q +JOIN k ON q.t = k.t AND q.d = k.d +JOIN v ON q.t = v.t AND q.d = v.d +ORDER BY q.t, q.d +""" + +# The same attention as ddx spells it: each row's max as a GROUP BY, +# joined back to the scores. Here ddx_stop_gradient isn't needed, but it +# is exact, and it spares ddx differentiating the max. ATTENTION_GROUP_BY = """ -scores AS ( +WITH scores AS ( SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s FROM q JOIN k ON q.d = k.d WHERE q.t >= k.t @@ -92,18 +138,51 @@ JOIN v ON w.j = v.t JOIN norm n ON w.i = n.i GROUP BY w.i, v.d) +SELECT t, d, x FROM attention ORDER BY t, d """ -SPELLINGS = {"window": ATTENTION_WINDOW, "group by": ATTENTION_GROUP_BY} +ATTENTION_GROUP_BY_GRADIENT = """ +WITH scores AS ( + SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s + FROM q JOIN k ON q.d = k.d + WHERE q.t >= k.t + GROUP BY q.t, k.t), +row_max AS ( + SELECT i, ddx_stop_gradient(MAX(s)) AS m FROM scores GROUP BY i), +weights AS ( + SELECT s.i, s.j, exp(s.s - r.m) AS w + FROM scores s JOIN row_max r ON s.i = r.i), +norm AS ( + SELECT i, SUM(w) AS z FROM weights GROUP BY i), +attention AS ( + SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x + FROM weights w + JOIN v ON w.j = v.t + JOIN norm n ON w.i = n.i + GROUP BY w.i, v.d), +loss AS ( + SELECT SUM(x * x) / 2 AS l FROM attention) +SELECT q.t, q.d, + grad(loss, q.x) AS dq, + grad(loss, k.x) AS dk, + grad(loss, v.x) AS dv +FROM q +JOIN k ON q.t = k.t AND q.d = k.d +JOIN v ON q.t = v.t AND q.d = v.d +ORDER BY q.t, q.d +""" -# A scalar of the output to differentiate: half its squared norm. -LOSS = "loss AS (SELECT SUM(x * x) / 2 AS l FROM attention)" +# Each spelling's forward query and gradient query. +SPELLINGS = { + "window": (ATTENTION_WINDOW, ATTENTION_WINDOW_GRADIENT), + "group by": (ATTENTION_GROUP_BY, ATTENTION_GROUP_BY_GRADIENT), +} # Nadaraya–Watson: at each target point, the average of the sources' y # weighted by a Gaussian kernel of their distance (bandwidth h = 0.5). # Structurally attention: a weighted average with weights exp(score). NADARAYA_WATSON = """ -kernel AS ( +WITH kernel AS ( SELECT t.p, s.q, exp(-(power(t.x - s.x, 2) + power(t.y - s.y, 2)) / 0.25) AS w FROM target t CROSS JOIN source s) SELECT k.p, SUM(k.w * s.v) / SUM(k.w) AS v @@ -155,8 +234,7 @@ def main(): for name, data in zip("qkv", (q, k, v)): host.load(name, table(data)) - for spelling, ctes in SPELLINGS.items(): - forward = f"WITH {ctes} SELECT * FROM attention ORDER BY t, d" + for spelling, (forward, gradient) in SPELLINGS.items(): rows = report.time( f"forward ({spelling})", host, n, host.query, forward ) @@ -167,16 +245,10 @@ def main(): if not host.differentiates: report.skip(step, host, n, "no ddx backend") continue - # One statement per input, sharing one backward pass. - gradients = [ - f"WITH {ctes}, {LOSS} " - f"SELECT t, d, grad(loss, x) AS x FROM {name} ORDER BY t, d" - for name in "qkv" - ] - grads = report.time(step, host, n, host.query_all, gradients) + rows = report.time(step, host, n, host.query, gradient) errors = [ - assert_close(f"d/d{name}", matrix(got), ref) - for name, got, ref in zip("qkv", grads, want_grads) + assert_close(f"d/d{name}", matrix(rows, f"d{name}"), ref) + for name, ref in zip("qkv", want_grads) ] report.check(max(errors)) @@ -198,10 +270,8 @@ def check_masks_are_absent_rows(host_names): for host in hosts(host_names): for name, data in zip("qkv", (q, k, v)): host.load(name, table(data)) - for spelling, ctes in SPELLINGS.items(): - out = matrix( - host.query(f"WITH {ctes} SELECT * FROM attention ORDER BY t, d") - ) + for spelling, (forward, _) in SPELLINGS.items(): + out = matrix(host.query(forward)) assert np.isfinite(out[:j]).all(), (host.name, spelling) assert np.isposinf(out[j:, 0]).all(), (host.name, spelling) assert np.isfinite(out[:, 1:]).all(), (host.name, spelling) @@ -233,7 +303,7 @@ def check_nadaraya_watson(host_names, report): host, len(targets), host.query, - f"WITH {NADARAYA_WATSON}", + NADARAYA_WATSON, ) got = xql.to_dataset(rows, dims=["p"])["v"].values report.check(assert_close("nadaraya-watson", got, want)) diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index 62df9ccb..225aa880 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -8,19 +8,23 @@ # [tool.uv.sources] # xarray-sql = { path = "../../", editable = true } # /// -"""Demo 2: a GPT in SQL, trained by gradient descent through ddx. +"""Demo 2: a GPT in SQL, trained by gradient descent. -nanoGPT is Andrej Karpathy's small, readable GPT-2. This is its model as -one SQL query: the cross-entropy loss of a batch of characters from -Shakespeare, computed by a transformer whose parameters are tables. +nanoGPT (https://github.com/karpathy/nanoGPT) is Andrej Karpathy's small, +readable implementation of GPT-2, the language model. This is its model, +``model.py``, as one SQL query: the cross-entropy loss of a batch of +characters from Shakespeare, computed by a transformer whose parameters +are tables. ddx (https://github.com/xqlsystems/ddx), which differentiates +SQL queries, takes the loss's gradient with respect to every parameter, +so a training step is SQL too. ================= ============================================ Model piece In SQL ================= ============================================ Embeddings a join of the batch with ``wte`` and ``wpe`` -Layer norm per-position ``AVG``s, joined back -Attention Demo 1's query, per head, causal -MLP two einsums around ReLU², ``power(greatest(x, 0), 2)`` +Layer norm per-position ``AVG``s, joined back, times a gain +Attention per head and causal, as in ``01_attention.py`` +MLP two einsums around GELU Output layer an einsum with ``wte`` again (tied, as in GPT-2) Loss log-sum-exp per position, minus the target's logit Gradients ddx: ``grad(loss, w)``, one statement per table @@ -28,14 +32,36 @@ ================= ============================================ Each training step updates every parameter table with its own -statement; the statements share one backward pass. The same model in JAX, -from the same initial parameters and the same batches, is the reference: -its loss at every step and its gradients must match the SQL's. - -This is Part A of einfold's Demo 2 (``docs/demos.md``). Part B, the four -tricks of the NanoGPT speedrun's sparsity record as a small diff to this -SQL, builds on it. Simplifications from nanoGPT: no dropout, no biases, -and layer norms without a learned gain. +statement; the statements share one backward pass. The same model in JAX +(Google's library for differentiable array programs), from the same +initial parameters and the same batches, is the reference: its loss at +every step, and its parameters after every step, must match the SQL's. + +This is Part A of the second target demo of einfold +(https://github.com/xqlsystems/einfold), a query optimizer that rewrites +SQL plans into array algorithms; the demo is described in einfold's +``docs/demos.md``. Part B builds on it: it writes the tricks behind a +record of the NanoGPT speedrun (a public race to train GPT-2 fastest) +as a small diff to this SQL. + +Differences from nanoGPT's ``model.py`` and ``train.py``: + +* **Size.** nanoGPT's Shakespeare model has 6 layers, 6 heads, 384 + channels and 256 characters of context, 10.65M parameters. This one has + 1 layer, 2 heads, 16 channels and 16 characters, 4,416 parameters (the + constants below). +* **GELU** is its tanh approximation, as GPT-2 computed it. nanoGPT's + ``nn.GELU()`` uses the exact form, through the error function ``erf``, + which DataFusion doesn't have. The two differ by less than 5e-4. +* **The optimizer** is SGD at a fixed learning rate. nanoGPT trains with + AdamW, warmup, cosine decay and gradient clipping. +* **No dropout**, nanoGPT's default, so that a run is deterministic and + can be checked. Its Shakespeare configuration trains with 0.2. +* **Q, K and V are three tables.** nanoGPT's ``c_attn`` holds the same + parameters as one matrix, split after the product. + +Like nanoGPT's default (``bias = False`` in ``train.py``), the model has +no biases, in linear layers or layer norms. Run ``uv run benchmarks/ml/02_nanogpt.py --help`` for the options. """ @@ -55,25 +81,34 @@ jax.config.update("jax_enable_x64", True) -# A tiny character-level GPT. -LAYERS = 1 # transformer blocks -EMBED = 16 # e: channels of the residual stream -HEADS = 2 # h: attention heads +# A tiny character-level GPT. The comments name each constant in nanoGPT's +# GPTConfig (https://github.com/karpathy/nanoGPT/blob/master/model.py) and +# its value in nanoGPT's Shakespeare model +# (https://github.com/karpathy/nanoGPT/blob/master/config/train_shakespeare_char.py). +LAYERS = 1 # n_layer = 6: transformer blocks +EMBED = 16 # n_embd = 384; e: channels of the residual stream +HEADS = 2 # n_head = 6; h: attention heads HEAD = EMBED // HEADS # c: channels per head -HIDDEN = 4 * EMBED # f: channels of the MLP -CONTEXT = 16 # pos: characters per sequence -LEARNING_RATE = 0.5 - -# Each parameter table's dimensions; its values are the column w. +HIDDEN = 4 * EMBED # f: channels of the MLP, 4 × n_embd as in nanoGPT +CONTEXT = 16 # block_size = 256; pos: characters per sequence +LEARNING_RATE = 0.5 # for SGD; nanoGPT's AdamW uses 1e-3 + +# Each parameter table's dimensions; its values are the column w. In +# nanoGPT, wq, wk and wv are attn.c_attn; wo is attn.c_proj; w_fc and +# w_proj are mlp.c_fc and mlp.c_proj; ln_1, ln_2 and ln_f are the layer +# norms' weights. PARAMS = { "wte": ("tok", "e"), "wpe": ("pos", "e"), + "ln_1": ("layer", "e"), "wq": ("layer", "e", "h", "c"), "wk": ("layer", "e", "h", "c"), "wv": ("layer", "e", "h", "c"), "wo": ("layer", "h", "c", "e"), + "ln_2": ("layer", "e"), "w_fc": ("layer", "e", "f"), "w_proj": ("layer", "f", "e"), + "ln_f": ("e",), } # -------------------------------------------------------------------------- @@ -82,8 +117,10 @@ # -------------------------------------------------------------------------- -def layer_norm(x: str, out: str) -> str: - """Normalize each position of ``x`` to zero mean and unit variance.""" +def layer_norm(x: str, out: str, gain: str, layer: int | None = None) -> str: + """Normalize each position of ``x`` to zero mean and unit variance, + then scale each channel by the table ``gain`` (of ``layer``, if set).""" + of_layer = "" if layer is None else f" AND g.layer = {layer}" return f""" {out}_mean AS ( SELECT seq, pos, AVG(val) AS mean FROM {x} GROUP BY seq, pos), @@ -92,10 +129,15 @@ def layer_norm(x: str, out: str) -> str: FROM {x} x JOIN {out}_mean m ON x.seq = m.seq AND x.pos = m.pos GROUP BY x.seq, x.pos), {out} AS ( - SELECT x.seq, x.pos, x.e, (x.val - m.mean) / sqrt(v.var + 1e-5) AS val + SELECT x.seq, x.pos, x.e, (x.val - m.mean) / sqrt(v.var + 1e-5) * g.w AS val FROM {x} x JOIN {out}_mean m ON x.seq = m.seq AND x.pos = m.pos - JOIN {out}_var v ON x.seq = v.seq AND x.pos = v.pos),""" + JOIN {out}_var v ON x.seq = v.seq AND x.pos = v.pos + JOIN {gain} g ON x.e = g.e{of_layer}),""" + + +# GELU's tanh approximation: x/2 · (1 + tanh(√(2/π) · (x + 0.044715 x³))). +GELU_SCALE = np.sqrt(2 / np.pi) def block(layer: int, x: str, out: str) -> str: @@ -111,7 +153,7 @@ def block(layer: int, x: str, out: str) -> str: GROUP BY x.seq, x.pos, p.h, p.c),""" for name in ("q", "k", "v") ) - return f"""{layer_norm(x, f"{l}_ln1")}{projections} + return f"""{layer_norm(x, f"{l}_ln1", "ln_1", layer)}{projections} {l}_scores AS ( SELECT q.seq, q.h, q.pos AS i, k.pos AS j, SUM(q.val * k.val) / {np.sqrt(HEAD)} AS s @@ -119,6 +161,9 @@ def block(layer: int, x: str, out: str) -> str: WHERE q.pos >= k.pos GROUP BY q.seq, q.h, q.pos, k.pos), {l}_weights AS ( + -- TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX + -- once ddx can differentiate a window aggregate. The max only keeps exp() + -- from overflowing, and cancels out of p / z, so its gradient is zero. SELECT seq, h, i, j, exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY seq, h, i))) AS p FROM {l}_scores), @@ -138,12 +183,17 @@ def block(layer: int, x: str, out: str) -> str: {l}_mid AS ( SELECT x.seq, x.pos, x.e, x.val + a.val AS val FROM {x} x JOIN {l}_attn a ON x.seq = a.seq AND x.pos = a.pos AND x.e = a.e), -{layer_norm(f"{l}_mid", f"{l}_ln2")} -{l}_hidden AS ( - SELECT x.seq, x.pos, p.f, power(greatest(SUM(x.val * p.w), 0.0), 2) AS val +{layer_norm(f"{l}_mid", f"{l}_ln2", "ln_2", layer)} +{l}_fc AS ( + SELECT x.seq, x.pos, p.f, SUM(x.val * p.w) AS val FROM {l}_ln2 x JOIN w_fc p ON x.e = p.e WHERE p.layer = {layer} GROUP BY x.seq, x.pos, p.f), +{l}_hidden AS ( + SELECT seq, pos, f, + 0.5 * val * (1 + tanh({GELU_SCALE} * (val + 0.044715 * val * val * val))) + AS val + FROM {l}_fc), {l}_mlp AS ( SELECT x.seq, x.pos, p.e, SUM(x.val * p.w) AS val FROM {l}_hidden x JOIN w_proj p ON x.f = p.f @@ -167,12 +217,13 @@ def model() -> str: JOIN wte te ON b.tok = te.tok JOIN wpe pe ON b.pos = pe.pos AND te.e = pe.e), {blocks} -{layer_norm(f"x{LAYERS}", "final")} +{layer_norm(f"x{LAYERS}", "final", "ln_f")} logits AS ( SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val FROM final x JOIN wte te ON x.e = te.e GROUP BY x.seq, x.pos, te.tok), shifted AS ( + -- TODO(xqlsystems/ddx#142): as in the attention weights above. SELECT seq, pos, tok, val, ddx_stop_gradient(MAX(val) OVER (PARTITION BY seq, pos)) AS m FROM logits), @@ -205,17 +256,17 @@ def loss_statement() -> str: # -------------------------------------------------------------------------- -def jax_layer_norm(x): +def jax_layer_norm(x, gain): mean = x.mean(-1, keepdims=True) var = ((x - mean) ** 2).mean(-1, keepdims=True) - return (x - mean) / jnp.sqrt(var + 1e-5) + return (x - mean) / jnp.sqrt(var + 1e-5) * gain def jax_loss(p, tok, nxt): x = p["wte"][tok] + p["wpe"][None] causal = jnp.tril(jnp.ones((CONTEXT, CONTEXT), dtype=bool)) for layer in range(LAYERS): - h = jax_layer_norm(x) + h = jax_layer_norm(x, p["ln_1"][layer]) q, k, v = ( jnp.einsum("bte,ehc->bhtc", h, p[w][layer]) for w in ("wq", "wk", "wv") @@ -224,10 +275,10 @@ def jax_loss(p, tok, nxt): a = jax.nn.softmax(jnp.where(causal, s, -jnp.inf), axis=-1) heads = jnp.einsum("bhij,bhjc->bhic", a, v) x = x + jnp.einsum("bhtc,hce->bte", heads, p["wo"][layer]) - h = jax_layer_norm(x) - hidden = jnp.maximum(h @ p["w_fc"][layer], 0.0) ** 2 + h = jax_layer_norm(x, p["ln_2"][layer]) + hidden = jax.nn.gelu(h @ p["w_fc"][layer], approximate=True) x = x + hidden @ p["w_proj"][layer] - logits = jax_layer_norm(x) @ p["wte"].T + logits = jax_layer_norm(x, p["ln_f"]) @ p["wte"].T target = jnp.take_along_axis(logits, nxt[..., None], -1)[..., 0] return jnp.mean(jax.nn.logsumexp(logits, -1) - target) @@ -273,7 +324,9 @@ def batches(text: str, size: int, rng): def init_params(vocab: int, rng) -> dict[str, np.ndarray]: - """GPT-2's initialization: normal, standard deviation 0.02.""" + """nanoGPT's initialization, from GPT-2: weights normal with standard + deviation 0.02, or 0.02/√(2 × layers) for the two projections that add + to the residual stream; layer norm gains 1.""" sizes = { "tok": vocab, "pos": CONTEXT, @@ -283,10 +336,16 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]: "c": HEAD, "f": HIDDEN, } - return { - name: rng.normal(0, 0.02, [sizes[d] for d in dims]) - for name, dims in PARAMS.items() - } + params = {} + for name, dims in PARAMS.items(): + shape = [sizes[d] for d in dims] + if name.startswith("ln_"): + params[name] = np.ones(shape) + elif name in ("wo", "w_proj"): + params[name] = rng.normal(0, 0.02 / np.sqrt(2 * LAYERS), shape) + else: + params[name] = rng.normal(0, 0.02, shape) + return params def as_table(name: str, values: np.ndarray) -> xr.Dataset: @@ -309,8 +368,12 @@ def main(): text = shakespeare() vocab = len(set(text)) lines = len(loss_statement().splitlines()) + count = sum( + v.size for v in init_params(vocab, np.random.default_rng(0)).values() + ) print( - f" {LAYERS}-layer GPT, vocab {vocab}: the loss is {lines} lines of SQL" + f" {LAYERS}-layer GPT, vocab {vocab}, {count:,} parameters: " + f"the loss is {lines} lines of SQL" ) for size in args.sizes: diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py index 3ca61684..8b0232c7 100644 --- a/benchmarks/ml/03_graphcast.py +++ b/benchmarks/ml/03_graphcast.py @@ -21,8 +21,10 @@ round; * a **decoder** maps the mesh back onto the grid, as the 6-hour change. -Here it is in miniature, as SQL over ERA5 read through xarray-sql from -WeatherBench 2's Zarr store (or a synthetic stand-in, offline). Every +Here it is in miniature, as SQL over ERA5, ECMWF's reanalysis of the +atmosphere since 1940. xarray-sql reads it from the Zarr store of +WeatherBench 2, a benchmark for weather models (or a synthetic stand-in, +offline). Every graph is a table of edges ``(edge, sender, receiver)``, every state a table of ``(id, f, val)``, and every stage the same relational steps: @@ -38,11 +40,17 @@ The geometry (the mesh, and which grid cells and mesh nodes are joined) is computed once with numpy; the model, its inputs and its forecast are -SQL. The script checks the forecast, and ddx's gradient of the forecast -error, against the same model in JAX. Then it measures what einfold's -Demo 3 is about (``docs/demos.md``): how much less a regional forecast -needs than a global one. It reports the receptive field: how many mesh -nodes and grid cells each stage needs for the region. +SQL. The script checks the forecast, and the gradient of the forecast +error, against the same model in JAX, Google's library for differentiable +array programs. ddx (https://github.com/xqlsystems/ddx), which +differentiates SQL queries, computes the gradient. + +Then it measures what the third target demo of einfold +(https://github.com/xqlsystems/einfold, a query optimizer that rewrites +SQL plans into array algorithms; the demo is in its ``docs/demos.md``) is +about: how much less a regional forecast needs than a global one. It +reports the receptive field: how many mesh nodes and grid cells each +stage needs for the region. Simplifications from GraphCast: random weights, a 64×32 grid and six variables, a mesh refined 2 times rather than 6, 2 processor rounds diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 053a31b0..5c286f41 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -1,16 +1,25 @@ # Machine learning in SQL -These scripts write three machine-learning models as SQL over -[xarray-sql](../../README.md), take their gradients with -[ddx](https://github.com/xqlsystems/ddx), and time the same SQL on several -engines. The models are the three target demos of -[einfold](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md): +These scripts write three machine-learning models as SQL over tables that +[xarray-sql](../../README.md) registers from xarray Datasets, and time the +same SQL on several engines. Two related projects shape them: + +- [ddx](https://github.com/xqlsystems/ddx) differentiates SQL queries: it + turns the gradient of a query into more SQL, which trains the models. +- [einfold](https://github.com/xqlsystems/einfold) is a query optimizer + that rewrites SQL plans into array algorithms. These models are its + three [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md), + and this SQL is the baseline its plans must beat. + +Each script checks its results against the same model in +[JAX](https://github.com/jax-ml/jax), Google's library for differentiable +array programs. | # | Script | Model | What it measures | |---|---|---|---| -| 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, in einfold's window and `GROUP BY` spellings | the naive plan FlashAttention must beat; gradients with respect to Q, K and V; masked pairs as absent rows; a Nadaraya–Watson control | -| 02 | [`02_nanogpt.py`](02_nanogpt.py) | a 1-layer character-level GPT on Tiny Shakespeare | the loss on each host; SGD training steps, every parameter table checked after every step | -| 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature GraphCast (encoder, processor, decoder on an icosahedral multi-mesh) over ERA5 in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | +| 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, in einfold's window and `GROUP BY` spellings | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average | +| 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, on Tiny Shakespeare | the loss on each host; SGD training steps, every parameter table checked after every step | +| 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature [GraphCast](https://arxiv.org/abs/2212.12794), DeepMind's graph neural network for weather forecasting, over [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5) reanalysis in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | The SQL favors legibility over speed. Each model reads like its definition, and the slow parts are left slow on purpose: they are what ddx, xarray-sql @@ -40,9 +49,7 @@ options are the same in all three: ``` Demo 2 adds `--steps`, the number of training steps. Demo 3 adds -`--rounds`, the number of processor rounds, and `--lazy`, which registers -all of ERA5 (65 years, 60 variables) rather than the three time steps the -forecast reads. +`--rounds`, the number of processor rounds. The output ends with one table: a row per step, host and size, with its median time, its error against JAX, and the number of operators in the @@ -70,7 +77,7 @@ SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w ``` `xarray_sql.ad.sql_all(ctx, statements)` runs several such statements for -one backward pass, which is how Demo 2 updates its eight parameter tables. +one backward pass, which is how Demo 2 updates its eleven parameter tables. ## Hosts @@ -96,9 +103,9 @@ with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. **ddx's gradients agree with JAX.** The gradients of attention with respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the -largest value at every size. After each of five SGD steps, all eight of +largest value at every size. After each of five SGD steps, all eleven of the GPT's parameter tables match JAX's to 1e-15, and the loss falls from -4.18 to 3.85 (each step on a different batch). +4.18 to 3.84 (each step on a different batch). **DataFusion computes a CTE once for every place that reads it.** A model reads its activations more than once: a layer norm reads its input @@ -109,11 +116,11 @@ read more than once a single time. | Query | DataFusion operators | DuckDB operators | DataFusion | DuckDB | |---|---:|---:|---:|---:| -| GPT loss, 1 layer, batch 4 | 12,166 | 120 | 8.4 s | 0.055 s | +| GPT loss, 1 layer, batch 4 | 12,794 | 126 | 9.4 s | 0.058 s | | GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | | GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | -A GPT training step, forward and backward through ddx, takes 2.7 s on +A GPT training step, forward and backward through ddx, takes 3.3 s on DataFusion: less than its forward query alone, because ddx stores the result of each step of its program as a table. Computing shared subexpressions once is the first thing xarray-sql's DataFusion path needs @@ -155,15 +162,12 @@ SQL gives a number. ## Next -- Part B of the NanoGPT demo: the four tricks of the speedrun's sparsity - record, as a small diff to `02_nanogpt.py`, with Adam's state as - versioned tables. +- Part B of the NanoGPT demo: the four tricks behind a record of the + [NanoGPT speedrun](https://github.com/KellerJordan/modded-nanogpt), a + public race to train GPT-2 fastest, as a small diff to `02_nanogpt.py`. + It needs nanoGPT's optimizer, AdamW, with its state as versioned tables. - Peak memory per host, from a fresh process per measurement. - An einfold host, once einfold produces plans. -- With `03_graphcast.py --lazy`, the bytes each forecast reads from ERA5. - -[`_bench.py`](_bench.py) holds the shared command line, timing, checks and -results table. [`_bench.py`](_bench.py) holds the shared command line, timing, checks and results table. diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py index 815db316..eed1f918 100644 --- a/benchmarks/ml/_bench.py +++ b/benchmarks/ml/_bench.py @@ -5,8 +5,10 @@ uv run benchmarks/ml/01_attention.py --hosts datafusion duckdb \\ --sizes 64 256 --reps 3 --csv attention.csv -It runs every step on every chosen host and size, checks each result -against a JAX reference, and ends with one table: a row per step, host +It runs every step on every chosen host (an engine that runs SQL; see +``_hosts.py``) and size, checks each result against the same model in +JAX, Google's library for differentiable array programs, and ends with +one table: a row per step, host and size, with the median time and whether the result matched. Checks come first: a wrong answer fails the run, however fast it was. """ diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index 9e41bbbd..8784a0e2 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -1,6 +1,8 @@ """The SQL engines ("hosts") the ML demos run on. -einfold calls the engine that runs a plan its *host*. Every demo writes +einfold (https://github.com/xqlsystems/einfold), a query optimizer that +rewrites SQL plans into array algorithms, calls the engine that runs a +plan its *host*. Every demo writes its model once, as SQL, and runs that same text on each host, so a difference in time is a difference between engines, never between queries. A host does four things: @@ -16,7 +18,8 @@ run, which shows, for one, whether a CTE read twice is computed twice. Forward queries run everywhere. Gradients (``grad``, ``vjp``, ``jvp``) -need ddx, which today runs its programs on DataFusion only; a host that +need ddx (https://github.com/xqlsystems/ddx), which differentiates SQL +queries and today runs its programs on DataFusion only; a host that cannot differentiate says so with ``differentiates = False``, and the demos report its gradient steps as not run rather than skip it silently. diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py index b8a9c5fd..9d692293 100644 --- a/xarray_sql/ad.py +++ b/xarray_sql/ad.py @@ -1,7 +1,8 @@ """Gradients of SQL queries: the optional ml extension, built on ddx. [ddx](https://github.com/xqlsystems/ddx) differentiates a whole SQL query -the way JAX differentiates a function. With it installed +the way [JAX](https://github.com/jax-ml/jax), Google's library for +differentiable array programs, differentiates a function. With it installed (``pip install xarray-sql[ml]``), [xarray_sql.Context][]'s ``.sql()`` understands three calls in a select list: @@ -24,8 +25,9 @@ SELECT feature, val - 0.1 * grad(loss, val) AS val FROM w The query being differentiated is ordinary SQL, over tables registered -with ``from_dataset`` or any other way. ddx plans the backward pass as -Substrait, DataFusion runs it, and the statement reads what it computed. +with ``from_dataset`` or any other way. ddx plans the backward pass in +[Substrait](https://substrait.io), a portable format for query plans; +DataFusion runs it, and the statement reads what it computed. The one function ddx gives a meaning to is ``ddx_stop_gradient(x)``, the identity whose derivative is zero (JAX's ``lax.stop_gradient``). It is From e3777c7bb10f55b3ba36218a4f4b0423ed05ae3f Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 20:56:40 -0700 Subject: [PATCH 03/11] Train the GPT with nanoGPT's AdamW behind --optimizer adamw SGD at a fixed learning rate stays the default. With --optimizer adamw, a step stores each parameter table's gradient as a table (one backward pass for all of them), then updates each table with one SQL statement: gradients clipped to a global norm of 1, Adam's moments as columns beside the weights, weight decay on all but the layer norms, and the learning rate read from a schedule(step, lr) table with nanoGPT's warmup and cosine decay, stretched to --steps. The JAX reference step matches torch.optim.AdamW after clip_grad_norm_ exactly, and the SQL matches it to 6e-15 after each of five steps. Also say in the attention demo that its three grad() calls share one backward pass. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/01_attention.py | 2 + benchmarks/ml/02_nanogpt.py | 206 +++++++++++++++++++++++++++++----- benchmarks/ml/README.md | 12 +- 3 files changed, 186 insertions(+), 34 deletions(-) diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index f3d1b539..16cb60c0 100644 --- a/benchmarks/ml/01_attention.py +++ b/benchmarks/ml/01_attention.py @@ -87,6 +87,8 @@ # The gradient of half the output's squared norm, with respect to every # value of Q, K and V: the same query, plus a loss and a different SELECT. +# The three grad() calls differentiate one loss, so ddx runs one backward +# pass for all three, as it would for three statements given to sql_all. ATTENTION_WINDOW_GRADIENT = """ WITH scores AS ( SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index 225aa880..bdac788d 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -28,11 +28,16 @@ Output layer an einsum with ``wte`` again (tied, as in GPT-2) Loss log-sum-exp per position, minus the target's logit Gradients ddx: ``grad(loss, w)``, one statement per table -Optimizer step ``w - lr * grad(loss, w)``, the new table +Optimizer step ``w - lr * grad(loss, w)``, the new table; or AdamW, + with Adam's moments as columns beside ``w`` and + the learning rate schedule as a table ================= ============================================ Each training step updates every parameter table with its own -statement; the statements share one backward pass. The same model in JAX +statement; the statements share one backward pass. By default the +optimizer is SGD at a fixed learning rate, the simplest step to read. +``--optimizer adamw`` trains as nanoGPT does instead: it stores the +gradients as tables, then updates each parameter table with AdamW. The same model in JAX (Google's library for differentiable array programs), from the same initial parameters and the same batches, is the reference: its loss at every step, and its parameters after every step, must match the SQL's. @@ -53,8 +58,13 @@ * **GELU** is its tanh approximation, as GPT-2 computed it. nanoGPT's ``nn.GELU()`` uses the exact form, through the error function ``erf``, which DataFusion doesn't have. The two differ by less than 5e-4. -* **The optimizer** is SGD at a fixed learning rate. nanoGPT trains with - AdamW, warmup, cosine decay and gradient clipping. +* **The optimizer**, by default, is SGD at a fixed learning rate. + ``--optimizer adamw`` matches nanoGPT's: AdamW with its Shakespeare + settings, weight decay on all but the layer norms, gradients clipped to + a norm of 1, and the learning rate warming up, then decaying along a + cosine. The schedule is stretched to ``--steps``, where nanoGPT's spans + 5,000. The JAX reference for this step matches ``torch.optim.AdamW`` + after ``clip_grad_norm_`` exactly. * **No dropout**, nanoGPT's default, so that a run is deterministic and can be checked. Its Shakespeare configuration trains with 0.2. * **Q, K and V are three tables.** nanoGPT's ``c_attn`` holds the same @@ -91,7 +101,18 @@ HEAD = EMBED // HEADS # c: channels per head HIDDEN = 4 * EMBED # f: channels of the MLP, 4 × n_embd as in nanoGPT CONTEXT = 16 # block_size = 256; pos: characters per sequence -LEARNING_RATE = 0.5 # for SGD; nanoGPT's AdamW uses 1e-3 +LEARNING_RATE = 0.5 # for SGD; nanoGPT's AdamW uses 1e-3, below + +# AdamW (``--optimizer adamw``), with the settings of nanoGPT's Shakespeare +# model: train.py's defaults, and config/train_shakespeare_char.py's. +MAX_LR = 1e-3 # learning_rate +MIN_LR = 1e-4 # min_lr +WARMUP = 100 / 5000 # warmup_iters, as a fraction of max_iters +BETA1 = 0.9 # beta1 +BETA2 = 0.99 # beta2 +EPS = 1e-8 # torch.optim.AdamW's default +WEIGHT_DECAY = 0.1 # weight_decay, on every table but the layer norms' gains +GRAD_CLIP = 1.0 # grad_clip: the largest norm of all gradients together # Each parameter table's dimensions; its values are the column w. In # nanoGPT, wq, wk and wv are attn.c_attn; wo is attn.c_proj; w_fc and @@ -247,6 +268,53 @@ def step_statements(lr: float) -> list[str]: ] +def gradient_statements() -> list[str]: + """The gradient of the loss with respect to each parameter table, as + the table ``{name}_grad``'s rows. The statements share one backward + pass.""" + return [ + f"WITH {model()} SELECT {', '.join(dims)}, grad(loss, w) AS g " + f"FROM {table} ORDER BY {', '.join(dims)}" + for table, dims in PARAMS.items() + ] + + +def adamw_statement(table: str, step: int) -> str: + """One AdamW step on a parameter table with columns ``w``, ``m`` and + ``v`` (the weight and Adam's two moments), from its gradient in + ``{table}_grad`` and the learning rate in ``schedule``. This is + ``torch.optim.AdamW``, after nanoGPT clips the gradients.""" + dims = PARAMS[table] + keys = ", ".join(f"p.{d}" for d in dims) + same_row = " AND ".join(f"p.{d} = g.{d}" for d in dims) + decay = 0.0 if table.startswith("ln_") else WEIGHT_DECAY + every_grad = "\n UNION ALL ".join( + f"SELECT g FROM {t}_grad" for t in PARAMS + ) + return f""" +WITH grad_norm AS ( + SELECT sqrt(SUM(g * g)) AS norm FROM ( + {every_grad})), +clip AS ( + SELECT LEAST(1.0, {GRAD_CLIP} / (norm + 1e-6)) AS c FROM grad_norm), +moments AS ( + SELECT {keys}, p.w, + {BETA1} * p.m + (1 - {BETA1}) * g.g * clip.c AS m, + {BETA2} * p.v + (1 - {BETA2}) * power(g.g * clip.c, 2) AS v + FROM {table} p + JOIN {table}_grad g ON {same_row} + CROSS JOIN clip) +SELECT {keys}, + p.w * (1 - s.lr * {decay}) + - s.lr * (p.m / (1 - power({BETA1}, s.step + 1))) + / (sqrt(p.v / (1 - power({BETA2}, s.step + 1))) + {EPS}) AS w, + p.m, + p.v +FROM moments p JOIN schedule s ON s.step = {step} +ORDER BY {keys} +""" + + def loss_statement() -> str: return f"WITH {model()} SELECT l FROM loss" @@ -283,6 +351,26 @@ def jax_loss(p, tok, nxt): return jnp.mean(jax.nn.logsumexp(logits, -1) - target) +def jax_adamw(params, moments, grads, step, lr): + """The reference AdamW step: new parameters and moments.""" + norm = np.sqrt(sum(float(jnp.sum(g**2)) for g in grads.values())) + clip = min(1.0, GRAD_CLIP / (norm + 1e-6)) + new_params, new_moments = {}, {} + for name, w in params.items(): + g = grads[name] * clip + m, v = moments[name] + m = BETA1 * m + (1 - BETA1) * g + v = BETA2 * v + (1 - BETA2) * g**2 + decay = 0.0 if name.startswith("ln_") else WEIGHT_DECAY + m_hat = m / (1 - BETA1 ** (step + 1)) + v_hat = v / (1 - BETA2 ** (step + 1)) + new_params[name] = w * (1 - lr * decay) - lr * m_hat / ( + jnp.sqrt(v_hat) + EPS + ) + new_moments[name] = (m, v) + return new_params, new_moments + + # -------------------------------------------------------------------------- # Data and parameters, as xarray Datasets. # -------------------------------------------------------------------------- @@ -348,22 +436,59 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]: return params -def as_table(name: str, values: np.ndarray) -> xr.Dataset: +def as_table(name: str, w: np.ndarray, m=None, v=None) -> xr.Dataset: + """A parameter as a table: its weights ``w``, and Adam's moments ``m`` + and ``v`` when they are given.""" dims = PARAMS[name] + columns = {"w": w, "m": m, "v": v} return xr.Dataset( - {"w": (dims, values)}, - coords={d: np.arange(n) for d, n in zip(dims, values.shape)}, + {k: (dims, x) for k, x in columns.items() if x is not None}, + coords={d: np.arange(n) for d, n in zip(dims, w.shape)}, ) -def as_array(name: str, rows) -> np.ndarray: - return xql.to_dataset(rows, dims=list(PARAMS[name]))["w"].values +def as_array(name: str, rows, column: str = "w") -> np.ndarray: + return xql.to_dataset(rows, dims=list(PARAMS[name]))[column].values + + +def schedule(steps: int) -> xr.Dataset: + """nanoGPT's learning rate at each step, as the table + ``schedule(step, lr)``: a linear warmup, then a cosine decay to the + minimum, with train.py's ``max_iters`` and ``lr_decay_iters`` both + ``steps``.""" + warmup = round(WARMUP * steps) + lr = [] + for it in range(steps): + if it < warmup: + lr.append(MAX_LR * (it + 1) / (warmup + 1)) + else: + decay = (it - warmup) / (steps - warmup) + coeff = 0.5 * (1 + np.cos(np.pi * decay)) + lr.append(MIN_LR + coeff * (MAX_LR - MIN_LR)) + return xr.Dataset({"lr": ("step", lr)}, coords={"step": np.arange(steps)}) + + +def adamw_step(host, step: int) -> list: + """One AdamW step on every parameter table: the gradients, in one + backward pass, registered as tables; then each table's update.""" + gradients = host.query_all(gradient_statements()) + for name, rows in zip(PARAMS, gradients): + host.load(f"{name}_grad", xql.to_dataset(rows, dims=list(PARAMS[name]))) + return [host.query(adamw_statement(name, step)) for name in PARAMS] def main(): parser = cli(__doc__.splitlines()[0], sizes=[4], size_name="batch") - parser.add_argument("--steps", type=int, default=5, help="SGD steps") + parser.add_argument("--steps", type=int, default=5, help="training steps") + parser.add_argument( + "--optimizer", + choices=["sgd", "adamw"], + default="sgd", + help="SGD at a fixed learning rate (default), or nanoGPT's AdamW", + ) args = parser.parse_args() + adamw = args.optimizer == "adamw" + lr = schedule(args.steps) report = Report("nanogpt", "batch", args.reps) text = shakespeare() vocab = len(set(text)) @@ -386,8 +511,13 @@ def main(): for host in hosts(args.hosts): reference = dict(init) + moments = { + n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items() + } for name, values in init.items(): - host.load(name, as_table(name, values)) + state = moments[name] if adamw else () + host.load(name, as_table(name, values, *state)) + host.load("schedule", lr) for step, batch in data: host.load("batch", batch) @@ -407,31 +537,47 @@ def main(): report.note("loss", round(float(want), 6)) report.note("operators", host.operators(loss_statement())) + train = f"train ({args.optimizer}), step {step}" if not host.differentiates: - report.skip( - f"train, step {step}", host, size, "no ddx backend" - ) + report.skip(train, host, size, "no ddx backend") continue - updated = report.time( - f"train, step {step}", - host, - size, - host.query_all, - step_statements(LEARNING_RATE), - ) grads = jax.grad(jax_loss)(reference, tok, nxt) + if adamw: + updated = report.time( + train, host, size, adamw_step, host, step + ) + reference, moments = jax_adamw( + reference, moments, grads, step, lr["lr"].values[step] + ) + else: + updated = report.time( + train, + host, + size, + host.query_all, + step_statements(LEARNING_RATE), + ) + reference = { + name: w - LEARNING_RATE * grads[name] + for name, w in reference.items() + } + errors = [] for name, rows in zip(PARAMS, updated): - reference[name] = ( - reference[name] - LEARNING_RATE * grads[name] - ) - trained = as_array(name, rows) - errors.append( - assert_close( - f"{name}, step {step}", trained, reference[name] + expected = {"w": reference[name]} + if adamw: + expected["m"], expected["v"] = moments[name] + for column, values in expected.items(): + errors.append( + assert_close( + f"{name}.{column}, step {step}", + as_array(name, rows, column), + values, + ) ) + host.load( + name, xql.to_dataset(rows, dims=list(PARAMS[name])) ) - host.load(name, as_table(name, trained)) report.check(max(errors)) report.finish(args.csv) diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 5c286f41..4b6e3902 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -18,7 +18,7 @@ array programs. | # | Script | Model | What it measures | |---|---|---|---| | 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, in einfold's window and `GROUP BY` spellings | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average | -| 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, on Tiny Shakespeare | the loss on each host; SGD training steps, every parameter table checked after every step | +| 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step | | 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature [GraphCast](https://arxiv.org/abs/2212.12794), DeepMind's graph neural network for weather forecasting, over [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5) reanalysis in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | The SQL favors legibility over speed. Each model reads like its definition, @@ -48,7 +48,8 @@ options are the same in all three: --csv results.csv # also write the results table ``` -Demo 2 adds `--steps`, the number of training steps. Demo 3 adds +Demo 2 adds `--steps`, the number of training steps, and `--optimizer`: +`sgd` (the default), or `adamw` to train as nanoGPT does. Demo 3 adds `--rounds`, the number of processor rounds. The output ends with one table: a row per step, host and size, with its @@ -105,7 +106,10 @@ with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the largest value at every size. After each of five SGD steps, all eleven of the GPT's parameter tables match JAX's to 1e-15, and the loss falls from -4.18 to 3.84 (each step on a different batch). +4.18 to 3.84 (each step on a different batch). With AdamW, the weights +and both of Adam's moments match a JAX transcription of +`torch.optim.AdamW` to 6e-15 after each of five steps, and that +transcription matches PyTorch exactly. **DataFusion computes a CTE once for every place that reads it.** A model reads its activations more than once: a layer norm reads its input @@ -165,7 +169,7 @@ SQL gives a number. - Part B of the NanoGPT demo: the four tricks behind a record of the [NanoGPT speedrun](https://github.com/KellerJordan/modded-nanogpt), a public race to train GPT-2 fastest, as a small diff to `02_nanogpt.py`. - It needs nanoGPT's optimizer, AdamW, with its state as versioned tables. + Its sparse optimizer updates build on `--optimizer adamw`. - Peak memory per host, from a fresh process per measurement. - An einfold host, once einfold produces plans. From 272900b042cefaf230d21e435de48e4b104ecdbd Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 21:32:38 -0700 Subject: [PATCH 04/11] Define attention once, as a view, and its loss separately Each spelling of attention is now the model alone: CREATE OR REPLACE VIEW attention AS .... The forward pass selects from the view, and the gradient statement's loss CTE reads it, as a Python loss function calls the model. The plans are the same as before (DataFusion and DuckDB inline the view; operator counts are unchanged), and every check still matches JAX. Hosts gain execute(sql), to run a statement such as CREATE VIEW for its effect. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/01_attention.py | 125 +++++++++++----------------------- benchmarks/ml/README.md | 2 +- benchmarks/ml/_hosts.py | 13 +++- 3 files changed, 54 insertions(+), 86 deletions(-) diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index 16cb60c0..f9f57bce 100644 --- a/benchmarks/ml/01_attention.py +++ b/benchmarks/ml/01_attention.py @@ -29,11 +29,11 @@ (``docs/demos.md`` in its repository). This script is the baseline that plan must match and then beat: -* **forward**: each spelling on each *host* (einfold's word for an engine +* **forward**: each spelling of the model on each *host* (einfold's word for an engine that runs a plan; here DataFusion and DuckDB), checked against the same attention in JAX, Google's library for differentiable array programs; -* **gradient**: the gradient of a loss with respect to every value of Q, - K and V, as one SQL statement. ddx (https://github.com/xqlsystems/ddx), +* **gradient**: a loss that reads the model, and its gradient with + respect to every value of Q, K and V, as one SQL statement. ddx (https://github.com/xqlsystems/ddx), which differentiates SQL queries, computes it, and ``jax.grad`` checks it; * **masks are absent rows**: an infinity in V reaches only the queries @@ -59,6 +59,11 @@ D = 64 # features per head; the scores are scaled by 1/√64 = 1/8. +# The model and its loss are separate, as in a neural network written in +# Python (compare attention() and loss() below). The model is a view, a +# named query: each spelling of attention defines the view ``attention``, +# and the forward pass and the loss read it as if it were a table. + # Attention in einfold's window spelling. The causal mask is a join # predicate: a masked (i, j) pair is a row that does not exist. Subtracting # each row's max only keeps exp() from overflowing: it cancels out of the @@ -66,6 +71,7 @@ # TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX # once ddx can differentiate a window aggregate; it refuses one today. ATTENTION_WINDOW = """ +CREATE OR REPLACE VIEW attention AS WITH scores AS ( SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s FROM q JOIN k ON q.d = k.d @@ -75,53 +81,19 @@ SELECT i, j, exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i))) AS w FROM scores), norm AS ( - SELECT i, SUM(w) AS z FROM weights GROUP BY i), -attention AS ( - SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x - FROM weights w - JOIN v ON w.j = v.t - JOIN norm n ON w.i = n.i - GROUP BY w.i, v.d) -SELECT t, d, x FROM attention ORDER BY t, d -""" - -# The gradient of half the output's squared norm, with respect to every -# value of Q, K and V: the same query, plus a loss and a different SELECT. -# The three grad() calls differentiate one loss, so ddx runs one backward -# pass for all three, as it would for three statements given to sql_all. -ATTENTION_WINDOW_GRADIENT = """ -WITH scores AS ( - SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s - FROM q JOIN k ON q.d = k.d - WHERE q.t >= k.t - GROUP BY q.t, k.t), -weights AS ( - SELECT i, j, exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i))) AS w - FROM scores), -norm AS ( - SELECT i, SUM(w) AS z FROM weights GROUP BY i), -attention AS ( - SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x - FROM weights w - JOIN v ON w.j = v.t - JOIN norm n ON w.i = n.i - GROUP BY w.i, v.d), -loss AS ( - SELECT SUM(x * x) / 2 AS l FROM attention) -SELECT q.t, q.d, - grad(loss, q.x) AS dq, - grad(loss, k.x) AS dk, - grad(loss, v.x) AS dv -FROM q -JOIN k ON q.t = k.t AND q.d = k.d -JOIN v ON q.t = v.t AND q.d = v.d -ORDER BY q.t, q.d + SELECT i, SUM(w) AS z FROM weights GROUP BY i) +SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x +FROM weights w +JOIN v ON w.j = v.t +JOIN norm n ON w.i = n.i +GROUP BY w.i, v.d """ # The same attention as ddx spells it: each row's max as a GROUP BY, # joined back to the scores. Here ddx_stop_gradient isn't needed, but it # is exact, and it spares ddx differentiating the max. ATTENTION_GROUP_BY = """ +CREATE OR REPLACE VIEW attention AS WITH scores AS ( SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s FROM q JOIN k ON q.d = k.d @@ -133,36 +105,25 @@ SELECT s.i, s.j, exp(s.s - r.m) AS w FROM scores s JOIN row_max r ON s.i = r.i), norm AS ( - SELECT i, SUM(w) AS z FROM weights GROUP BY i), -attention AS ( - SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x - FROM weights w - JOIN v ON w.j = v.t - JOIN norm n ON w.i = n.i - GROUP BY w.i, v.d) -SELECT t, d, x FROM attention ORDER BY t, d + SELECT i, SUM(w) AS z FROM weights GROUP BY i) +SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x +FROM weights w +JOIN v ON w.j = v.t +JOIN norm n ON w.i = n.i +GROUP BY w.i, v.d """ -ATTENTION_GROUP_BY_GRADIENT = """ -WITH scores AS ( - SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s - FROM q JOIN k ON q.d = k.d - WHERE q.t >= k.t - GROUP BY q.t, k.t), -row_max AS ( - SELECT i, ddx_stop_gradient(MAX(s)) AS m FROM scores GROUP BY i), -weights AS ( - SELECT s.i, s.j, exp(s.s - r.m) AS w - FROM scores s JOIN row_max r ON s.i = r.i), -norm AS ( - SELECT i, SUM(w) AS z FROM weights GROUP BY i), -attention AS ( - SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x - FROM weights w - JOIN v ON w.j = v.t - JOIN norm n ON w.i = n.i - GROUP BY w.i, v.d), -loss AS ( +SPELLINGS = {"window": ATTENTION_WINDOW, "group by": ATTENTION_GROUP_BY} + +# The forward pass: run the model. +FORWARD = "SELECT t, d, x FROM attention ORDER BY t, d" + +# The loss, half the output's squared norm, calls the model by reading the +# view. The SELECT asks for its gradient with respect to every value of Q, +# K and V. The three grad() calls differentiate one loss, so ddx runs one +# backward pass for all three, as jax.grad(loss, argnums=(0, 1, 2)) does. +GRADIENT = """ +WITH loss AS ( SELECT SUM(x * x) / 2 AS l FROM attention) SELECT q.t, q.d, grad(loss, q.x) AS dq, @@ -174,12 +135,6 @@ ORDER BY q.t, q.d """ -# Each spelling's forward query and gradient query. -SPELLINGS = { - "window": (ATTENTION_WINDOW, ATTENTION_WINDOW_GRADIENT), - "group by": (ATTENTION_GROUP_BY, ATTENTION_GROUP_BY_GRADIENT), -} - # Nadaraya–Watson: at each target point, the average of the sources' y # weighted by a Gaussian kernel of their distance (bandwidth h = 0.5). # Structurally attention: a weighted average with weights exp(score). @@ -236,18 +191,19 @@ def main(): for name, data in zip("qkv", (q, k, v)): host.load(name, table(data)) - for spelling, (forward, gradient) in SPELLINGS.items(): + for spelling, model in SPELLINGS.items(): + host.execute(model) rows = report.time( - f"forward ({spelling})", host, n, host.query, forward + f"forward ({spelling})", host, n, host.query, FORWARD ) report.check(assert_close("forward", matrix(rows), want)) - report.note("operators", host.operators(forward)) + report.note("operators", host.operators(FORWARD)) step = f"gradient ({spelling})" if not host.differentiates: report.skip(step, host, n, "no ddx backend") continue - rows = report.time(step, host, n, host.query, gradient) + rows = report.time(step, host, n, host.query, GRADIENT) errors = [ assert_close(f"d/d{name}", matrix(rows, f"d{name}"), ref) for name, ref in zip("qkv", want_grads) @@ -272,8 +228,9 @@ def check_masks_are_absent_rows(host_names): for host in hosts(host_names): for name, data in zip("qkv", (q, k, v)): host.load(name, table(data)) - for spelling, (forward, _) in SPELLINGS.items(): - out = matrix(host.query(forward)) + for spelling, model in SPELLINGS.items(): + host.execute(model) + out = matrix(host.query(FORWARD)) assert np.isfinite(out[:j]).all(), (host.name, spelling) assert np.isposinf(out[j:, 0]).all(), (host.name, spelling) assert np.isfinite(out[:, 1:]).all(), (host.name, spelling) diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 4b6e3902..306ba65b 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -17,7 +17,7 @@ array programs. | # | Script | Model | What it measures | |---|---|---|---| -| 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, in einfold's window and `GROUP BY` spellings | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average | +| 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, as a view (the model) in einfold's window and `GROUP BY` spellings, and a loss that reads it | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average | | 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step | | 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature [GraphCast](https://arxiv.org/abs/2212.12794), DeepMind's graph neural network for weather forecasting, over [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5) reanalysis in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index 8784a0e2..54a246a1 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -5,11 +5,13 @@ plan its *host*. Every demo writes its model once, as SQL, and runs that same text on each host, so a difference in time is a difference between engines, never between -queries. A host does four things: +queries. A host does five things: * ``load(name, data, **options)`` registers an ``xr.Dataset`` as the table ``name`` through xarray-sql, replacing any table of that name (``options`` such as ``table_names`` go to xarray-sql); +* ``execute(sql)`` runs a statement for its effect, such as + ``CREATE VIEW``; * ``query(sql)`` runs a statement and returns its rows as Arrow; * ``query_all(sqls)`` runs several statements, and, on a host that differentiates, shares the backward pass between those that take @@ -53,6 +55,9 @@ class Host: def load(self, name: str, data: xr.Dataset, **options) -> None: raise NotImplementedError + def execute(self, sql: str) -> None: + raise NotImplementedError + def query(self, sql: str) -> pa.Table: raise NotImplementedError @@ -76,6 +81,9 @@ def load(self, name, data, **options): self.ctx.deregister_table(name) self.ctx.from_dataset(name, data, **options) + def execute(self, sql): + self.ctx.sql(sql).collect() + def query(self, sql): return self.ctx.sql(sql).to_arrow_table() @@ -99,6 +107,9 @@ def __init__(self): def load(self, name, data, **options): xql.register(self.con, name, data, **options) + def execute(self, sql): + self.con.execute(sql) + def query(self, sql): return self.con.sql(sql).fetch_arrow_table() From b1b25036259fe2d696c013f5f7f4a772b3b7db4b Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 21:46:06 -0700 Subject: [PATCH 05/11] Split the GPT into a model view and a separate loss The model is the view logits; the cross-entropy loss is a LOSS of CTEs that reads it, as nanoGPT's F.cross_entropy reads the model's logits. The loss reads the view once, into a CTE, because DuckDB expands a view at every read but computes a CTE read twice only once (232 operators and 0.097 s reading it twice, 126 and 0.06 s reading it once). A DataFusion view keeps reading the tables that existed when it was created (#271), so each step recreates the view after registering the batch and the parameters. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/02_nanogpt.py | 53 ++++++++++++++++++++++++++----------- benchmarks/ml/README.md | 2 +- 2 files changed, 39 insertions(+), 16 deletions(-) diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index bdac788d..45947a56 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -33,6 +33,10 @@ the learning rate schedule as a table ================= ============================================ +As in nanoGPT, the model and the loss are separate. The model is the +view ``logits``; the loss is SQL that reads it, as nanoGPT's +``F.cross_entropy`` reads the model's logits. + Each training step updates every parameter table with its own statement; the statements share one backward pass. By default the optimizer is SGD at a fixed learning rate, the simplest step to read. @@ -226,28 +230,46 @@ def block(layer: int, x: str, out: str) -> str: def model() -> str: - """The CTEs of the whole model, ending in ``loss``: the mean - cross-entropy of predicting each next character of the batch.""" + """The model, as the view ``logits``: for each position of each + sequence in ``batch``, a score for every character that could come + next. A DataFusion view reads the tables that existed when it was + created (DuckDB's looks them up by name), so it is recreated whenever + the batch or the parameters are.""" blocks = "".join( block(layer, f"x{layer}", f"x{layer + 1}") for layer in range(LAYERS) ) return f""" -x0 AS ( +CREATE OR REPLACE VIEW logits AS +WITH x0 AS ( SELECT b.seq, b.pos, te.e, te.w + pe.w AS val FROM batch b JOIN wte te ON b.tok = te.tok JOIN wpe pe ON b.pos = pe.pos AND te.e = pe.e), {blocks} -{layer_norm(f"x{LAYERS}", "final", "ln_f")} -logits AS ( - SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val - FROM final x JOIN wte te ON x.e = te.e - GROUP BY x.seq, x.pos, te.tok), +{layer_norm(f"x{LAYERS}", "final", "ln_f").rstrip(",")} +SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val +FROM final x JOIN wte te ON x.e = te.e +GROUP BY x.seq, x.pos, te.tok +""" + + +# The loss, separate from the model, as in nanoGPT: the mean cross-entropy +# of the model's logits against each next character of the batch. +# TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX +# once ddx can differentiate a window aggregate. The max only keeps exp() +# from overflowing, and cancels out of the log-sum-exp, so its gradient is +# zero. +# +# The loss reads the model's view once, into the CTE ``predicted``, and +# uses that twice: a host expands a view at every read, but DuckDB +# computes a CTE that is read twice only once. +LOSS = """ +predicted AS ( + SELECT seq, pos, tok, val FROM logits), shifted AS ( - -- TODO(xqlsystems/ddx#142): as in the attention weights above. SELECT seq, pos, tok, val, ddx_stop_gradient(MAX(val) OVER (PARTITION BY seq, pos)) AS m - FROM logits), + FROM predicted), log_sum_exp AS ( SELECT seq, pos, MAX(m) + ln(SUM(exp(val - m))) AS lse FROM shifted GROUP BY seq, pos), @@ -255,13 +277,13 @@ def model() -> str: SELECT AVG(s.lse - l.val) AS l FROM log_sum_exp s JOIN batch b ON s.seq = b.seq AND s.pos = b.pos - JOIN logits l ON l.seq = b.seq AND l.pos = b.pos AND l.tok = b.next)""" + JOIN predicted l ON l.seq = b.seq AND l.pos = b.pos AND l.tok = b.next)""" def step_statements(lr: float) -> list[str]: """One SGD step: each parameter table's next version, as a statement.""" return [ - f"WITH {model()} SELECT {', '.join(dims)}, " + f"WITH {LOSS} SELECT {', '.join(dims)}, " f"w - {lr} * grad(loss, w) AS w FROM {table} " f"ORDER BY {', '.join(dims)}" for table, dims in PARAMS.items() @@ -273,7 +295,7 @@ def gradient_statements() -> list[str]: the table ``{name}_grad``'s rows. The statements share one backward pass.""" return [ - f"WITH {model()} SELECT {', '.join(dims)}, grad(loss, w) AS g " + f"WITH {LOSS} SELECT {', '.join(dims)}, grad(loss, w) AS g " f"FROM {table} ORDER BY {', '.join(dims)}" for table, dims in PARAMS.items() ] @@ -316,7 +338,7 @@ def adamw_statement(table: str, step: int) -> str: def loss_statement() -> str: - return f"WITH {model()} SELECT l FROM loss" + return f"WITH {LOSS} SELECT l FROM loss" # -------------------------------------------------------------------------- @@ -492,7 +514,7 @@ def main(): report = Report("nanogpt", "batch", args.reps) text = shakespeare() vocab = len(set(text)) - lines = len(loss_statement().splitlines()) + lines = len(model().splitlines()) + len(loss_statement().splitlines()) count = sum( v.size for v in init_params(vocab, np.random.default_rng(0)).values() ) @@ -521,6 +543,7 @@ def main(): for step, batch in data: host.load("batch", batch) + host.execute(model()) tok, nxt = batch["tok"].values, batch["next"].values rows = report.time( diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 306ba65b..c6ccaa6d 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -18,7 +18,7 @@ array programs. | # | Script | Model | What it measures | |---|---|---|---| | 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, as a view (the model) in einfold's window and `GROUP BY` spellings, and a loss that reads it | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average | -| 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step | +| 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, as a view, and a cross-entropy loss that reads it, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step | | 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature [GraphCast](https://arxiv.org/abs/2212.12794), DeepMind's graph neural network for weather forecasting, over [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5) reanalysis in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | The SQL favors legibility over speed. Each model reads like its definition, From a913421a4c6c934acc4b0a2314d0e852b8b62bc8 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 22:02:43 -0700 Subject: [PATCH 06/11] Fix lint: format the GraphCast script, and type two values for mypy Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/03_graphcast.py | 6 ++---- xarray_sql/ad.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py index 8b0232c7..c2704518 100644 --- a/benchmarks/ml/03_graphcast.py +++ b/benchmarks/ml/03_graphcast.py @@ -508,7 +508,7 @@ def open_era5() -> xr.Dataset: rng = np.random.default_rng(3) surface = rng.normal(size=(len(TIMES), 64, 32)) atmosphere = rng.normal(size=(len(TIMES), 2, 64, 32)) - data = { + data: dict[str, tuple[tuple[str, ...], np.ndarray]] = { name: ( ("time", "longitude", "latitude"), surface + i, @@ -688,9 +688,7 @@ def main(): ] ) regional = rows["change"].to_numpy().reshape(-1, len(VARIABLES)) - report.check( - assert_close("regional", regional, want[in_region]) - ) + report.check(assert_close("regional", regional, want[in_region])) report.note("operators", host.operators(statement)) step = "gradient of the error" diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py index 9d692293..edf4b09e 100644 --- a/xarray_sql/ad.py +++ b/xarray_sql/ad.py @@ -128,6 +128,6 @@ def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list: """ if isinstance(statements, str): raise TypeError("statements is a list of SQL statements") - frames = _ddx.sql_all(ctx, list(statements)) + frames: list = _ddx.sql_all(ctx, list(statements)) wrap = getattr(ctx, "_wrap", None) return [wrap(f) for f in frames] if wrap else frames From 04d74a0782ec4dcf90b00572ba90461c66cc60fc Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 18 Oct 2026 08:15:55 -0700 Subject: [PATCH 07/11] Differentiate the softmax's window MAX directly with ddxdb 0.3.1 ddxdb 0.3.1 differentiates window aggregates (xqlsystems/ddx#142), so attention and the GPT's cross-entropy no longer wrap the row max in ddx_stop_gradient. Gradients still match JAX (1e-15), and a GPT training step takes the same time. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/01_attention.py | 12 +++++------- benchmarks/ml/02_nanogpt.py | 14 +++++--------- benchmarks/ml/README.md | 6 +----- benchmarks/ml/_hosts.py | 5 ----- pyproject.toml | 2 +- 5 files changed, 12 insertions(+), 27 deletions(-) diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index f9f57bce..ecd59dfb 100644 --- a/benchmarks/ml/01_attention.py +++ b/benchmarks/ml/01_attention.py @@ -67,9 +67,8 @@ # Attention in einfold's window spelling. The causal mask is a join # predicate: a masked (i, j) pair is a row that does not exist. Subtracting # each row's max only keeps exp() from overflowing: it cancels out of the -# ratio w / z, so it carries no gradient, which ddx_stop_gradient says. -# TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX -# once ddx can differentiate a window aggregate; it refuses one today. +# ratio w / z, so its gradient sums to zero, and ddx differentiates it +# like any other window. ATTENTION_WINDOW = """ CREATE OR REPLACE VIEW attention AS WITH scores AS ( @@ -78,7 +77,7 @@ WHERE q.t >= k.t GROUP BY q.t, k.t), weights AS ( - SELECT i, j, exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY i))) AS w + SELECT i, j, exp(s - MAX(s) OVER (PARTITION BY i)) AS w FROM scores), norm AS ( SELECT i, SUM(w) AS z FROM weights GROUP BY i) @@ -90,8 +89,7 @@ """ # The same attention as ddx spells it: each row's max as a GROUP BY, -# joined back to the scores. Here ddx_stop_gradient isn't needed, but it -# is exact, and it spares ddx differentiating the max. +# joined back to the scores. ATTENTION_GROUP_BY = """ CREATE OR REPLACE VIEW attention AS WITH scores AS ( @@ -100,7 +98,7 @@ WHERE q.t >= k.t GROUP BY q.t, k.t), row_max AS ( - SELECT i, ddx_stop_gradient(MAX(s)) AS m FROM scores GROUP BY i), + SELECT i, MAX(s) AS m FROM scores GROUP BY i), weights AS ( SELECT s.i, s.j, exp(s.s - r.m) AS w FROM scores s JOIN row_max r ON s.i = r.i), diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index 45947a56..ae670d29 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -186,11 +186,9 @@ def block(layer: int, x: str, out: str) -> str: WHERE q.pos >= k.pos GROUP BY q.seq, q.h, q.pos, k.pos), {l}_weights AS ( - -- TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX - -- once ddx can differentiate a window aggregate. The max only keeps exp() - -- from overflowing, and cancels out of p / z, so its gradient is zero. + -- The max only keeps exp() from overflowing: it cancels out of p / z. SELECT seq, h, i, j, - exp(s - ddx_stop_gradient(MAX(s) OVER (PARTITION BY seq, h, i))) AS p + exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p FROM {l}_scores), {l}_norm AS ( SELECT seq, h, i, SUM(p) AS z FROM {l}_weights GROUP BY seq, h, i), @@ -255,10 +253,8 @@ def model() -> str: # The loss, separate from the model, as in nanoGPT: the mean cross-entropy # of the model's logits against each next character of the batch. -# TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX -# once ddx can differentiate a window aggregate. The max only keeps exp() -# from overflowing, and cancels out of the log-sum-exp, so its gradient is -# zero. +# Subtracting each row's max only keeps exp() from overflowing: it cancels +# out of the log-sum-exp. # # The loss reads the model's view once, into the CTE ``predicted``, and # uses that twice: a host expands a view at every read, but DuckDB @@ -268,7 +264,7 @@ def model() -> str: SELECT seq, pos, tok, val FROM logits), shifted AS ( SELECT seq, pos, tok, val, - ddx_stop_gradient(MAX(val) OVER (PARTITION BY seq, pos)) AS m + MAX(val) OVER (PARTITION BY seq, pos) AS m FROM predicted), log_sum_exp AS ( SELECT seq, pos, MAX(m) + ln(SUM(exp(val - m))) AS lse diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index c6ccaa6d..4980c1db 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -100,7 +100,7 @@ einfold's optimizer or DuckDB with a GPU extension, write one class in ## Findings Measured on 2026-10-17 on an Intel Core i7-8700 (6 cores, 12 threads) -with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. +with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.1, one repetition each. **ddx's gradients agree with JAX.** The gradients of attention with respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the @@ -147,10 +147,6 @@ region's filter back through the graph is einfold's Demo 3. **What ddx cannot differentiate yet.** Each of these raised a typed error naming the construct, never a wrong number: -- A window `MAX(s) OVER (...)` used as a value. Attention and the - cross-entropy subtract the row maximum before `exp`, and that shift - cancels out of the result, so the demos wrap it in `ddx_stop_gradient`, - which is exact. ddx could treat this as a rule. - A `GROUP BY` key that depends on a differentiated column. einfold's window spelling groups by the normalizer `z`; the demos divide inside the sum instead, `SUM(w / z * v)`, which einfold also lists. diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index 54a246a1..c52fb71b 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -25,10 +25,6 @@ cannot differentiate says so with ``differentiates = False``, and the demos report its gradient steps as not run rather than skip it silently. -On every host, ``ddx_stop_gradient(x)`` is the identity. It is how a -model marks a value whose derivative is zero (the max subtracted before a -softmax's ``exp``), and it changes nothing in a forward pass. - Adding a host (einfold's optimizer on DataFusion, DuckDB with a GPU extension, or DuckDB running ddx's programs once ddx supports it) is one class here and one entry in ``HOSTS``. @@ -102,7 +98,6 @@ class DuckDB(Host): def __init__(self): self.con = duckdb.connect() - self.con.sql("CREATE MACRO ddx_stop_gradient(x) AS x") def load(self, name, data, **options): xql.register(self.con, name, data, **options) diff --git a/pyproject.toml b/pyproject.toml index f9384fb2..9c2bc3b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,7 @@ geo = [ # and breaks its API between minor versions. It reads and writes the # Substrait of DataFusion 54, so bump it together with that pin. ml = [ - "ddxdb>=0.3.0,<0.4; sys_platform != 'emscripten'", + "ddxdb>=0.3.1,<0.4; sys_platform != 'emscripten'", ] test = [ "adbc-driver-flightsql>=1.12", From f2a7622eacd128ec6e0038c98fea8d42e9c3ec9d Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 18 Oct 2026 13:35:54 -0700 Subject: [PATCH 08/11] Require ddxdb 0.3.2 and record its faster GPT steps ddxdb 0.3.2 no longer plans a gradient statement's unused objective CTEs (xqlsystems/ddx#160), and differentiates two-operand products without rebuilding the forward join (xqlsystems/ddx#126). On DataFusion, a GPT training step goes from 3.0 s to 2.7 s with SGD and from 3.4 s to 3.1 s with AdamW (median of three). Attention's gradient is unchanged: its products are under a causal mask, which the new rule doesn't cover yet (xqlsystems/ddx#153). Every check still passes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/README.md | 12 ++++++++---- pyproject.toml | 2 +- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 4980c1db..0fa98ad5 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -100,7 +100,9 @@ einfold's optimizer or DuckDB with a GPU extension, write one class in ## Findings Measured on 2026-10-17 on an Intel Core i7-8700 (6 cores, 12 threads) -with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.1, one repetition each. +with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. +The demos' gradients were timed again on 2026-10-18 with ddxdb 0.3.2, +taking the median of three repetitions. **ddx's gradients agree with JAX.** The gradients of attention with respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the @@ -124,8 +126,9 @@ read more than once a single time. | GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | | GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | -A GPT training step, forward and backward through ddx, takes 3.3 s on -DataFusion: less than its forward query alone, because ddx stores the +A GPT training step, forward and backward through ddx, takes 2.7 s on +DataFusion with SGD and 3.1 s with AdamW (3.0 s and 3.4 s with ddxdb +0.3.1, which planned the loss again for every parameter table): less than its forward query alone, because ddx stores the result of each step of its program as a table. Computing shared subexpressions once is the first thing xarray-sql's DataFusion path needs for these models, and einfold should plan them as graphs rather than @@ -134,7 +137,8 @@ trees. **Attention stores every score.** Causal attention's forward pass at 512 tokens takes 0.68 s on DataFusion in the window spelling and 1.3 s in the `GROUP BY` spelling; DuckDB takes 0.52 s and 0.26 s. Both grow with the -square of the tokens. The gradients take 3 to 4 times the forward pass. +square of the tokens. The gradients take 2 to 4 times the forward pass: 2.7 s and 2.6 s on +DataFusion at 512 tokens, the same with ddxdb 0.3.1 and 0.3.2. **A regional forecast costs what a global one does.** For the contiguous United States, the forecast covers 55 of 2,048 grid cells. Walking the diff --git a/pyproject.toml b/pyproject.toml index 9c2bc3b4..cbea0304 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,7 @@ geo = [ # and breaks its API between minor versions. It reads and writes the # Substrait of DataFusion 54, so bump it together with that pin. ml = [ - "ddxdb>=0.3.1,<0.4; sys_platform != 'emscripten'", + "ddxdb>=0.3.2,<0.4; sys_platform != 'emscripten'", ] test = [ "adbc-driver-flightsql>=1.12", From 84d1e733d9ab1d6bfb61dc8daf3837198a872a85 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 20 Oct 2026 17:23:08 -0700 Subject: [PATCH 09/11] Require ddxdb 0.3.3 and record its faster gradients ddxdb 0.3.3 keeps each gradient program between calls (xqlsystems/ddx#164), plans each CTE an objective reads twice on its own (xqlsystems/ddx#168), and differentiates attention's masked product as jax.grad does (xqlsystems/ddx#162). Alternating runs with 0.3.2 on DataFusion: - GPT training step, SGD: 2.8 s -> 1.2 s (first step 2.0 s) - GPT training step, AdamW: 3.2 s -> 1.55 s - attention gradient at 512 tokens: 2.8 s -> 1.1 s (window), 2.7 s -> 1.3 s (GROUP BY) Every check against JAX still passes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/README.md | 23 +++++++++++++---------- pyproject.toml | 2 +- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 0fa98ad5..73632dfc 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -101,8 +101,8 @@ einfold's optimizer or DuckDB with a GPU extension, write one class in Measured on 2026-10-17 on an Intel Core i7-8700 (6 cores, 12 threads) with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. -The demos' gradients were timed again on 2026-10-18 with ddxdb 0.3.2, -taking the median of three repetitions. +The demos' gradients were timed again on 2026-10-20 with ddxdb 0.3.3, +alternating runs with ddxdb 0.3.2 on the same machine. **ddx's gradients agree with JAX.** The gradients of attention with respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the @@ -126,19 +126,22 @@ read more than once a single time. | GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | | GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | -A GPT training step, forward and backward through ddx, takes 2.7 s on -DataFusion with SGD and 3.1 s with AdamW (3.0 s and 3.4 s with ddxdb -0.3.1, which planned the loss again for every parameter table): less than its forward query alone, because ddx stores the -result of each step of its program as a table. Computing shared -subexpressions once is the first thing xarray-sql's DataFusion path needs -for these models, and einfold should plan them as graphs rather than +A GPT training step, forward and backward through ddx, takes 1.2 s on +DataFusion with SGD and 1.55 s with AdamW, an eighth of its forward query +alone (2.8 s and 3.2 s with ddxdb 0.3.2). ddx builds the gradient program +on the first step, which takes 2.0 s, and runs it again on later ones. It +plans each CTE the loss reads twice on its own, and stores the result of +each step of its program as a table, so it computes each shared +subexpression once. Doing the same for forward queries is the first +thing xarray-sql's DataFusion path needs for these models, and einfold should plan them as graphs rather than trees. **Attention stores every score.** Causal attention's forward pass at 512 tokens takes 0.68 s on DataFusion in the window spelling and 1.3 s in the `GROUP BY` spelling; DuckDB takes 0.52 s and 0.26 s. Both grow with the -square of the tokens. The gradients take 2 to 4 times the forward pass: 2.7 s and 2.6 s on -DataFusion at 512 tokens, the same with ddxdb 0.3.1 and 0.3.2. +square of the tokens. The gradients take 1 to 2 times the forward pass: 1.1 s and 1.3 s on +DataFusion at 512 tokens, against 2.8 s and 2.7 s with ddxdb 0.3.2, +since ddx differentiates the masked product as `jax.grad` does. **A regional forecast costs what a global one does.** For the contiguous United States, the forecast covers 55 of 2,048 grid cells. Walking the diff --git a/pyproject.toml b/pyproject.toml index cbea0304..ec9a6457 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,7 @@ geo = [ # and breaks its API between minor versions. It reads and writes the # Substrait of DataFusion 54, so bump it together with that pin. ml = [ - "ddxdb>=0.3.2,<0.4; sys_platform != 'emscripten'", + "ddxdb>=0.3.3,<0.4; sys_platform != 'emscripten'", ] test = [ "adbc-driver-flightsql>=1.12", From 2fb20a68bc7cdd8b127bdb51abc50d060dffa2e6 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 21 Oct 2026 02:14:36 -0700 Subject: [PATCH 10/11] Require ddxdb 0.4.0, on DataFusion 55 ddxdb 0.4.0 exchanges Substrait with datafusion 55.0.0, the version xarray-sql now pins. Its gradients, programs and API are unchanged from 0.3.3. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- pyproject.toml | 4 ++-- uv.lock | 22 ++++++++++++++++++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ec9a6457..77bb1b0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,9 +60,9 @@ geo = [ ] # Gradients in SQL: grad(), vjp() and jvp() from ddx. ddxdb is pre-1.0 # and breaks its API between minor versions. It reads and writes the -# Substrait of DataFusion 54, so bump it together with that pin. +# Substrait of DataFusion 55, so bump it together with that pin. ml = [ - "ddxdb>=0.3.3,<0.4; sys_platform != 'emscripten'", + "ddxdb>=0.4.0,<0.5; sys_platform != 'emscripten'", ] test = [ "adbc-driver-flightsql>=1.12", diff --git a/uv.lock b/uv.lock index f98b1b8c..879bc3e7 100644 --- a/uv.lock +++ b/uv.lock @@ -520,6 +520,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4b/d0/bdc7ea7072c3925f3977e13e1e754e145fa42eb40bb5c1abff8028e0f791/datafusion-55.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b7191ff60defee8d9d9e5ff6aeefdd961bc200897cd60440b02d2a0f316902ef", size = 45733022, upload-time = "2026-10-12T16:20:42.607Z" }, ] +[[package]] +name = "ddxdb" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/9b/e545bc8537ec63c8c41d0642e52798956eda3393feaeccb0406a3a5950c7/ddxdb-0.4.0.tar.gz", hash = "sha256:0c0d4a14b82af02d6b83b9ed243014c64401a167f123c8a84d9fc19429d5f1ed", size = 291824, upload-time = "2026-10-21T09:04:11.42Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/7e/acbeeca83355880e4404a7b15d765ed82d85efa4a79952941c9401931885/ddxdb-0.4.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:be538a452d2c45577fb060a2be3134631ea68658caa9c3a2a854b5b13a0f21d9", size = 6725712, upload-time = "2026-10-21T09:04:00.956Z" }, + { url = "https://files.pythonhosted.org/packages/23/b9/0fc5bbb27bd323533089c6372adc80f79875679f7741a9d4e738eabff1b2/ddxdb-0.4.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:8c4a752f27db4219ad0bd8d6b301e88bea192bc3753b54bc909dd81bbf562781", size = 6644600, upload-time = "2026-10-21T09:04:03.238Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8a/b76f44799346a5d6b9f88c1dde5d2aae19934ce31b6c992b27629eee8709/ddxdb-0.4.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0636c7d4d52e64e0d3bc71dcf1bb3a15d1e1e494c553e6b184da192df147be69", size = 6739521, upload-time = "2026-10-21T09:04:05.205Z" }, + { url = "https://files.pythonhosted.org/packages/f5/cc/2b80f26e102beb2c655dc2b60bf2e427b0d9e0e588c723343c8ea300dac4/ddxdb-0.4.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b2f745b7cd057f0a7fc550f0f5984a836683a6e3252a677fc18e6f6045dbb80d", size = 7179153, upload-time = "2026-10-21T09:04:07.276Z" }, + { url = "https://files.pythonhosted.org/packages/d6/dc/0231a760cbea1a86a8c71d0526ec9ff45d7850573ea73aacddcf3d9916b6/ddxdb-0.4.0-cp310-abi3-win_amd64.whl", hash = "sha256:f38c950fdc8ec1a4c347d80715c6b3c04499922f2e71dac004ef51b2adc71101", size = 7133649, upload-time = "2026-10-21T09:04:09.43Z" }, +] + [[package]] name = "decorator" version = "5.2.1" @@ -2865,6 +2878,9 @@ geo = [ { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] +ml = [ + { name = "ddxdb", marker = "sys_platform != 'emscripten'" }, +] polars = [ { name = "polars" }, ] @@ -2874,6 +2890,7 @@ test = [ { name = "adbc-driver-sqlite" }, { name = "cftime" }, { name = "dask" }, + { name = "ddxdb", marker = "sys_platform != 'emscripten'" }, { name = "duckdb" }, { name = "gcsfs" }, { name = "polars" }, @@ -2900,6 +2917,7 @@ requires-dist = [ { name = "cftime", marker = "extra == 'test'" }, { name = "dask", marker = "extra == 'dask'", specifier = ">=2024.8.0" }, { name = "datafusion", marker = "sys_platform != 'emscripten'", specifier = "==55.0.0" }, + { name = "ddxdb", marker = "sys_platform != 'emscripten' and extra == 'ml'", specifier = ">=0.4.0,<0.5" }, { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.4.0" }, { name = "gcsfs", marker = "extra == 'test'" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" }, @@ -2912,11 +2930,11 @@ requires-dist = [ { name = "watchfiles", marker = "extra == 'dev'" }, { name = "xarray", specifier = ">=2024.7.0" }, { name = "xarray", extras = ["io"], marker = "extra == 'test'" }, - { name = "xarray-sql", extras = ["adbc", "dask", "duckdb", "polars", "geo"], marker = "extra == 'test'" }, + { name = "xarray-sql", extras = ["adbc", "dask", "duckdb", "polars", "geo", "ml"], marker = "extra == 'test'" }, { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" }, { name = "zensical", marker = "extra == 'docs'" }, ] -provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "polars", "test"] +provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "ml", "polars", "test"] [package.metadata.requires-dev] dev = [ From 92dcde04fbaa279082c594a1ead3d90b798f9fd8 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 21 Oct 2026 02:16:58 -0700 Subject: [PATCH 11/11] Record the demos on ddxdb 0.4.0 and in-memory tables On DataFusion 55 with ddxdb 0.4.0 and xarray-sql's in-memory tables, the GPT's forward loss takes 3.0 s (9.4 s), an SGD step 0.75 s (1.2 s) and an AdamW step 0.85 s (1.55 s). Attention's forward pass at 512 tokens takes 0.28 s and 0.37 s, and its gradients 0.92 s and 0.93 s. Every result still matches JAX. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/README.md | 40 ++++++++++++++++++++++++---------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index 73632dfc..40a92a6a 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -101,8 +101,9 @@ einfold's optimizer or DuckDB with a GPU extension, write one class in Measured on 2026-10-17 on an Intel Core i7-8700 (6 cores, 12 threads) with DataFusion 54.0.0, DuckDB 1.5.5 and ddxdb 0.3.0, one repetition each. -The demos' gradients were timed again on 2026-10-20 with ddxdb 0.3.3, -alternating runs with ddxdb 0.3.2 on the same machine. +The GPT and attention demos were timed again on 2026-10-21 with +DataFusion 55.0.0 and ddxdb 0.4.0, after xarray-sql began registering +small in-memory Datasets as DataFusion in-memory tables. **ddx's gradients agree with JAX.** The gradients of attention with respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the @@ -122,26 +123,33 @@ read more than once a single time. | Query | DataFusion operators | DuckDB operators | DataFusion | DuckDB | |---|---:|---:|---:|---:| -| GPT loss, 1 layer, batch 4 | 12,794 | 126 | 9.4 s | 0.058 s | +| GPT loss, 1 layer, batch 4 | 13,756 | 126 | 3.0 s | 0.057 s | | GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | | GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | -A GPT training step, forward and backward through ddx, takes 1.2 s on -DataFusion with SGD and 1.55 s with AdamW, an eighth of its forward query -alone (2.8 s and 3.2 s with ddxdb 0.3.2). ddx builds the gradient program -on the first step, which takes 2.0 s, and runs it again on later ones. It -plans each CTE the loss reads twice on its own, and stores the result of -each step of its program as a table, so it computes each shared -subexpression once. Doing the same for forward queries is the first -thing xarray-sql's DataFusion path needs for these models, and einfold should plan them as graphs rather than +The GraphCast rows are from 2026-10-17, before in-memory tables. + +Setting up each lazy read of a Dataset takes xarray-sql about 1.5 ms, so +the GPT's 2,989 reads in one forward pass took 6.7 s of its 9.4 s. A small +in-memory Dataset is now an in-memory table, which takes a quarter of a +millisecond per read. A GPT training step, forward and backward through +ddx, takes 0.75 s on DataFusion with SGD and 0.85 s with AdamW, a quarter +of its forward query alone (1.2 s and 1.55 s with ddxdb 0.3.3 on +DataFusion 54). ddx builds the gradient program on the first step, which +takes 1.2 s, and runs it again on later ones. It plans each CTE the loss +reads twice on its own, and stores the result of each step of its program +as a table, so it computes each shared subexpression once. Doing the same +for forward queries is the first thing xarray-sql's DataFusion path needs +for these models, and einfold should plan them as graphs rather than trees. **Attention stores every score.** Causal attention's forward pass at 512 -tokens takes 0.68 s on DataFusion in the window spelling and 1.3 s in the -`GROUP BY` spelling; DuckDB takes 0.52 s and 0.26 s. Both grow with the -square of the tokens. The gradients take 1 to 2 times the forward pass: 1.1 s and 1.3 s on -DataFusion at 512 tokens, against 2.8 s and 2.7 s with ddxdb 0.3.2, -since ddx differentiates the masked product as `jax.grad` does. +tokens takes 0.28 s on DataFusion in the window spelling and 0.37 s in the +`GROUP BY` spelling; DuckDB takes 0.48 s and 0.27 s. Both grow with the +square of the tokens. The gradients take 2.5 to 3.5 times the forward +pass: 0.92 s and 0.93 s on DataFusion at 512 tokens, against 1.1 s and 1.3 +s with ddxdb 0.3.3 on DataFusion 54, and 2.8 s and 2.7 s with ddxdb 0.3.2, +before ddx differentiated the masked product as `jax.grad` does. **A regional forecast costs what a global one does.** For the contiguous United States, the forecast covers 55 of 2,048 grid cells. Walking the