Skip to content

Commit 5f92e69

Browse files
alxmrsclaude
alxmrs
and
claude
committed
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 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent 21f99d6 commit 5f92e69

3 files changed

Lines changed: 186 additions & 34 deletions

File tree

‎benchmarks/ml/01_attention.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,8 @@
8787

8888
# The gradient of half the output's squared norm, with respect to every
8989
# value of Q, K and V: the same query, plus a loss and a different SELECT.
90+
# The three grad() calls differentiate one loss, so ddx runs one backward
91+
# pass for all three, as it would for three statements given to sql_all.
9092
ATTENTION_WINDOW_GRADIENT = """
9193
WITH scores AS (
9294
SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s

‎benchmarks/ml/02_nanogpt.py‎

Lines changed: 176 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,16 @@
2828
Output layer an einsum with ``wte`` again (tied, as in GPT-2)
2929
Loss log-sum-exp per position, minus the target's logit
3030
Gradients ddx: ``grad(loss, w)``, one statement per table
31-
Optimizer step ``w - lr * grad(loss, w)``, the new table
31+
Optimizer step ``w - lr * grad(loss, w)``, the new table; or AdamW,
32+
with Adam's moments as columns beside ``w`` and
33+
the learning rate schedule as a table
3234
================= ============================================
3335
3436
Each training step updates every parameter table with its own
35-
statement; the statements share one backward pass. The same model in JAX
37+
statement; the statements share one backward pass. By default the
38+
optimizer is SGD at a fixed learning rate, the simplest step to read.
39+
``--optimizer adamw`` trains as nanoGPT does instead: it stores the
40+
gradients as tables, then updates each parameter table with AdamW. The same model in JAX
3641
(Google's library for differentiable array programs), from the same
3742
initial parameters and the same batches, is the reference: its loss at
3843
every step, and its parameters after every step, must match the SQL's.
@@ -53,8 +58,13 @@
5358
* **GELU** is its tanh approximation, as GPT-2 computed it. nanoGPT's
5459
``nn.GELU()`` uses the exact form, through the error function ``erf``,
5560
which DataFusion doesn't have. The two differ by less than 5e-4.
56-
* **The optimizer** is SGD at a fixed learning rate. nanoGPT trains with
57-
AdamW, warmup, cosine decay and gradient clipping.
61+
* **The optimizer**, by default, is SGD at a fixed learning rate.
62+
``--optimizer adamw`` matches nanoGPT's: AdamW with its Shakespeare
63+
settings, weight decay on all but the layer norms, gradients clipped to
64+
a norm of 1, and the learning rate warming up, then decaying along a
65+
cosine. The schedule is stretched to ``--steps``, where nanoGPT's spans
66+
5,000. The JAX reference for this step matches ``torch.optim.AdamW``
67+
after ``clip_grad_norm_`` exactly.
5868
* **No dropout**, nanoGPT's default, so that a run is deterministic and
5969
can be checked. Its Shakespeare configuration trains with 0.2.
6070
* **Q, K and V are three tables.** nanoGPT's ``c_attn`` holds the same
@@ -91,7 +101,18 @@
91101
HEAD = EMBED // HEADS # c: channels per head
92102
HIDDEN = 4 * EMBED # f: channels of the MLP, 4 × n_embd as in nanoGPT
93103
CONTEXT = 16 # block_size = 256; pos: characters per sequence
94-
LEARNING_RATE = 0.5 # for SGD; nanoGPT's AdamW uses 1e-3
104+
LEARNING_RATE = 0.5 # for SGD; nanoGPT's AdamW uses 1e-3, below
105+
106+
# AdamW (``--optimizer adamw``), with the settings of nanoGPT's Shakespeare
107+
# model: train.py's defaults, and config/train_shakespeare_char.py's.
108+
MAX_LR = 1e-3 # learning_rate
109+
MIN_LR = 1e-4 # min_lr
110+
WARMUP = 100 / 5000 # warmup_iters, as a fraction of max_iters
111+
BETA1 = 0.9 # beta1
112+
BETA2 = 0.99 # beta2
113+
EPS = 1e-8 # torch.optim.AdamW's default
114+
WEIGHT_DECAY = 0.1 # weight_decay, on every table but the layer norms' gains
115+
GRAD_CLIP = 1.0 # grad_clip: the largest norm of all gradients together
95116

96117
# Each parameter table's dimensions; its values are the column w. In
97118
# 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]:
247268
]
248269

249270

271+
def gradient_statements() -> list[str]:
272+
"""The gradient of the loss with respect to each parameter table, as
273+
the table ``{name}_grad``'s rows. The statements share one backward
274+
pass."""
275+
return [
276+
f"WITH {model()} SELECT {', '.join(dims)}, grad(loss, w) AS g "
277+
f"FROM {table} ORDER BY {', '.join(dims)}"
278+
for table, dims in PARAMS.items()
279+
]
280+
281+
282+
def adamw_statement(table: str, step: int) -> str:
283+
"""One AdamW step on a parameter table with columns ``w``, ``m`` and
284+
``v`` (the weight and Adam's two moments), from its gradient in
285+
``{table}_grad`` and the learning rate in ``schedule``. This is
286+
``torch.optim.AdamW``, after nanoGPT clips the gradients."""
287+
dims = PARAMS[table]
288+
keys = ", ".join(f"p.{d}" for d in dims)
289+
same_row = " AND ".join(f"p.{d} = g.{d}" for d in dims)
290+
decay = 0.0 if table.startswith("ln_") else WEIGHT_DECAY
291+
every_grad = "\n UNION ALL ".join(
292+
f"SELECT g FROM {t}_grad" for t in PARAMS
293+
)
294+
return f"""
295+
WITH grad_norm AS (
296+
SELECT sqrt(SUM(g * g)) AS norm FROM (
297+
{every_grad})),
298+
clip AS (
299+
SELECT LEAST(1.0, {GRAD_CLIP} / (norm + 1e-6)) AS c FROM grad_norm),
300+
moments AS (
301+
SELECT {keys}, p.w,
302+
{BETA1} * p.m + (1 - {BETA1}) * g.g * clip.c AS m,
303+
{BETA2} * p.v + (1 - {BETA2}) * power(g.g * clip.c, 2) AS v
304+
FROM {table} p
305+
JOIN {table}_grad g ON {same_row}
306+
CROSS JOIN clip)
307+
SELECT {keys},
308+
p.w * (1 - s.lr * {decay})
309+
- s.lr * (p.m / (1 - power({BETA1}, s.step + 1)))
310+
/ (sqrt(p.v / (1 - power({BETA2}, s.step + 1))) + {EPS}) AS w,
311+
p.m,
312+
p.v
313+
FROM moments p JOIN schedule s ON s.step = {step}
314+
ORDER BY {keys}
315+
"""
316+
317+
250318
def loss_statement() -> str:
251319
return f"WITH {model()} SELECT l FROM loss"
252320

@@ -283,6 +351,26 @@ def jax_loss(p, tok, nxt):
283351
return jnp.mean(jax.nn.logsumexp(logits, -1) - target)
284352

285353

354+
def jax_adamw(params, moments, grads, step, lr):
355+
"""The reference AdamW step: new parameters and moments."""
356+
norm = np.sqrt(sum(float(jnp.sum(g**2)) for g in grads.values()))
357+
clip = min(1.0, GRAD_CLIP / (norm + 1e-6))
358+
new_params, new_moments = {}, {}
359+
for name, w in params.items():
360+
g = grads[name] * clip
361+
m, v = moments[name]
362+
m = BETA1 * m + (1 - BETA1) * g
363+
v = BETA2 * v + (1 - BETA2) * g**2
364+
decay = 0.0 if name.startswith("ln_") else WEIGHT_DECAY
365+
m_hat = m / (1 - BETA1 ** (step + 1))
366+
v_hat = v / (1 - BETA2 ** (step + 1))
367+
new_params[name] = w * (1 - lr * decay) - lr * m_hat / (
368+
jnp.sqrt(v_hat) + EPS
369+
)
370+
new_moments[name] = (m, v)
371+
return new_params, new_moments
372+
373+
286374
# --------------------------------------------------------------------------
287375
# Data and parameters, as xarray Datasets.
288376
# --------------------------------------------------------------------------
@@ -348,22 +436,59 @@ def init_params(vocab: int, rng) -> dict[str, np.ndarray]:
348436
return params
349437

350438

351-
def as_table(name: str, values: np.ndarray) -> xr.Dataset:
439+
def as_table(name: str, w: np.ndarray, m=None, v=None) -> xr.Dataset:
440+
"""A parameter as a table: its weights ``w``, and Adam's moments ``m``
441+
and ``v`` when they are given."""
352442
dims = PARAMS[name]
443+
columns = {"w": w, "m": m, "v": v}
353444
return xr.Dataset(
354-
{"w": (dims, values)},
355-
coords={d: np.arange(n) for d, n in zip(dims, values.shape)},
445+
{k: (dims, x) for k, x in columns.items() if x is not None},
446+
coords={d: np.arange(n) for d, n in zip(dims, w.shape)},
356447
)
357448

358449

359-
def as_array(name: str, rows) -> np.ndarray:
360-
return xql.to_dataset(rows, dims=list(PARAMS[name]))["w"].values
450+
def as_array(name: str, rows, column: str = "w") -> np.ndarray:
451+
return xql.to_dataset(rows, dims=list(PARAMS[name]))[column].values
452+
453+
454+
def schedule(steps: int) -> xr.Dataset:
455+
"""nanoGPT's learning rate at each step, as the table
456+
``schedule(step, lr)``: a linear warmup, then a cosine decay to the
457+
minimum, with train.py's ``max_iters`` and ``lr_decay_iters`` both
458+
``steps``."""
459+
warmup = round(WARMUP * steps)
460+
lr = []
461+
for it in range(steps):
462+
if it < warmup:
463+
lr.append(MAX_LR * (it + 1) / (warmup + 1))
464+
else:
465+
decay = (it - warmup) / (steps - warmup)
466+
coeff = 0.5 * (1 + np.cos(np.pi * decay))
467+
lr.append(MIN_LR + coeff * (MAX_LR - MIN_LR))
468+
return xr.Dataset({"lr": ("step", lr)}, coords={"step": np.arange(steps)})
469+
470+
471+
def adamw_step(host, step: int) -> list:
472+
"""One AdamW step on every parameter table: the gradients, in one
473+
backward pass, registered as tables; then each table's update."""
474+
gradients = host.query_all(gradient_statements())
475+
for name, rows in zip(PARAMS, gradients):
476+
host.load(f"{name}_grad", xql.to_dataset(rows, dims=list(PARAMS[name])))
477+
return [host.query(adamw_statement(name, step)) for name in PARAMS]
361478

362479

363480
def main():
364481
parser = cli(__doc__.splitlines()[0], sizes=[4], size_name="batch")
365-
parser.add_argument("--steps", type=int, default=5, help="SGD steps")
482+
parser.add_argument("--steps", type=int, default=5, help="training steps")
483+
parser.add_argument(
484+
"--optimizer",
485+
choices=["sgd", "adamw"],
486+
default="sgd",
487+
help="SGD at a fixed learning rate (default), or nanoGPT's AdamW",
488+
)
366489
args = parser.parse_args()
490+
adamw = args.optimizer == "adamw"
491+
lr = schedule(args.steps)
367492
report = Report("nanogpt", "batch", args.reps)
368493
text = shakespeare()
369494
vocab = len(set(text))
@@ -386,8 +511,13 @@ def main():
386511

387512
for host in hosts(args.hosts):
388513
reference = dict(init)
514+
moments = {
515+
n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items()
516+
}
389517
for name, values in init.items():
390-
host.load(name, as_table(name, values))
518+
state = moments[name] if adamw else ()
519+
host.load(name, as_table(name, values, *state))
520+
host.load("schedule", lr)
391521

392522
for step, batch in data:
393523
host.load("batch", batch)
@@ -407,31 +537,47 @@ def main():
407537
report.note("loss", round(float(want), 6))
408538
report.note("operators", host.operators(loss_statement()))
409539

540+
train = f"train ({args.optimizer}), step {step}"
410541
if not host.differentiates:
411-
report.skip(
412-
f"train, step {step}", host, size, "no ddx backend"
413-
)
542+
report.skip(train, host, size, "no ddx backend")
414543
continue
415-
updated = report.time(
416-
f"train, step {step}",
417-
host,
418-
size,
419-
host.query_all,
420-
step_statements(LEARNING_RATE),
421-
)
422544
grads = jax.grad(jax_loss)(reference, tok, nxt)
545+
if adamw:
546+
updated = report.time(
547+
train, host, size, adamw_step, host, step
548+
)
549+
reference, moments = jax_adamw(
550+
reference, moments, grads, step, lr["lr"].values[step]
551+
)
552+
else:
553+
updated = report.time(
554+
train,
555+
host,
556+
size,
557+
host.query_all,
558+
step_statements(LEARNING_RATE),
559+
)
560+
reference = {
561+
name: w - LEARNING_RATE * grads[name]
562+
for name, w in reference.items()
563+
}
564+
423565
errors = []
424566
for name, rows in zip(PARAMS, updated):
425-
reference[name] = (
426-
reference[name] - LEARNING_RATE * grads[name]
427-
)
428-
trained = as_array(name, rows)
429-
errors.append(
430-
assert_close(
431-
f"{name}, step {step}", trained, reference[name]
567+
expected = {"w": reference[name]}
568+
if adamw:
569+
expected["m"], expected["v"] = moments[name]
570+
for column, values in expected.items():
571+
errors.append(
572+
assert_close(
573+
f"{name}.{column}, step {step}",
574+
as_array(name, rows, column),
575+
values,
576+
)
432577
)
578+
host.load(
579+
name, xql.to_dataset(rows, dims=list(PARAMS[name]))
433580
)
434-
host.load(name, as_table(name, trained))
435581
report.check(max(errors))
436582

437583
report.finish(args.csv)

‎benchmarks/ml/README.md‎

Lines changed: 8 additions & 4 deletions
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, 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 |
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; SGD training steps, 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, 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,
@@ -48,7 +48,8 @@ options are the same in all three:
4848
--csv results.csv # also write the results table
4949
```
5050

51-
Demo 2 adds `--steps`, the number of training steps. Demo 3 adds
51+
Demo 2 adds `--steps`, the number of training steps, and `--optimizer`:
52+
`sgd` (the default), or `adamw` to train as nanoGPT does. Demo 3 adds
5253
`--rounds`, the number of processor rounds.
5354

5455
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.
105106
respect to Q, K and V, in both spellings, match `jax.grad` to 3e-15 of the
106107
largest value at every size. After each of five SGD steps, all eleven of
107108
the GPT's parameter tables match JAX's to 1e-15, and the loss falls from
108-
4.18 to 3.84 (each step on a different batch).
109+
4.18 to 3.84 (each step on a different batch). With AdamW, the weights
110+
and both of Adam's moments match a JAX transcription of
111+
`torch.optim.AdamW` to 6e-15 after each of five steps, and that
112+
transcription matches PyTorch exactly.
109113

110114
**DataFusion computes a CTE once for every place that reads it.** A
111115
model reads its activations more than once: a layer norm reads its input
@@ -165,7 +169,7 @@ SQL gives a number.
165169
- Part B of the NanoGPT demo: the four tricks behind a record of the
166170
[NanoGPT speedrun](https://github.com/KellerJordan/modded-nanogpt), a
167171
public race to train GPT-2 fastest, as a small diff to `02_nanogpt.py`.
168-
It needs nanoGPT's optimizer, AdamW, with its state as versioned tables.
172+
Its sparse optimizer updates build on `--optimizer adamw`.
169173
- Peak memory per host, from a fresh process per measurement.
170174
- An einfold host, once einfold produces plans.
171175

0 commit comments

Comments
 (0)