Assignment 3 · Scaling

What you build

This assignment is unlike the others: there is almost no infrastructure to implement, and no test suite to pass. The deliverable is judgment. The setting: you are training ClosedAI’s next-generation model on a fixed compute budget of 48 B200-hours, and you must commit — in advance — to the model size and hyperparameters you believe will attain the lowest validation loss. To inform that bet, you get a scaling-laws budget of an additional 12 B200-hours (25% of the big run’s FLOPs budget) to spend on smaller experiments, queried through a training API that runs real jobs on B200 GPUs.

Two problems, very unevenly weighted:

• chinchilla_isoflops (5 points) — a warm-up on synthetic data. Reproduce the Chinchilla IsoFLOPs method from Hoffmann et al. on the runs given in data/isoflops_curves.json, fit power laws $N_{\text{opt}} \propto C^a$ and $D_{\text{opt}} \propto C^b$, and extrapolate the compute-optimal model and dataset sizes to $10^{23}$ and $10^{24}$ FLOPs. Deliverables are two plots plus one-sentence predictions.
• scaling_laws (50 points) — the real thing. Design your own suite of training runs, fit your own scaling law, and predict the optimal configuration and its final validation loss for the 48 B200-hour budget.

You submit writeup.pdf (a methodology description detailed enough to reproduce your results) and code.zip to Gradescope, and your predicted optimal hyperparameters plus predicted final loss to the API. Part of your grade is determined by how well your predicted model actually performs.

The training API, and what your budget actually buys

The API accepts a full training configuration — architecture (hidden_size, num_hidden_layers, num_attention_heads, and friends, on the Assignment 1 architecture with slight differences), optimizer (AdamW or SGD with a warmup–cosine schedule), batch size, total_train_tokens, a model_seed, and a max_runtime_seconds between 1 second and 12 hours — and returns the sequence of validation losses from actually training that model on tokenized DCLM data (32K vocabulary, context length 512, fixed deterministic data order, up to 500B tokens).

The budget mechanics reward planning. A queued run reserves its full max_runtime_seconds against your 43,200-second budget; if it finishes early, the unused time is refunded, but a run that times out is charged in full. Submitting an identical configuration twice returns a 409 and costs nothing. Once the budget is exhausted, the API refuses further requests — a hard cap, not a suggestion. Endpoints cover the whole loop: POST /submit to queue, GET /budget to check your balance, GET /experiment/1 to poll until completion, and POST /final_submission (resubmittable; latest wins) for your final answer. You need to be on the Stanford network, and your API key is your student ID.

Which lectures feed it

This is the assignment for the scaling-laws stretch of the course (Lectures 9–12): how loss relates to model size, data, and compute; the Chinchilla IsoFLOPs methodology; and how hyperparameters should move as models grow. The handout’s own reading list is the right rewatch guide — Kaplan et al. (2020) and Hoffmann et al. (2022) for the two canonical ways of fitting the compute trade-off, and Yang et al.’s Tensor Programs V (muP) if you’re deciding how learning rate and width should co-vary. When stuck on the warm-up, go back to the IsoFLOPs discussion: for each compute budget $C$, train varying model sizes $N$ with $D = C/6N$, find the loss-minimizing $N$ per budget, and fit power laws through those minima. (The handout suggests simply taking the lowest-loss run per budget rather than fitting a quadratic to each profile, and points at scipy.optimize.curve_fit for the fitting itself.)

Order of attack

1. Do the warm-up first, completely. It is the dress rehearsal for the method at zero API cost — every fitting decision you make there transfers to the real problem.
2. Plan your run suite before touching the API. Decide which compute budgets, which model sizes per budget, and which hyperparameters you’ll hold fixed versus vary. The handout’s estimate for non-embedding parameters, $12 \, n_{\text{layer}} d_{\text{model}}^2$, is what connects a hyperparameter configuration to a point on your scaling law.
3. Spend the budget deliberately. Reserve tight max_runtime_seconds (refunds only come from runs that finish), and remember failed timeouts return partial losses but are charged in full.
4. Write as you go. The handout lists the questions your writeup must answer — how you chose runs, how you fit the law, how well it fits, what it predicts, and which hyperparameters you’d use at the predicted size. That list is effectively the rubric; design decisions are much easier to document when they’re fresh.

The hard part is not the curve fitting — it’s experiment design under a budget that punishes waste, and having the discipline to commit to a prediction (including the predicted loss itself) that the leaderboard will then test against reality.