Tuning Shampoo on a small language model

Tune Meta's distributed_shampoo on a 12M-parameter GPT by editing only the optimizer region of a fixed training script, scored by validation loss on hidden seeds on a line where the shipped configuration is 0 and a tuned recipe is 1.

Final score per graded run

Claude Opus 5
GPT-5.6 Sol
Gemini 3.7 Flash
Kimi K3
Grok 4.6

Configuring an optimizer well matters as much as choosing it, and Shampoo has an above average number of settings to get right: what it preconditions and how often, how it stabilises the inverse roots, grafting, the betas, the decay and the schedule. The best values at any given scale are not written down anywhere, and a weak recipe can make a strong optimizer look ordinary. This task hands the agent a small GPT, a fixed training script, Meta's distributed_shampoo package, an L40S and four hours, and asks how low the validation loss can go in a fixed number of steps when the only thing that changes is how Shampoo is configured.

This task is derived from the NanoGPT speedrun. The model was made small enough that a training run takes only a few minutes on one GPU, so the agent can iterate on many candidate configurations over the course of four hours.

The task

/app/src/train_gpt.py trains a 12M-parameter GPT on FineWeb for 1,500 steps. The model, the data, the batch and the step budget are fixed. Everything between the two EDITABLE REGION markers is the agent's: the parameter initialization, the optimizers and their hyperparameters, and the learning-rate schedule. Nothing outside the markers may change.

As shipped, the region trains the transformer's matrices with DistributedShampoo in a plain configuration and reaches a validation loss of about 4.20. A tuned Shampoo configuration reaches 4.13 in the same steps, and that gap is the task. The hidden matrices must stay on DistributedShampoo with a Shampoo-family preconditioner, so the optimizer cannot be swapped for the one it is being tuned to match, and a small set of rules on what the region may import and touch is checked by validate.py, the same check the grader applies.

Environment

The agent works in an isolated container with no general internet access, on an L40S with four hours on the clock. /app/src is a git repo holding the script, and the instruction asks for one idea per commit: keep the checked-out commit as the best so far, reset a failed idea, and let the log read as a staircase of accepted improvements.

Three tools do the work. train checks the rules, runs the script once with a chosen seed, streams the validation loss every 250 steps, and ends with the final loss and its score on the task's line; it appends every attempt, including the failures, to results.tsv and to the experiment journal. A run takes about four minutes as shipped and longer if Shampoo preconditions more often. submit validates the file, seals it with a snapshot of the source tree, and sends it for grading. clock reports elapsed and remaining time. Local runs and scored submissions are both unlimited.

README.md77 lines

# Working in this environment

Everything lives under `/app`:

```
/app/src/train_gpt.py the training script; the EDITABLE REGION is yours, the rest is frozen
/app/task/ validate.py (the rules), metric.py (the score), spec.json, the pristine train_gpt.py
/app/tools/ train, submit, clock
/app/docs/ the distributed_shampoo README
/app/experiments/ your lab notebook lives here (history.jsonl)
/app/submissions/ your submission record, managed by `submit`; never edit by hand
/opt/data/fineweb16k/ the data (read-only); tokenizer.json and manifest.json describe it
```

One GPU, PyTorch and `distributed_shampoo` are preinstalled. There is no network access beyond
the leaderboard.

## Training

```
/app/tools/train --label freq1-clipped --seed 0 > run.log 2>&1 &
grep "^dev_score:" run.log
```

`train` checks the rules, runs `train_gpt.py` with the seed, streams its output (`step:N/T
val_loss:<x>` every 250 steps), and ends with `dev_loss:<x>` and `dev_score:<y>`; it appends one
line to `/app/src/results.tsv` and `/app/experiments/history.jsonl`. Plain `SEED=1 python
train_gpt.py` works too. `MICRO_BATCH` (default 32) only changes memory, never results. The seed
selects the initialization; the grader uses seeds you do not know, so what counts is a
configuration that works across seeds, not one lucky run.

