diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index 30ede227..5c435adf 100644 --- a/benchmarks/ml/01_attention.py +++ b/benchmarks/ml/01_attention.py @@ -185,7 +185,7 @@ def main(): want = attention(q, k, v) want_grads = jax.grad(loss, argnums=(0, 1, 2))(q, k, v) - for host in hosts(args.hosts): + for host in hosts(args.hosts, args.memory): for name, data in zip("qkv", (q, k, v)): host.load(name, table(data)) diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index b0cd15f3..cf7c203e 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -569,7 +569,7 @@ def main(): ) ) - for host in hosts(args.hosts): + for host in hosts(args.hosts, args.memory): reference = dict(init) moments = { n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items() diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py new file mode 100644 index 00000000..08dfc661 --- /dev/null +++ b/benchmarks/ml/03_graphcast.py @@ -0,0 +1,531 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[ml,duckdb]", +# "weathernext @ git+https://github.com/google-deepmind/graphcast", +# "gcsfs", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Demo 3: GraphCast in SQL, with DeepMind's published weights. + +GraphCast (Lam et al., 2023, https://arxiv.org/abs/2212.12794) is +DeepMind's graph neural network for weather forecasting. It forecasts the +atmosphere 6 hours ahead from the two previous states. This runs +GraphCast_small, as DeepMind publishes it (1° grid, 13 pressure levels, a +mesh refined 5 times, 16 rounds of message passing, 36 million weights), +as a program of SQL statements over tables xarray-sql registers: + +* ``era5``: DeepMind's example inputs, two states 6 hours apart from ERA5 + (ECMWF's reanalysis of the atmosphere since 1940), and the forcings at + the forecast's time: inputs known in advance, such as the sunlight + reaching the top of the atmosphere; +* ``weights``, ``biases``, ``norms`` and ``nets``: every parameter in + DeepMind's checkpoint, one row each; +* ``grid2mesh``, ``mesh`` and ``mesh2grid``: the three graphs, as tables of + edges and of edge features, built by DeepMind's code. + +The model has three stages: + +* an **encoder** embeds each grid cell's 183 input channels, and passes + one round of messages from the grid to the mesh; +* a **processor** passes 16 rounds of messages along the multi-mesh, the + icosahedron's edges at every level of refinement; +* a **decoder** passes one round back to the grid, and turns each cell's + state into 83 output channels: the 6-hour change of each variable. + +Every network is the same SQL: an MLP (linear, swish, linear, layer norm) +over a table of ``(id, f, val)``, one row per node or edge and feature. +Swish is the activation ``x · sigmoid(x)``. Each stage is one statement, +and its result is stored as a table, so no host (an engine that runs the +SQL; see ``_hosts.py``) computes a stage twice. + +The check is DeepMind's own model. Its code, in JAX (Google's library for +differentiable array programs), run on the same inputs, +gives the stages' outputs; the SQL must match them. A plain JAX +transcription of the model (``_graphcast.mirror``), checked against +DeepMind's, also gives each processor round, so a run can stop after a +few rounds and still be checked. A full forecast is about two trillion +multiply-adds, which hash joins take hours over; DeepMind's JAX code +takes about 15 seconds on a 6-core CPU, including its compilation. + +Run ``uv run benchmarks/ml/03_graphcast.py --help`` for the options. +``--sizes`` is the number of processor rounds to run; with 16, the +forecast is complete. +""" + +from __future__ import annotations + +import jax +import numpy as np + +import _graphcast as gc +from _bench import Report, assert_close, cli, hosts + +jax.config.update("jax_enable_x64", True) + +INPUTS = 183 # input channels: the variables at each time and level +ROUNDS = 16 # processor rounds in GraphCast_small + +# -------------------------------------------------------------------------- +# The networks, in SQL. A state is a table of (id, f, val). +# -------------------------------------------------------------------------- + + +def weights(name: str) -> str: + """The weight matrix of one linear layer, by DeepMind's name.""" + return f"""( + SELECT w.f, w.g, w.w FROM nets n JOIN weights w ON w.net = n.net + WHERE n.name = '{name}')""" + + +def biases(name: str) -> str: + """The bias vector of one linear layer, by DeepMind's name.""" + return f"""( + SELECT b.g, b.b FROM nets n JOIN biases b ON b.net = n.net + WHERE n.name = '{name}')""" + + +def mlp(out: str, x: str, net: str, layer_norm: bool = True) -> str: + """DeepMind's MLP over the features of each ``id`` of ``x``: linear, + swish, linear, and (unless it is the decoder's last) a layer norm with + a learned scale and offset. A linear layer is a matrix product, a + join and a sum, and then its bias.""" + first, second = f"{net}_mlp/linear_0", f"{net}_mlp/linear_1" + sql = f""" +{out}_product_0 AS ( + SELECT x.id, w.g AS f, SUM(x.val * w.w) AS val + FROM {x} x JOIN {weights(first)} w ON x.f = w.f + GROUP BY x.id, w.g), +{out}_hidden AS ( + SELECT h.id, h.f, (h.val + b.b) / (1 + exp(-(h.val + b.b))) AS val + FROM {out}_product_0 h JOIN {biases(first)} b ON h.f = b.g), +{out}_product_1 AS ( + SELECT h.id, w.g AS f, SUM(h.val * w.w) AS val + FROM {out}_hidden h JOIN {weights(second)} w ON h.f = w.f + GROUP BY h.id, w.g), +{out}_linear AS ( + SELECT y.id, y.f, y.val + b.b AS val + FROM {out}_product_1 y JOIN {biases(second)} b ON y.f = b.g),""" + if not layer_norm: + return sql + f"\n{out} AS (SELECT * FROM {out}_linear)," + return ( + sql + + f""" +{out} AS ( + SELECT y.id, y.f, + (y.val - AVG(y.val) OVER (PARTITION BY y.id)) + / sqrt(VAR_POP(y.val) OVER (PARTITION BY y.id) + 1e-5) + * n.scale + n.offset AS val + FROM {out}_linear y + JOIN norms n ON n.g = y.f + JOIN nets m ON m.net = n.net + WHERE m.name = '{net}_layer_norm'),""" + ) + + +def select(ctes: str, out: str) -> str: + return f"WITH {ctes.rstrip(',')}\nSELECT id, f, val FROM {out}" + + +def embed(x: str, net: str) -> str: + """A statement: an MLP of every row of ``x``.""" + return select(mlp("embedded", x, net), "embedded") + + +def messages(net: str, graph: str, edges: str, senders: str, receivers: str): + """A statement: each edge's message, an MLP of the edge's state, its + sender's and its receiver's, concatenated.""" + return select( + f""" +edge_in AS ( + SELECT x.id, x.f, x.val FROM {edges} x JOIN {graph} e ON x.id = e.edge + UNION ALL + SELECT e.edge AS id, s.f + 512 AS f, s.val + FROM {graph} e JOIN {senders} s ON e.sender = s.id + UNION ALL + SELECT e.edge AS id, r.f + 1024 AS f, r.val + FROM {graph} e JOIN {receivers} r ON e.receiver = r.id),""" + + mlp("message", "edge_in", net), + "message", + ) + + +def update(net: str, graph: str, nodes: str, messages: str | None) -> str: + """A statement: each node plus an MLP of itself and the sum of the + messages it receives (none, for the encoder's grid cells).""" + node_in = f"SELECT id, f, val FROM {nodes}" + if messages: + node_in += f""" + UNION ALL + SELECT e.receiver AS id, m.f + 512 AS f, SUM(m.val) AS val + FROM {messages} m JOIN {graph} e ON m.id = e.edge + GROUP BY e.receiver, m.f""" + return select( + f"\nnode_in AS (\n {node_in})," + + mlp("updated", "node_in", net) + + f""" +new AS ( + SELECT x.id, x.f, x.val + u.val AS val + FROM {nodes} x JOIN updated u ON x.id = u.id AND x.f = u.f),""", + "new", + ) + + +def residual(state: str, change: str) -> str: + """A statement: ``state + change``, row by row.""" + return f"""SELECT x.id, x.f, x.val + d.val AS val +FROM {state} x JOIN {change} d ON x.id = d.id AND x.f = d.f""" + + +ERA5_TABLES = { + ("time", "lat", "lon"): "surface", + ("time", "level", "lat", "lon"): "atmosphere", + ("lat", "lon"): "static", + ("time", "lon"): "day", + ("time",): "year", +} +# How each era5 table's rows reach the grid's cells. +CELLS = { + "surface": "JOIN cells c ON x.lat = c.lat AND x.lon = c.lon", + "atmosphere": "JOIN cells c ON x.lat = c.lat AND x.lon = c.lon", + "static": "JOIN cells c ON x.lat = c.lat AND x.lon = c.lon", + "day": "JOIN cells c ON x.lon = c.lon", + "year": "CROSS JOIN cells c", +} + + +def grid_inputs(era5) -> str: + """A statement: each cell's input channels, normalized, then its + position. Every variable is read from its ``era5`` table at the times + and levels the channels ask for.""" + parts = [] + for name in sorted(era5.data_vars): + dims = era5[name].dims + table = ERA5_TABLES[dims] + time = "ch.time = x.time" if "time" in dims else "ch.time = -1" + level = "ch.level = x.level" if "level" in dims else "ch.level = 0" + parts.append( + f"""SELECT c.cell AS id, ch.channel AS f, + (x."{name}" - ch.mean) / ch.std AS val +FROM era5.{table} x +{CELLS[table]} +JOIN channels ch ON ch.variable = '{name}' AND {time} AND {level}""" + ) + parts.append( + f"SELECT cell AS id, f + {INPUTS} AS f, val FROM cell_features" + ) + return "\nUNION ALL\n".join(parts) + + +def forecast(era5) -> str: + """A statement: the forecast of each variable at each cell and level, + the last input plus the output's normalized 6-hour change.""" + targets = sorted(set(gc.outputs_table().variable.values)) + last = "\n UNION ALL\n".join( + f""" SELECT c.cell AS id, '{name}' AS variable, + {"x.level" if "level" in era5[name].dims else "0"} AS level, + x."{name}" AS value + FROM era5.{ERA5_TABLES[era5[name].dims]} x + {CELLS[ERA5_TABLES[era5[name].dims]]} + WHERE x.time = 0""" + for name in targets + ) + return f"""WITH last_input AS ( +{last}) +SELECT c.lat, c.lon, o.variable, o.level, + x.value + y.val * o.diff_std AS value +FROM output y +JOIN outputs o ON o.channel = y.f +JOIN cells c ON c.cell = y.id +JOIN last_input x + ON x.id = y.id AND x.variable = o.variable AND x.level = o.level +ORDER BY o.variable, o.level, c.lat, c.lon""" + + +def flatten(ds) -> np.ndarray: + """A forecast Dataset in ``forecast``'s row order: by variable, level, + latitude and longitude.""" + return np.concatenate( + [ + ds[name].transpose(..., "lat", "lon").values.ravel() + for name in sorted(ds.data_vars) + ] + ) + + +# -------------------------------------------------------------------------- +# The program: one stored table per stage. +# -------------------------------------------------------------------------- + + +# A region to forecast: the contiguous United States, roughly. +BOUNDS = {"lat": (25, 50), "lon": (235, 295)} +REGION = " AND ".join( + f"c.{coord} BETWEEN {low} AND {high}" + for coord, (low, high) in BOUNDS.items() +) + + +def receptive_field(rounds: int): + """The cells and mesh nodes a regional forecast needs, as stored + ``(table, statement)``s, found by walking the graphs backward from the + region: semi-joins of each graph's edges with what the next stage + needs.""" + yield "needed_cells", f"SELECT cell AS id FROM cells c WHERE {REGION}" + yield ( + f"needed_mesh_{rounds}", + """SELECT DISTINCT sender AS id +FROM mesh2grid.edges WHERE receiver IN (SELECT id FROM needed_cells)""", + ) + for r in reversed(range(rounds)): + yield ( + f"needed_mesh_{r}", + f"""SELECT id FROM needed_mesh_{r + 1} +UNION +SELECT sender AS id FROM mesh.edges +WHERE receiver IN (SELECT id FROM needed_mesh_{r + 1})""", + ) + yield ( + "needed_grid", + """SELECT id FROM needed_cells +UNION +SELECT sender AS id FROM grid2mesh.edges +WHERE receiver IN (SELECT id FROM needed_mesh_0)""", + ) + + +def program(era5, rounds: int, regional: bool = False): + """The stages in order, as ``(table, statement)``; ``(table, None)`` + marks a table the reference checks. A regional program runs the same + stages on the rows its receptive field needs.""" + + def within(table: str, column: str, needed: str) -> str: + if not regional: + return table + return f"(SELECT * FROM {table} WHERE {column} IN (SELECT id FROM {needed}))" + + def features(graph: str, edges: str) -> str: + return f"""(SELECT x.edge AS id, x.f, x.val FROM {graph}.features x + WHERE x.edge IN (SELECT edge FROM {edges}))""" + + checked = not regional + yield ( + "grid_in", + ( + f"SELECT * FROM ({grid_inputs(era5)}) g\n" + "WHERE id IN (SELECT id FROM needed_grid)" + if regional + else grid_inputs(era5) + ), + ) + if checked: + yield "grid_in", None + # Encoder: embed the grid, the mesh and the grid-to-mesh edges, then + # one round of messages from the grid to the mesh. + g2m = "grid2mesh_gnn" + g2m_graph = within("grid2mesh.edges", "receiver", "needed_mesh_0") + yield "grid", embed("grid_in", f"{g2m}/encoder_nodes_grid_nodes") + # DeepMind gives the mesh nodes 183 zero inputs; zeros are absent rows. + yield ( + "mesh", + embed( + f"""(SELECT node AS id, f + {INPUTS} AS f, val + FROM {within("node_features", "node", "needed_mesh_0")} n)""", + f"{g2m}/encoder_nodes_mesh_nodes", + ), + ) + yield ( + "g2m_edges", + embed( + features("grid2mesh", g2m_graph), f"{g2m}/encoder_edges_grid2mesh" + ), + ) + yield ( + "g2m_messages", + messages( + f"{g2m}/processor_edges_0_grid2mesh", + g2m_graph, + "g2m_edges", + "grid", + "mesh", + ), + ) + yield ( + "mesh_0", + update( + f"{g2m}/processor_nodes_0_mesh_nodes", + g2m_graph, + "mesh", + "g2m_messages", + ), + ) + yield ( + "grid_0", + update( + f"{g2m}/processor_nodes_0_grid_nodes", + g2m_graph, + within("grid", "id", "needed_cells"), + None, + ), + ) + if checked: + yield "mesh_0", None + yield "grid_0", None + # Processor: rounds of messages along the multi-mesh, edges and nodes + # each updated with a residual. + graphs = [ + within("mesh.edges", "receiver", f"needed_mesh_{r + 1}") + for r in range(rounds) + ] + yield ( + "mesh_edges_0", + embed( + features("mesh", graphs[0] if rounds else "mesh.edges"), + "mesh_gnn/encoder_edges_mesh", + ), + ) + for r in range(rounds): + yield ( + f"mesh_messages_{r}", + messages( + f"mesh_gnn/processor_edges_{r}_mesh", + graphs[r], + f"mesh_edges_{r}", + f"mesh_{r}", + f"mesh_{r}", + ), + ) + yield ( + f"mesh_{r + 1}", + update( + f"mesh_gnn/processor_nodes_{r}_mesh_nodes", + graphs[r], + within(f"mesh_{r}", "id", f"needed_mesh_{r + 1}"), + f"mesh_messages_{r}", + ), + ) + yield ( + f"mesh_edges_{r + 1}", + residual(f"mesh_edges_{r}", f"mesh_messages_{r}"), + ) + if checked: + yield f"mesh_{r + 1}", None + if rounds < ROUNDS: + return + # Decoder: one round of messages from the mesh to the grid, then each + # cell's 83 output channels. + m2g = "mesh2grid_gnn" + m2g_graph = within("mesh2grid.edges", "receiver", "needed_cells") + yield ( + "m2g_edges", + embed( + features("mesh2grid", m2g_graph), f"{m2g}/encoder_edges_mesh2grid" + ), + ) + yield ( + "m2g_messages", + messages( + f"{m2g}/processor_edges_0_mesh2grid", + m2g_graph, + "m2g_edges", + f"mesh_{rounds}", + "grid_0", + ), + ) + yield ( + "grid_1", + update( + f"{m2g}/processor_nodes_0_grid_nodes", + m2g_graph, + "grid_0", + "m2g_messages", + ), + ) + yield ( + "output", + select( + mlp( + "out", + "grid_1", + f"{m2g}/decoder_nodes_grid_nodes", + layer_norm=False, + ), + "out", + ), + ) + if checked: + yield "output", None + + +def matrix(host, table: str) -> np.ndarray: + """A stored ``(id, f, val)`` table as an (id, f) matrix.""" + rows = host.query(f"SELECT id, f, val FROM {table} ORDER BY id, f") + n = rows["id"][-1].as_py() + 1 + return rows["val"].to_numpy().reshape(n, -1) + + +def main(): + parser = cli(__doc__.splitlines()[0], sizes=[1], size_name="rounds") + args = parser.parse_args() + report = Report("graphcast", "rounds", args.reps) + + print(" loading DeepMind's checkpoint, example and graphs") + reference = gc.deepmind_stages() + era5 = gc.era5() + tables = { + "era5": era5, + "channels": gc.channels_table(), + "outputs": gc.outputs_table(), + **gc.graph_tables(), + **gc.param_tables(), + } + + for rounds in args.sizes: + mirror = gc.mirror(reference["features"], rounds=rounds) + expected = { + "grid_in": np.concatenate( + [reference["features"], gc.graphs()["grid_nodes"]], axis=1 + ), + "grid_0": reference["grid"], + "mesh_0": reference["mesh"], + **{f"mesh_{r + 1}": mirror[f"round_{r}"] for r in range(rounds)}, + "output": reference["output"], + } + for host in hosts(args.hosts, args.memory): + for name, table in tables.items(): + options = {} + if name == "era5": + options = {"table_names": ERA5_TABLES} + elif "edge" in table.dims: + options = {"table_names": gc.GRAPH_TABLES} + host.load(name, table, **options) + + for table, statement in program(era5, rounds): + if statement is None: + got = matrix(host, table) + # The reference runs in float32, the SQL in float64. + error = assert_close( + table, got, expected[table], rtol=1e-4, atol=1e-3 + ) + report.check(error, step=table) + continue + report.time(table, host, rounds, host.store, table, statement) + + if rounds == ROUNDS: + rows = report.time( + "forecast", host, rounds, host.query, forecast(era5) + ) + want = flatten(gc.deepmind_forecast().isel(time=0)) + error = assert_close( + "forecast", rows["value"].to_numpy(), want, rtol=1e-4 + ) + report.check(error) + + report.finish(args.csv) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index ae5b5950..e0e080d2 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -1,14 +1,14 @@ # Machine learning in SQL -These scripts write machine-learning models as SQL over tables that +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. Two of these models - are its [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md), + 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 @@ -19,6 +19,7 @@ array programs. |---|---|---|---| | 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, 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, and the slow parts are left slow on purpose: they are what ddx, xarray-sql @@ -38,15 +39,16 @@ mismatch fails the run. ```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 each: +options are the same in all three: ```shell --hosts datafusion duckdb # the engines to run on (default: all) ---sizes 64 256 # problem sizes: tokens, or batch size +--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 ``` @@ -55,7 +57,8 @@ Demo 2 adds `--steps`, the number of training steps, and `--optimizer`: `sgd` (the default), or `adamw` to train as nanoGPT does, and the model's size: `--layers`, `--embed`, `--heads` and `--context`, nanoGPT's `n_layer`, `n_embd`, `n_head` and `block_size` (1, 16, 2 and 16 by -default; nanoGPT's Shakespeare model is 6, 384, 6 and 256). +default; nanoGPT's Shakespeare model is 6, 384, 6 and 256). Demo 3 adds +`--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 @@ -130,13 +133,20 @@ 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 -three times, a residual connection twice. DataFusion inlines each read, +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 | 13,756 | 126 | 3.0 s | 0.063 s | +| GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | +| GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | + +The GraphCast rows are from 2026-10-17, before in-memory tables. + +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 @@ -193,12 +203,23 @@ against 0.20 s for one step), but copying Q, K and V into native tables only brings the gradient to 1.7 s, and slows the forward pass (measured on 2026-10-24, with ddx's development branch and DuckDB 1.5.6). +**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 `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. - On DuckDB, a view whose CTE computes a window and is read more than once, as the softmax's weights are in attention's window spelling and in the GPT, so neither trains on DuckDB yet. ddx reads DuckDB's plans diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py index 5a3235e0..af6b209d 100644 --- a/benchmarks/ml/_bench.py +++ b/benchmarks/ml/_bench.py @@ -50,13 +50,19 @@ def cli(description: str, sizes: Sequence[int], size_name: str = "n"): default=1, help="timed repetitions per step; the median is reported", ) + parser.add_argument( + "--memory", + type=float, + default=6, + help="each host's memory limit, in GB, past which it spills to disk", + ) 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 hosts(names: Sequence[str], memory: float = 6) -> list[Host]: + """A fresh instance of each named host, limited to ``memory`` GB.""" + return [HOSTS[name](int(memory * 2**30)) for name in names] def assert_close(name: str, got, want, *, rtol=1e-9, atol=1e-9) -> float: @@ -116,10 +122,11 @@ def attempt(self, step: str, host: Host, size: int, run: Callable, *args): self.skip(step, host, size, f"refused: {reason}") return None - 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 check(self, error: float, step: str | None = None) -> None: + """Mark a row's result as checked, with its error (from + :func:`assert_close`): the last row, or the last row of ``step``.""" + row = next(r for r in reversed(self.rows) if step in (None, r["step"])) + row["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.""" diff --git a/benchmarks/ml/_graphcast.py b/benchmarks/ml/_graphcast.py new file mode 100644 index 00000000..cef6de7d --- /dev/null +++ b/benchmarks/ml/_graphcast.py @@ -0,0 +1,490 @@ +"""GraphCast_small, as DeepMind publishes it, for the SQL demo to match. + +Everything here is DeepMind's code or data, or a direct transcription of +it into tables: + +* :func:`checkpoint`, :func:`example` and :func:`stats` fetch the + published weights, an example of the model's inputs from ERA5 (ECMWF's + reanalysis of the atmosphere) on a 1° grid, and the normalization + statistics from DeepMind's public bucket, cached in + ``~/.cache/xarray-sql``; +* :func:`graphs` builds the three graphs with DeepMind's own code: which + cells and mesh nodes are joined, and their geometric features; +* :func:`deepmind_stages` runs DeepMind's model, in float32, and returns + the output of each stage: the reference the SQL must match; +* :func:`mirror` is the same model written plainly in JAX, from the same + weights. It is checked against :func:`deepmind_stages`, and it can stop + after any processor round, which DeepMind's code cannot, so each SQL + stage can be checked on its own; +* :func:`param_tables`, :func:`channels_table` and :func:`graph_tables` + turn the weights, the input channels and the graphs into Datasets that + xarray-sql registers as tables. +""" + +from __future__ import annotations + +import dataclasses +import functools +import pathlib + +import gcsfs +import haiku as hk +import jax +import jax.numpy as jnp +import numpy as np +import xarray as xr +from weathernext.utils import checkpoint as checkpoint_lib +from weathernext.utils import data_utils, normalization +from weathernext.weathernext1_graph import graphcast + +CACHE = pathlib.Path.home() / ".cache" / "xarray-sql" / "graphcast" +BUCKET = "dm_graphcast/graphcast" +PARAMS = ( + "params/GraphCast_small - ERA5 1979-2015 - resolution 1.0 - " + "pressure levels 13 - mesh 2to5 - precipitation input and output.npz" +) +EXAMPLE = "dataset/source-era5_date-2022-01-01_res-1.0_levels-13_steps-01.nc" +STATS = { + "diffs_stddev": "stats/diffs_stddev_by_level.nc", + "mean": "stats/mean_by_level.nc", + "stddev": "stats/stddev_by_level.nc", +} + + +def fetch(path: str) -> pathlib.Path: + """A file from DeepMind's bucket, downloaded once.""" + local = CACHE / pathlib.Path(path).name + if not local.exists(): + CACHE.mkdir(parents=True, exist_ok=True) + gcsfs.GCSFileSystem(token="anon").get(f"{BUCKET}/{path}", str(local)) + return local + + +@functools.cache +def checkpoint() -> graphcast.CheckPoint: + with open(fetch(PARAMS), "rb") as f: + return checkpoint_lib.load(f, graphcast.CheckPoint) + + +@functools.cache +def stats() -> dict[str, xr.Dataset]: + return { + name: xr.load_dataset(fetch(path)).compute() + for name, path in STATS.items() + } + + +@functools.cache +def example() -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + """The inputs, the 6-hour target and the forcings, as DeepMind's + notebook extracts them from the example, without the batch dim.""" + data = xr.load_dataset(fetch(EXAMPLE), decode_timedelta=True).compute() + task = dataclasses.asdict(checkpoint().task_config) + inputs, targets, forcings = data_utils.extract_inputs_targets_forcings( + data, target_lead_times=slice("6h", "6h"), **task + ) + return tuple( + ds.isel(batch=0).drop_vars("datetime", errors="ignore") + for ds in (inputs, targets, forcings) + ) + + +# -------------------------------------------------------------------------- +# DeepMind's model: the graphs it builds and the output of each stage. +# -------------------------------------------------------------------------- + + +@functools.cache +def _model() -> graphcast.GraphCast: + """A GraphCast with its graphs built for the example's grid.""" + ckpt = checkpoint() + built = {} + + def build(): + model = graphcast.GraphCast(ckpt.model_config, ckpt.task_config) + model._maybe_init(example()[0]) + built["model"] = model + + hk.transform(build).init(None) + return built["model"] + + +@functools.cache +def graphs() -> dict: + """The grid-to-mesh, mesh and mesh-to-grid graphs: their edges' + senders and receivers, and the nodes' and edges' features.""" + model = _model() + out = {} + for name, graph in ( + ("grid2mesh", model._grid2mesh_graph_structure), + ("mesh", model._mesh_graph_structure), + ("mesh2grid", model._mesh2grid_graph_structure), + ): + (edges,) = graph.edges.values() + out[name] = { + "senders": np.asarray(edges.indices.senders), + "receivers": np.asarray(edges.indices.receivers), + "features": np.asarray(edges.features), + } + nodes = model._grid2mesh_graph_structure.nodes + out["grid_nodes"] = np.asarray(nodes["grid_nodes"].features) + out["mesh_nodes"] = np.asarray(nodes["mesh_nodes"].features) + out["grid_lat"] = model._grid_nodes_lat + out["grid_lon"] = model._grid_nodes_lon + out["mesh_lat"] = model._mesh_nodes_lat + out["mesh_lon"] = model._mesh_nodes_lon + return out + + +def _normalized(inputs, forcings): + s = stats() + return ( + normalization.normalize(inputs, s["stddev"], s["mean"]), + normalization.normalize(forcings, s["stddev"], s["mean"]), + ) + + +@functools.cache +def deepmind_stages() -> dict[str, np.ndarray]: + """DeepMind's model, in float32, stage by stage: the grid features it + reads, the latent grid and mesh after the encoder, the mesh after the + processor, and the decoder's normalized output.""" + ckpt = checkpoint() + inputs, _, forcings = example() + + @hk.transform_with_state + def stages(inputs, forcings): + model = graphcast.GraphCast(ckpt.model_config, ckpt.task_config) + inputs, forcings = _normalized(inputs, forcings) + model._maybe_init(inputs) + features = model._inputs_to_grid_node_features( + inputs.expand_dims("batch"), forcings.expand_dims("batch") + ) + mesh, grid = model._run_grid2mesh_gnn(features) + processed = model._run_mesh_gnn(mesh) + output = model._run_mesh2grid_gnn(processed, grid) + return dict( + features=features, + grid=grid, + mesh=mesh, + processed=processed, + output=output, + ) + + out, _ = jax.jit(stages.apply)( + ckpt.params, + {}, + None, + inputs.astype(np.float32), + forcings.astype(np.float32), + ) + return {k: np.asarray(v)[:, 0] for k, v in out.items()} + + +@functools.cache +def deepmind_forecast() -> xr.Dataset: + """DeepMind's 6-hour forecast, in float32, with its normalization and + residual: what the notebook's ``run_forward`` returns, without the + bfloat16 cast.""" + ckpt = checkpoint() + inputs, targets, forcings = example() + s = stats() + + @hk.transform_with_state + def forward(inputs, targets, forcings): + model = normalization.InputsAndResiduals( + graphcast.GraphCast(ckpt.model_config, ckpt.task_config), + diffs_stddev_by_level=s["diffs_stddev"], + mean_by_level=s["mean"], + stddev_by_level=s["stddev"], + ) + return model(inputs, targets_template=targets, forcings=forcings) + + batch = [ + ds.expand_dims("batch").astype(np.float32) + for ds in (inputs, targets * np.nan, forcings) + ] + out, _ = forward.apply(ckpt.params, {}, None, *batch) + return out.isel(batch=0).compute() + + +# -------------------------------------------------------------------------- +# The same model, written plainly: the oracle for each stage. +# -------------------------------------------------------------------------- + + +def weights(net: str) -> dict[str, np.ndarray]: + """One network's parameters, by DeepMind's name for it, e.g. + ``mesh_gnn/processor_edges_3_mesh``.""" + params = checkpoint().params + gnn, block = net.split("/") + prefix = f"{gnn}/~_networks_builder/{block}" + out = {} + for layer in ("linear_0", "linear_1"): + for p, v in params[f"{prefix}_mlp/~/{layer}"].items(): + out[f"{layer}:{p}"] = np.asarray(v, np.float64) + if f"{prefix}_layer_norm" in params: + for p, v in params[f"{prefix}_layer_norm"].items(): + out[f"layer_norm:{p}"] = np.asarray(v, np.float64) + return out + + +def mlp(net: str, x): + """Linear, swish, linear, then a layer norm, as DeepMind's MLPs are.""" + p = weights(net) + h = jax.nn.swish(x @ p["linear_0:w"] + p["linear_0:b"]) + y = h @ p["linear_1:w"] + p["linear_1:b"] + if "layer_norm:scale" not in p: + return y + mean = y.mean(-1, keepdims=True) + var = ((y - mean) ** 2).mean(-1, keepdims=True) + return (y - mean) / jnp.sqrt(var + 1e-5) * p["layer_norm:scale"] + p[ + "layer_norm:offset" + ] + + +def interaction( + gnn, step, edge_set, graph, edges, senders, receivers, node_set +): + """One message-passing step: new edges, then the receivers' update, + each with its residual. Returns the new edges and receivers.""" + s, r = graph["senders"], graph["receivers"] + message = mlp( + f"{gnn}/processor_edges_{step}_{edge_set}", + jnp.concatenate([edges, senders[s], receivers[r]], axis=1), + ) + total = jax.ops.segment_sum(message, r, len(receivers)) + update = mlp( + f"{gnn}/processor_nodes_{step}_{node_set}", + jnp.concatenate([receivers, total], axis=1), + ) + return edges + message, receivers + update + + +def mirror(features: np.ndarray, rounds: int | None = None) -> dict: + """The model from the grid features DeepMind's model reads, in float64. + With ``rounds``, it stops after that many processor rounds.""" + g = graphs() + x = jnp.asarray(features, jnp.float64) + grid_in = jnp.concatenate([x, g["grid_nodes"]], axis=1) + mesh_in = jnp.concatenate( + [jnp.zeros((len(g["mesh_nodes"]), x.shape[1])), g["mesh_nodes"]], axis=1 + ) + grid = mlp("grid2mesh_gnn/encoder_nodes_grid_nodes", grid_in) + mesh = mlp("grid2mesh_gnn/encoder_nodes_mesh_nodes", mesh_in) + edges = mlp( + "grid2mesh_gnn/encoder_edges_grid2mesh", g["grid2mesh"]["features"] + ) + _, mesh = interaction( + "grid2mesh_gnn", + 0, + "grid2mesh", + g["grid2mesh"], + edges, + grid, + mesh, + "mesh_nodes", + ) + grid = grid + mlp("grid2mesh_gnn/processor_nodes_0_grid_nodes", grid) + out = {"grid": grid, "mesh": mesh} + + edges = mlp("mesh_gnn/encoder_edges_mesh", g["mesh"]["features"]) + total = checkpoint().model_config.gnn_msg_steps + for step in range(total if rounds is None else rounds): + edges, mesh = interaction( + "mesh_gnn", step, "mesh", g["mesh"], edges, mesh, mesh, "mesh_nodes" + ) + out[f"round_{step}"] = mesh + if rounds is not None and rounds < total: + return out + out["processed"] = mesh + + edges = mlp( + "mesh2grid_gnn/encoder_edges_mesh2grid", g["mesh2grid"]["features"] + ) + _, grid = interaction( + "mesh2grid_gnn", + 0, + "mesh2grid", + g["mesh2grid"], + edges, + mesh, + grid, + "grid_nodes", + ) + out["output"] = mlp("mesh2grid_gnn/decoder_nodes_grid_nodes", grid) + return out + + +# -------------------------------------------------------------------------- +# Tables: the example, the input channels, the graphs and the weights. +# -------------------------------------------------------------------------- + + +def _hours(ds: xr.Dataset) -> xr.Dataset: + """``time`` as whole hours from the forecast's start.""" + return ds.assign_coords(time=(ds.time / np.timedelta64(1, "h")).astype(int)) + + +def era5() -> xr.Dataset: + """The model's inputs at -6 h and 0 h, and the forcings at +6 h.""" + inputs, _, forcings = example() + return xr.merge( + [_hours(inputs), _hours(forcings)], compat="no_conflicts", join="outer" + ) + + +def _stacking(ds: xr.Dataset, time_offset: int = 0) -> list[tuple]: + """DeepMind's channel order for ``ds``: its variables sorted by name, + each over its own dimensions other than lat and lon, as + ``(variable, time, level)``; ``level`` is 0 for a surface variable.""" + channels = [] + for name in sorted(ds.data_vars): + var = ds[name] + times = var.time.values if "time" in var.dims else [None] + levels = var.level.values if "level" in var.dims else [0] + for t in times: + for level in levels: + time = None if t is None else int(t / np.timedelta64(1, "h")) + channels.append((name, time, int(level))) + return channels + + +def channels_table() -> xr.Dataset: + """Each input channel: its variable, time and level, and the mean and + standard deviation that normalize it.""" + inputs, _, forcings = example() + rows = _stacking(inputs) + _stacking(forcings) + s = stats() + + def stat(kind, name, level): + value = s[kind][name] + return float(value.sel(level=level) if "level" in value.dims else value) + + return xr.Dataset( + { + "variable": ("channel", [r[0] for r in rows]), + # Static variables have no time; -1 never matches a row's time. + "time": ("channel", [-1 if r[1] is None else r[1] for r in rows]), + "level": ("channel", [r[2] for r in rows]), + "mean": ("channel", [stat("mean", r[0], r[2]) for r in rows]), + "std": ("channel", [stat("stddev", r[0], r[2]) for r in rows]), + }, + coords={"channel": np.arange(len(rows))}, + ) + + +def outputs_table() -> xr.Dataset: + """Each output channel: its variable and level, and the standard + deviation of its 6-hour change, which turns the output into one.""" + _, targets, _ = example() + rows = _stacking(targets.isel(time=0, drop=True)) + diffs = stats()["diffs_stddev"] + + def std(name, level): + value = diffs[name] + return float(value.sel(level=level) if "level" in value.dims else value) + + return xr.Dataset( + { + "variable": ("channel", [r[0] for r in rows]), + "level": ("channel", [r[2] for r in rows]), + "diff_std": ("channel", [std(r[0], r[2]) for r in rows]), + }, + coords={"channel": np.arange(len(rows))}, + ) + + +def _long(name: str, values: np.ndarray, key: str) -> xr.Dataset: + """A (key, feature) matrix as the table ``(key, f, val)``.""" + return xr.Dataset( + {"val": ((key, "f"), values.astype(np.float64))}, + coords={key: np.arange(len(values)), "f": np.arange(values.shape[1])}, + ) + + +GRAPH_TABLES = {("edge",): "edges", ("edge", "f"): "features"} + + +def graph_tables() -> dict[str, xr.Dataset]: + """The grid, the mesh, and the three graphs, as Datasets. Grid cells + are numbered as DeepMind numbers them: by latitude, then longitude.""" + g = graphs() + tables = { + "cells": xr.Dataset( + {"lat": ("cell", g["grid_lat"]), "lon": ("cell", g["grid_lon"])}, + coords={"cell": np.arange(len(g["grid_lat"]))}, + ), + "cell_features": _long("cell_features", g["grid_nodes"], "cell"), + "node_features": _long("node_features", g["mesh_nodes"], "node"), + } + for name in ("grid2mesh", "mesh", "mesh2grid"): + graph = g[name] + tables[name] = xr.Dataset( + { + "sender": ("edge", graph["senders"]), + "receiver": ("edge", graph["receivers"]), + "val": (("edge", "f"), graph["features"].astype(np.float64)), + }, + coords={ + "edge": np.arange(len(graph["senders"])), + "f": np.arange(graph["features"].shape[1]), + }, + ) + return tables + + +def _short(module: str) -> str: + """DeepMind's module path without Haiku's scaffolding: + ``mesh_gnn/processor_edges_3_mesh_mlp/linear_0``.""" + return module.replace("~_networks_builder/", "").replace("~/", "") + + +def param_tables() -> dict[str, xr.Dataset]: + """Every parameter, long form: ``weights(net, f, g, w)``, + ``biases(net, g, b)``, ``norms(net, g, scale, offset)``, and + ``nets(net, name)`` naming each network by DeepMind's path.""" + params = checkpoint().params + names = sorted(_short(m) for m in params) + ids = {name: i for i, name in enumerate(names)} + weights, biases, norms = [], [], [] + for module, p in params.items(): + net = ids[_short(module)] + if "w" in p: + f, g = np.indices(p["w"].shape).reshape(2, -1) + weights.append((np.full(f.size, net), f, g, p["w"].ravel())) + biases.append( + (np.full(p["b"].size, net), np.arange(p["b"].size), p["b"]) + ) + else: + norms.append( + ( + np.full(p["scale"].size, net), + np.arange(p["scale"].size), + p["scale"], + p["offset"], + ) + ) + + def table(dim, columns, parts): + arrays = [np.concatenate(c) for c in zip(*parts)] + return xr.Dataset( + { + name: ( + dim, + a.astype(np.float64 if a.dtype.kind == "f" else np.int32), + ) + for name, a in zip(columns, arrays) + }, + coords={dim: np.arange(len(arrays[0]))}, + ) + + return { + "nets": xr.Dataset( + {"name": ("net", np.array(names, dtype=object))}, + coords={"net": np.arange(len(names), dtype=np.int32)}, + ), + "weights": table("weight", ("net", "f", "g", "w"), weights), + "biases": table("bias", ("net", "g", "b"), biases), + "norms": table("norm", ("net", "g", "scale", "offset"), norms), + } diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index d5bcadcf..3929b24b 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -5,7 +5,7 @@ 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 five things: +queries. A host does these things: * ``load(name, data, **options)`` registers an ``xr.Dataset`` as the table ``name`` through xarray-sql, replacing any table of that name @@ -17,7 +17,10 @@ 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. + run, which shows, for one, whether a CTE read twice is computed twice; +* ``store(name, sql)`` runs a statement and keeps its rows as the table + ``name``, and ``drop(name)`` removes it: a model too large for one query + runs as a program of stored stages. Forward queries run everywhere. Gradients (``grad``, ``vjp``, ``jvp``) need ddx (https://github.com/xqlsystems/ddx), which differentiates SQL @@ -27,6 +30,10 @@ A gradient that ddx (or the engine under it) refuses to plan raises one of the host's ``refusals``, and the demos report that step with the reason. +Every host gets the same memory limit, in bytes, and spills to +``~/.cache/xarray-sql/spill`` past it, so a stage too large for memory +slows down rather than taking the machine down. + Adding a host (einfold's optimizer on DataFusion, or DuckDB with a GPU extension) is one class here and one entry in ``HOSTS``. """ @@ -34,17 +41,22 @@ from __future__ import annotations import json +import pathlib from collections.abc import Sequence import duckdb import pyarrow as pa import xarray as xr +from datafusion import RuntimeEnvBuilder from ddxdb import DdxError from ddxdb import duckdb as ddx_duckdb import xarray_sql as xql +SPILL = pathlib.Path.home() / ".cache" / "xarray-sql" / "spill" + + class Host: """An engine that runs the demos' SQL. Subclasses fill in the rest.""" @@ -52,6 +64,9 @@ class Host: differentiates = False refusals: tuple[type[Exception], ...] = () + def __init__(self, memory: int): + SPILL.mkdir(parents=True, exist_ok=True) + def load(self, name: str, data: xr.Dataset, **options) -> None: raise NotImplementedError @@ -67,6 +82,12 @@ def query_all(self, sqls: Sequence[str]) -> list[pa.Table]: def operators(self, sql: str) -> int: raise NotImplementedError + def store(self, name: str, sql: str) -> None: + raise NotImplementedError + + def drop(self, name: str) -> None: + raise NotImplementedError + class DataFusion(Host): """xarray-sql's ``Context``, with ddx's gradients (the ``ml`` extra).""" @@ -75,8 +96,14 @@ class DataFusion(Host): differentiates = True refusals = (DdxError,) - def __init__(self): - self.ctx = xql.Context() + def __init__(self, memory): + super().__init__(memory) + runtime = ( + RuntimeEnvBuilder() + .with_disk_manager_specified(str(SPILL)) + .with_fair_spill_pool(memory) + ) + self.ctx = xql.Context(runtime=runtime) def load(self, name, data, **options): self.ctx.deregister_table(name) @@ -95,6 +122,14 @@ def operators(self, sql): plan = self.ctx.sql(sql).execution_plan().display_indent() return len(plan.splitlines()) + def store(self, name, sql): + partitions = self.ctx.sql(sql).collect_partitioned() + self.ctx.deregister_table(name) + self.ctx.register_record_batches(name, partitions) + + def drop(self, name): + self.ctx.deregister_table(name) + class DuckDB(Host): """DuckDB, reading the same Datasets through xarray-sql's adapter, with @@ -107,8 +142,14 @@ class DuckDB(Host): # plan it can't write. refusals = (DdxError, duckdb.NotImplementedException) - def __init__(self): - self.con = duckdb.connect() + def __init__(self, memory): + super().__init__(memory) + self.con = duckdb.connect( + config={ + "memory_limit": f"{memory // 2**20}MB", + "temp_directory": str(SPILL), + } + ) def load(self, name, data, **options): xql.register(self.con, name, data, **options) @@ -133,6 +174,12 @@ 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 store(self, name, sql): + self.con.sql(f"CREATE OR REPLACE TABLE {name} AS {sql}") + + def drop(self, name): + self.con.sql(f"DROP TABLE IF EXISTS {name}") + def _count(node: dict) -> int: return 1 + sum(_count(child) for child in node.get("children", []))