Skip to content

Commit 77b024b

Browse files
alxmrsclaude
andcommitted
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 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent 659e555 commit 77b024b

2 files changed

Lines changed: 39 additions & 16 deletions

File tree

‎benchmarks/ml/02_nanogpt.py‎

Lines changed: 38 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@
3333
the learning rate schedule as a table
3434
================= ============================================
3535
36+
As in nanoGPT, the model and the loss are separate. The model is the
37+
view ``logits``; the loss is SQL that reads it, as nanoGPT's
38+
``F.cross_entropy`` reads the model's logits.
39+
3640
Each training step updates every parameter table with its own
3741
statement; the statements share one backward pass. By default the
3842
optimizer is SGD at a fixed learning rate, the simplest step to read.
@@ -226,42 +230,60 @@ def block(layer: int, x: str, out: str) -> str:
226230

227231

228232
def model() -> str:
229-
"""The CTEs of the whole model, ending in ``loss``: the mean
230-
cross-entropy of predicting each next character of the batch."""
233+
"""The model, as the view ``logits``: for each position of each
234+
sequence in ``batch``, a score for every character that could come
235+
next. A DataFusion view reads the tables that existed when it was
236+
created (DuckDB's looks them up by name), so it is recreated whenever
237+
the batch or the parameters are."""
231238
blocks = "".join(
232239
block(layer, f"x{layer}", f"x{layer + 1}") for layer in range(LAYERS)
233240
)
234241
return f"""
235-
x0 AS (
242+
CREATE OR REPLACE VIEW logits AS
243+
WITH x0 AS (
236244
SELECT b.seq, b.pos, te.e, te.w + pe.w AS val
237245
FROM batch b
238246
JOIN wte te ON b.tok = te.tok
239247
JOIN wpe pe ON b.pos = pe.pos AND te.e = pe.e),
240248
{blocks}
241-
{layer_norm(f"x{LAYERS}", "final", "ln_f")}
242-
logits AS (
243-
SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val
244-
FROM final x JOIN wte te ON x.e = te.e
245-
GROUP BY x.seq, x.pos, te.tok),
249+
{layer_norm(f"x{LAYERS}", "final", "ln_f").rstrip(",")}
250+
SELECT x.seq, x.pos, te.tok, SUM(x.val * te.w) AS val
251+
FROM final x JOIN wte te ON x.e = te.e
252+
GROUP BY x.seq, x.pos, te.tok
253+
"""
254+
255+
256+
# The loss, separate from the model, as in nanoGPT: the mean cross-entropy
257+
# of the model's logits against each next character of the batch.
258+
# TODO(xqlsystems/ddx#142): remove ddx_stop_gradient from the window MAX
259+
# once ddx can differentiate a window aggregate. The max only keeps exp()
260+
# from overflowing, and cancels out of the log-sum-exp, so its gradient is
261+
# zero.
262+
#
263+
# The loss reads the model's view once, into the CTE ``predicted``, and
264+
# uses that twice: a host expands a view at every read, but DuckDB
265+
# computes a CTE that is read twice only once.
266+
LOSS = """
267+
predicted AS (
268+
SELECT seq, pos, tok, val FROM logits),
246269
shifted AS (
247-
-- TODO(xqlsystems/ddx#142): as in the attention weights above.
248270
SELECT seq, pos, tok, val,
249271
ddx_stop_gradient(MAX(val) OVER (PARTITION BY seq, pos)) AS m
250-
FROM logits),
272+
FROM predicted),
251273
log_sum_exp AS (
252274
SELECT seq, pos, MAX(m) + ln(SUM(exp(val - m))) AS lse
253275
FROM shifted GROUP BY seq, pos),
254276
loss AS (
255277
SELECT AVG(s.lse - l.val) AS l
256278
FROM log_sum_exp s
257279
JOIN batch b ON s.seq = b.seq AND s.pos = b.pos
258-
JOIN logits l ON l.seq = b.seq AND l.pos = b.pos AND l.tok = b.next)"""
280+
JOIN predicted l ON l.seq = b.seq AND l.pos = b.pos AND l.tok = b.next)"""
259281

260282

261283
def step_statements(lr: float) -> list[str]:
262284
"""One SGD step: each parameter table's next version, as a statement."""
263285
return [
264-
f"WITH {model()} SELECT {', '.join(dims)}, "
286+
f"WITH {LOSS} SELECT {', '.join(dims)}, "
265287
f"w - {lr} * grad(loss, w) AS w FROM {table} "
266288
f"ORDER BY {', '.join(dims)}"
267289
for table, dims in PARAMS.items()
@@ -273,7 +295,7 @@ def gradient_statements() -> list[str]:
273295
the table ``{name}_grad``'s rows. The statements share one backward
274296
pass."""
275297
return [
276-
f"WITH {model()} SELECT {', '.join(dims)}, grad(loss, w) AS g "
298+
f"WITH {LOSS} SELECT {', '.join(dims)}, grad(loss, w) AS g "
277299
f"FROM {table} ORDER BY {', '.join(dims)}"
278300
for table, dims in PARAMS.items()
279301
]
@@ -316,7 +338,7 @@ def adamw_statement(table: str, step: int) -> str:
316338

317339

318340
def loss_statement() -> str:
319-
return f"WITH {model()} SELECT l FROM loss"
341+
return f"WITH {LOSS} SELECT l FROM loss"
320342

321343

322344
# --------------------------------------------------------------------------
@@ -492,7 +514,7 @@ def main():
492514
report = Report("nanogpt", "batch", args.reps)
493515
text = shakespeare()
494516
vocab = len(set(text))
495-
lines = len(loss_statement().splitlines())
517+
lines = len(model().splitlines()) + len(loss_statement().splitlines())
496518
count = sum(
497519
v.size for v in init_params(vocab, np.random.default_rng(0)).values()
498520
)
@@ -521,6 +543,7 @@ def main():
521543

522544
for step, batch in data:
523545
host.load("batch", batch)
546+
host.execute(model())
524547
tok, nxt = batch["tok"].values, batch["next"].values
525548

526549
rows = report.time(

‎benchmarks/ml/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ array programs.
1818
| # | Script | Model | What it measures |
1919
|---|---|---|---|
2020
| 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 |
21-
| 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 |
21+
| 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 |
2222
| 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 |
2323

2424
The SQL favors legibility over speed. Each model reads like its definition,

0 commit comments

Comments
 (0)