## The rules

`python3 /app/task/validate.py src/train_gpt.py` says whether a file obeys them; the grader applies
exactly the same check, so a file that passes here cannot be rejected for structural reasons.
In short: the frozen regions must be untouched; inside the editable region only `torch`,
`distributed_shampoo` and `math` may be imported; the data, model, step budget and seed may not be
referenced or rebound; no classes, attribute assignment, or in-place tensor ops inside functions;
and the hidden matrices stay on `DistributedShampoo` with a Shampoo-family preconditioner.

## Scoring

`/app/task/metric.py` is the exact scoring code: `score(val_loss)` puts a validation loss on a
straight line where the shipped configuration is 0 and the tuned reference is 1, with no ceiling.
The grader averages the loss over its seeds before scoring.

## Submitting

```
/app/tools/submit --file src/train_gpt.py --code src --name clipped-eig-freq1 --notes "lr 0.01, wd 0.1"
```

- Validates first (a rejection costs nothing), then seals the submission and queues it; the
leaderboard trains it on one hidden seed and returns the **public** score, and on two more for
the **private** score you never see. `submit --wait <id>` blocks until the grade is in,
`submit --list` shows all your submissions.
- `--code` is required: that directory is zipped and sealed alongside the file for
reproducibility. It is archived, not executed.
- Submissions are not rationed; submit whenever your best improves. Your best public submission
is the one judged on the private seeds.

## Experiment discipline

`/app/src` is a git repo. One idea per commit; keep the checked-out commit your best-so-far;
`git reset --hard HEAD~1` rejects a failed idea. Append one line per run to
`/app/experiments/history.jsonl`:

```json
{"exp": "freq1-clipped", "hypothesis": "eigenvalue clipping lets lr go up", "seed": 1, "val_loss": 3.41, "verdict": "keep"}
```

## Time

`/app/tools/clock` prints elapsed and remaining time (a safe lower bound). Runs take minutes, not
seconds: launch them in the background and poll the log. Stop a background run by its pid, never
with `pkill -f`: your own session carries this text in its arguments.

Figure 1. The agent's workspace at /app, with the files it is given. src/train_gpt.py starts as a copy of the pristine script in a fresh git repo; the data under /opt/data/fineweb16k ships in the container and is omitted here, and the hidden seeds exist nowhere in the environment. The task instruction is delivered as the agent's prompt rather than as a file, and is on the Prompt tab.

The editable region is the whole of the agent's deliverable, so it is worth seeing exactly where it sits.

⋮188 lines, frozen: constants, data, model
189torch.manual_seed(seed)
190print(f"seed:{seed} layers:{NUM_LAYERS} dim:{MODEL_DIM} steps:{train_steps}")
191 
192 
193# ===================== BEGIN EDITABLE REGION =====================
194 
195from distributed_shampoo import (
196    AdamPreconditionerConfig,
197    DistributedShampoo,
198    SingleDeviceDistributedConfig,
199    WeightDecayType,
200)
201 
202# initialize model parameters
203for name, p in model.named_parameters():
204    w = p.data
205    if name.endswith("weight"):
206        if "proj" in name:
207            w.zero_()
208        elif "embed" in name:
209            w.normal_()
210        else:
211            w.normal_(std=0.33**0.5 / w.size(-1)**0.5)
212    elif name.endswith("bias"):
213        w.zero_()
214    elif name.endswith("gains"):
215        w.normal_(mean=1, std=0)
216    else:
217        raise Exception(f"Uninitialized parameter: {name}")
218 
219# create the optimizer(s)
220optimizer1 = AdamW([dict(params=[model.embed.weight], lr=0.3),
221                    dict(params=[model.proj.weight], lr=1/320),
222                    dict(params=[p for p in model.parameters() if p.ndim < 2], lr=0.01)],
223                   betas=(0.8, 0.95), eps=1e-10, weight_decay=0, fused=True)
224shampoo_beta2 = 0.95
225optimizer2 = DistributedShampoo(
226    [p for p in model.blocks.parameters() if p.ndim >= 2],
227    lr=0.0015,
228    betas=(0.9, shampoo_beta2),
229    weight_decay=0.2,
230    weight_decay_type=WeightDecayType.DECOUPLED,
231    max_preconditioner_dim=8192,
232    precondition_frequency=5,
233    start_preconditioning_step=-1,
234    grafting_config=AdamPreconditionerConfig(
235        beta2=shampoo_beta2,
236        epsilon=1e-10,
237    ),
238    distributed_config=SingleDeviceDistributedConfig(),
239)
240optimizers = [optimizer1, optimizer2]
241for opt in optimizers:
242    for group in opt.param_groups:
243        group["initial_lr"] = group["lr"]
244 
245# learning rate schedule: stable then decay
246def set_hparams(step, cooldown_frac=0.7):
247    progress = step / train_steps
248    assert 0 <= progress < 1
249    if progress < 1 - cooldown_frac:
250        eta = 1.0
251    else:
252        eta = (1 - progress) / cooldown_frac
253    for opt in optimizers:
254        for group in opt.param_groups:
255            group["lr"] = group["initial_lr"] * eta
256 
257# ====================== END EDITABLE REGION ======================
258 
259from distributed_shampoo import DistributedShampoo as _Shampoo
260assert set(p for opt in optimizers for group in opt.param_groups
261           for p in group["params"]) == set(model.parameters()), "every parameter must belong to exactly one optimizer"
262_hidden = set(p for p in model.blocks.parameters() if p.ndim >= 2)
263for opt in optimizers:
264    _params = set(p for group in opt.param_groups for p in group["params"])
265    assert isinstance(opt, _Shampoo) or not (_params & _hidden), "hidden matrices must be optimized only by DistributedShampoo"
266    _hidden -= _params if isinstance(opt, _Shampoo) else set()
267assert not _hidden, "hidden matrices must be optimized by DistributedShampoo"
268assert all("Shampoo" in type(group["preconditioner_config"]).__name__
269           for opt in optimizers if isinstance(opt, _Shampoo) for group in opt.param_groups), \
270    "DistributedShampoo must use a Shampoo-family preconditioner"
271optimizers = tuple(optimizers)
⋮42 lines, frozen: training loop, validation, the final loss

Figure 2. train_gpt.py as shipped, 313 lines. The highlighted 63 lines between the markers are the whole of what the agent may change: the initialization, the two optimizers and the learning-rate schedule. The frozen tail below the end marker checks that every parameter belongs to exactly one optimizer and that the hidden matrices are trained by DistributedShampoo with a Shampoo-family preconditioner; the frozen head above holds the constants, the data loader and the model.

Evaluation

A submission is graded by training it on hidden seeds and reading the validation loss after the last step. That loss is placed on a straight line by two anchors. Writing $L$ for the loss, $L_0 = 4.2039$ for the shipped configuration and $L_1 = 4.1331$ for the tuned reference, both means over five seeds measured at calibration,

s = \frac{L_0 - L}{L_0 - L_1}.

The shipped configuration scores 0, the reference scores 1, and the line carries on past both: nothing is clipped. Three hidden seeds are used. The public score the agent sees is the score of the loss on one of them. The private score it never sees is the score of the mean loss over the other two, and it decides the result. The grader averages the losses before scoring, so seed noise sits in the loss rather than in the scale. A run that crashes, diverges or exceeds 25 minutes on a seed scores $-1$ for that split.

The gap between the two anchors is 0.071 nats, which is 5.6 standard deviations of the shipped configuration's spread across seeds, so a score of 1 is a real result rather than a lucky draw. For orientation, one seed each at calibration: Muon on the hidden matrices lands at 1.11 on this line, eigenvalue-corrected Shampoo, usually called SOAP, at 1.18, and everything on AdamW at $-1.83$ . Neither Muon nor AdamW is a legal submission; they say what the scale means.