Scaling Laws

Lectures behind it: L09 Scaling laws 1 (opens A3), L11 Scaling laws 2 (closes it) · supporting: L02 resource accounting for C ≈ 6ND, L03 architectures, hyperparameters for which knobs are safe to fix.

A3 is the odd one out. Every other assignment in CS336 hands you a test suite and asks you to make it pass; this one hands you a budget and asks you to make a decision. The course's argument up to here has been that you can predict the cost of a training run from first principles (L02), that the architecture choices are mostly settled (L03), and that the hardware sets the shape of what is efficient (L05–L08). A3 asks the question those chapters were building toward: given a fixed pile of FLOPs, how do you split it between a bigger model and more tokens? What is forbidden here is not a library call — it is training the model you are reasoning about. You never see the 10¹⁹-FLOP run. You get a surrogate, a small allowance of small runs, and a power law you have to trust two orders of magnitude past where you measured it.

Map of the assignment

Two graded problems, 5 + 50 = 55 points. The handout does not print a total, and unlike A1, A2 and A4 there is no tests/ directory, no adapters.py and no pytest run in the repo — the entire spring2025 tree is a README, a pyproject, a reference model.py, one JSON data file and the handout. Everything is graded by a human reading your write-up, plus the accuracy of the configuration you submit to the form.

Setup: environment, data, the API key

Clone the repo and let uv handle the environment. The README is four lines long: install uv, uv add <package> whatever you need, uv run <command> to run in the environment, uv run which python if your editor wants an interpreter path. The pyproject pins Python ≥ 3.12 and declares only requests and torch>=2.7.0 — torch is there purely so the reference model.py imports, not because you train anything. You will want to uv add scipy matplotlib yourself. (Two cosmetic oddities in that file, both harmless: the setuptools stanza still says include = ["cs336_basics"], a leftover from A1, and the CHANGELOG stamps its 1.0.0 and 0.0.0 entries with 2024 dates in a 2025 repo.)

There is no test harness to learn. What there is instead:

• data/isoflops_curves.json — the only input to problem 1. A flat JSON array of run records with exactly three keys each: parameters, compute_budget, final_loss. Note what is not in there: no token count, no hyperparameters, no loss curve. Everything else you need is derived.
• cs336_scaling/model.py — the exact Transformer the training API was fit against, given for reference only. Worth reading before you choose hyperparameters: BasicsTransformerLM.__init__ takes (vocab_size, context_length, d_model, num_layers, num_heads, d_ff, attn_pdrop, residual_pdrop), and get_num_params(non_embedding=True) is the code's own definition of the N in your scaling law — it subtracts both the token and the position embedding tables.
• The training API at hyperturing.stanford.edu:8000, documented in §3.1. Your API key is your SSH public key from the start of the quarter, newlines stripped.

you must be on the Stanford network to query this API, so you may have to use a VPN.A3 handout §3.1 · cs336_spring2025_assignment3_scaling.pdf

If you are working through this outside the course, that is where the assignment stops being fully reproducible: http://hyperturing.stanford.edu:8000/docs is unreachable from the open internet (it times out rather than refusing), and the host is not maintained between quarters. Problem 1 runs fine on the checked-in JSON. Problem 2 you can only do as a design exercise on paper — or against the 2026 restatement of it, which the repo's main branch has since rolled forward to and which prices experiments in B200-hours instead of FLOPs.

1 · Assignment overview

The framing the handout gives you is a role-play: you run training for a lab about to spend a small country's GDP on one model, and you get one shot. Concretely there are two budgets and you should keep them straight from the first minute, because most of the confusion in this assignment comes from conflating them.

A 20% overhead for the science before the run is not an arbitrary number — it is roughly what a real lab budgets for a scaling study, and the reason the assignment feels tight is that it is supposed to. The relationship tying the two together is the one from L02: C ≈ 6ND, compute equals six times parameters times tokens. Every trade-off in this assignment is a movement along that hyperbola. Fix C, and choosing N chooses D.

2 · Scaling laws review

Section 2 is exposition, not work — but it is the exposition that tells you which of Chinchilla's three methods you are being asked to implement, so do not skim it. Hoffmann et al. (2022) fit their law three different ways; the one here is approach 2, the IsoFLOP profile. The other two (fixed-model training-curve envelopes, and a parametric fit of L(N, D)) are not what problem 1 wants, and a solution that quietly implements one of them will not match the deliverable.

2.1 Scaling laws from IsoFLOPs profiles

The mental model is a family of U-shaped curves. Pick a compute budget C. Train several models of different sizes N against that same budget, letting the token count fall out of D = C/(6N) — a small model gets a lot of tokens, a big model gets few. Plot final loss against model size and you get a convex curve: tiny models are underparameterised and plateau high no matter how much data you pour in; enormous models run out of gradient steps before they converge, and in the limit a model so large that one step exhausts C stops at initialisation loss. Somewhere between the two is a minimum, and that minimum is the compute-optimal model size for that budget.

Do this for a ladder of budgets C₁ … C_k and you get a sequence of minima ⟨C_i, N_opt(C_i)⟩, and by C = 6ND a matching ⟨C_i, D_opt(C_i)⟩. Fit power laws N_opt ∝ C^a and D_opt ∝ C^b through those minima, and you can extrapolate to a budget nobody has trained at. That extrapolation — from a handful of small runs to a budget you cannot afford to rehearse — is the entire product.

The handout also makes a deliberate simplification to Hoffmann's recipe, and it is worth knowing you are being let off:

While Hoffmann et al. [2022] fits a quadratic function to each IsoFLOP profile to find its minimum, we instead recommend you simply take the run with the lowest training loss for each compute budget as the minimum.A3 handout, Problem (chinchilla_isoflops) · handout PDF

Problem (chinchilla_isoflops): 5 points

Deliverable: a script that reads data/isoflops_curves.json, plus two plots (model size vs compute, dataset size vs compute) each showing the fitted power law, the data points it was fit from, and the extrapolation out to at least 10²⁴ FLOPs — and, for each plot, a one-sentence answer giving your predicted optimum at 10²³ and 10²⁴ FLOPs.

Mechanically this is an hour of work, and the five points are a reading comprehension check on §2.1. Group the records by compute_budget; take the argmin of final_loss within each group; that record's parameters is N_opt(C_i); derive D_opt(C_i) = C_i/(6·N_opt); fit each of the two point sets to a power law; evaluate the fits at 10²³ and 10²⁴. The shape of the supplied file, so you can check your parse:

Four places people lose time on this, in rough order of frequency. First, fitting in linear space. Handing scipy.optimize.curve_fit a model like N = k * C**a with raw values spanning 10¹⁸ to 10²¹ on one axis and 10⁸ to 10¹⁰ on the other gives you a least-squares problem whose residuals are dominated entirely by the largest budget, and an optimiser that stalls or returns nonsense depending on your initial guess. Take logs first and fit a straight line: log N = a·log C + const. The power law is a line in log-log; fit it as one. curve_fit is what the handout suggests and it is fine, but on log-transformed data a plain least-squares fit is enough.

Second, forgetting that D is derived. There is no token count in the JSON, and the second plot is not a second lookup — it is the first answer pushed through C = 6ND. Which gives you the best free correctness check in the assignment: because D_opt is defined as C/(6·N_opt), the two exponents must satisfy a + b = 1 exactly. If you fit them independently and they do not sum to one, you have a bug — you computed D from something other than C/(6N). For reference, Chinchilla's own reported exponents are a ≈ 0.49 and b ≈ 0.51; a fit on this synthetic set lands in the same neighbourhood of ½, with an R² around 0.98. If yours are far from a half, look for the bug before you look for a story.

Third, treating the argmin as clean. Eight sizes per budget is a coarse grid, and the minima are noisy enough to be non-monotone — in this file the argmin model size at 10²¹ is actually slightly smaller than at 6×10²⁰, which is physically impossible and is pure grid-and-noise artefact. The handout tells you to take the raw min anyway, so take it; but say in your write-up that you noticed, and note that fitting a quadratic in log N per profile (what Hoffmann actually did) is the standard fix, because it interpolates a minimum that need not sit on a grid point. Knowing why the shortcut is a shortcut is most of the five points.

Fourth, over-reading the numbers. This data is synthetic and it is not the same model family as the API in problem 2. Its optimal token-to-parameter ratios sit in the low single digits, nowhere near Chinchilla's famous ≈20:1. Do not carry any absolute number from problem 1 into problem 2 — carry the method.

• data/isoflops_curves.json — the input, 72 records; the handout prints its first two for format.
• scipy.optimize.curve_fit — suggested, not required; uv add scipy matplotlib to get it and a plotting backend.
• Training Compute-Optimal Large Language Models — approach 2, §3.2, is the procedure you are reimplementing.
• L09 Scaling laws 1 — the derivation of the U-curve and why the minimum moves, with the IsoFLOP construction drawn out.

3 · Constructing scaling laws

Now the same method, but you generate the data, and the data costs money. Your target is a compute-optimal configuration at 10¹⁹ FLOPs; your instrument is a GET /loss endpoint that returns the final training loss for an architecture and a FLOP count; your constraint is that the FLOP counts you request are summed and capped.

you may not query more than 2e18 FLOPs worth of experiments for fitting your scaling law. This is hard cap that will be enforced by the API.A3 handout, Problem (scaling_laws) · handout PDF (typo in the original)

Problem (scaling_laws): 50 points

Deliverable: a typeset write-up describing your method for fitting the scaling law and for turning it into a prediction — complete enough to reproduce, with commentary on why you made each design decision — plus a submission to the Google form of (1) predicted optimal model size, (2) the training hyperparameters to use, including a batch size of 128 or 256, and (3) the predicted training loss. Part of the 50 points rides on how well that configuration actually performs.

The handout seeds your write-up with five questions it wants answered: how you decided which runs to query, how you fit the law, how well it fits, what optimum and loss it predicts at 10¹⁹, and what hyperparameters you would use at that size. Treat them as the section headings of the report — they are as close to a rubric as you are given.

The real content of this problem is allocation, and the arithmetic is small enough to do before you write any code. The API's most expensive single run is 10¹⁸ FLOPs, so your entire budget is exactly two of them. Here is what each price tier buys you:

Two consequences fall straight out of that table. The cheap end is effectively free — anything at 10¹⁵ or below costs a rounding error, so there is no excuse for not doing your hyperparameter exploration down there. And the expensive end is where your extrapolation lever lives, because a power law fit through minima at 10¹⁵, 10¹⁶ and 10¹⁷ is reaching two decades to 10¹⁹, while one that includes a 10¹⁸ profile is reaching only one. The tension is that a single 10¹⁸ profile with enough sizes in it to have a credible minimum would blow the whole budget on its own. Resolving that tension, explicitly, is the assignment.

One rule makes planning much more forgiving than it first looks:

Note that querying this endpoint with a previously-queried experimental configuration does not incur extra FLOPs used toward your budget.A3 handout §3.1 · handout PDF

The cap is on distinct configurations, not on HTTP requests. A crashed script that re-runs its first twenty queries costs nothing, and GET /previous_runs gives you back everything you have ever asked for, so you can rebuild your dataset from the server rather than from a local file you forgot to flush. Cache locally anyway — but the failure mode people fear here (losing budget to a bug) is not actually the one that gets them. The one that gets them is a well-organised sweep of the wrong thing.

The competing demands. You are fitting two different objects with one budget. The scaling law N_opt(C) needs several distinct compute tiers, each with several model sizes bracketing a minimum — that is breadth. The hyperparameters at the predicted optimum — learning rate above all, then the depth/width split and head count — need sweeps at fixed size, which is depth. The write-up's fifth question is explicitly about the second, and the form asks for the hyperparameters as a graded field, so a solution that spends 100% of the budget mapping N_opt and then guesses a learning rate for the final configuration has left points on the table. The natural resolution is sequential: use the near-free tiers (10¹³–10¹⁵) to establish a hyperparameter recipe — a learning rate that works, an aspect ratio you will hold fixed, a head count rule — then hold that recipe constant while you spend the expensive tiers on IsoFLOP profiles. Constant-recipe profiles are also what makes the fit interpretable: if the architecture rule changes between tiers, the trend you fit is a mixture of two effects and the extrapolation inherits both.

The search box is small, and its top corner matters. The API bounds d_model to [64, 1024] and num_layers to [2, 24], and the handout gives you the bridge from those knobs to N:

To estimate the number of non-embedding parameters for a given model hyperparameter configuration, use 12·n_layer·d_model².A3 handout, Problem (scaling_laws) · handout PDF

Run that formula against the corners and the box is: minimum 12·2·64² ≈ 98K non-embedding parameters, maximum 12·24·1024² ≈ 302M. Now apply Chinchilla's rule of thumb in the other direction — if the optimum sits near 20 tokens per parameter then C = 6·20·N² and N_opt = √(C/120), which at C = 10¹⁹ is about 289M. That is within a few percent of the box's ceiling. Two things follow. It is a hint that the problem was sized to be solvable inside the allowed ranges. And it is a warning: if your law predicts an optimum near the top of the box, you cannot bracket it from above with a real query, so the claim rests entirely on the fit rather than on a measurement. Say so in the write-up rather than hoping nobody checks. (The corresponding token count, D = C/6N, is ≈5.5×10⁹ tokens at N = 302M — and at context length 512 with batch 128, that is roughly 84,000 optimizer steps, which is the sanity check that the configuration you are reporting is a real training run and not an arithmetic artefact.)

Treat the API as a smooth surrogate, not an oracle. The CHANGELOG records that version 1.0.0 switched the API from a scikit-learn decision-tree regressor to an XGBoost one. Whatever it returns is therefore a fitted model over a grid of real runs, not a fresh training run: broad trends are trustworthy, small differences between neighbouring configurations may be regressor artefacts, and the surface can have kinks no real optimiser would produce. The practical consequence is not to chase third-decimal-place improvements in loss, and to be suspicious of any hyperparameter trend that is within the noise of a flat line — which is exactly the caution L11 draws out of the DeepSeek scaling paper.

The one constraint that is not negotiable in the reported answer:

We place essentially no constraints on the hyperparameters you may report under the FLOPs budget of 1e19, other than the following requirement: your batch size must be either 128 or 256.A3 handout, Problem (scaling_laws) · handout PDF

The stated reason is model FLOPs utilisation — smaller batches would leave the GPUs idle enough that the FLOP accounting stops meaning what it is supposed to mean. If the staff hit OOM running your configuration they will use gradient accumulation or add data-parallel GPUs rather than change your batch size, so you do not need to reason about memory. It is also the only knob where the API's accepted range and the report's allowed range coincide exactly: batch_size is already restricted to {128, 256} on the endpoint.

• Hoffmann et al., Training Compute-Optimal LLMs — the IsoFLOP method and the ≈20 tokens/parameter result the handout expects you to compare against.
• Kaplan et al., Scaling Laws for Neural Language Models — the handout explicitly invites you to borrow from it; its L(N), L(D), L(C) decomposition is an alternative fitting target if you prefer to model loss directly rather than only the minima.
• Yang et al., Tensor Programs V (µP) — the third reference, and the one that speaks to the hyperparameter half of the problem: it is the argument that a learning rate tuned at small width transfers to large width, which is the assumption your cheap-tier sweep is silently making.
• L11 Scaling laws 2 · what you build with this — fit quadratics per compute scale, join the minima, distrust trends that could be flat lines.

3.1 Training API

Three GET endpoints on http://hyperturing.stanford.edu:8000, all taking api_key (your SSH public key, no newlines) as an ordinary query parameter. Interactive docs at /docs. Parameters and their accepted ranges, from the handout:

Two details the handout does not spell out but which fall out of the reference model. num_heads is bounded independently of d_model, but the model splits d_model across heads — its own docstring says d_model must be evenly divisible by num_heads — so a combination like d_model=100, num_heads=16 is nonsense even if the range check lets it through. Pick head counts that divide your width. And train_flops is a discrete ladder, not a continuous knob: your IsoFLOP tiers must come from those sixteen values, which is what makes the budget arithmetic in the table above exact rather than approximate. Before your first real query, hit /total_flops_used as a key sanity check; the handout suggests exactly this.

3.2 Training run details

The runs behind the API were produced with a Transformer close to the one you built in A1, but not identical — and the differences matter, because the loss surface you are fitting is this model's, not yours. Everything below is fixed by the staff; you cannot query any of it.

The one to actually think about is the last row. With no warm-up and a cosine cycle pinned to the total number of steps, the learning rate you request interacts with the run length — a rate that is stable for a long 10¹⁸-FLOP run may diverge on a short 10¹⁴-FLOP one, and vice versa. That is a real reason a learning rate tuned in the cheap tiers may not transfer upward cleanly, and it is the concrete version of the caveat µP exists to address. The dropout is the second: 0.1 on both attention and residual is unusual for a modern pretraining run and it shifts the optimal token-to-parameter ratio somewhat, so do not expect this model family to land exactly on Chinchilla's 20:1.

• cs336_scaling/model.py — the reference implementation; FFN and TransformerBlock are where the deltas above live in code.
• SlimPajama: a 627B-token cleaned and deduplicated version of RedPajama — Cerebras (2023) · the pretraining corpus behind every number the API returns.

What you hand in

Two files to Gradescope and one form. The form is not optional and it is not a formality — it is where the graded prediction lives.

• writeup.pdf — a complete description of your methodology for fitting the scaling law and using it to predict the optimal model size at the given budget, detailed enough to reproduce your results. This carries both problems: the plots and one-sentence answers from chinchilla_isoflops, and the full design narrative for scaling_laws including why you chose the runs you chose.
• code.zip — all the code you wrote to fit the law and compute the estimates. There is no autograder, so this is read, not run; make it legible.
• The Google form — (1) predicted optimal model size, (2) the training hyperparameters to use, with batch size 128 or 256, (3) the model's predicted training loss. Part of your grade is determined by how the reported configuration actually performs. Requires a Google sign-in; the link is the one printed twice in the handout.

There is no compute budget to report and no leaderboard repository for A3 — a check of stanford-cs336/assignment3-scaling-leaderboard returns 404, unlike the A1, A2 and A4 leaderboards. The form is the competitive surface here, and the only enforcement mechanism you will actually feel during the work is the API refusing you once total_flops_used crosses 2×10¹⁸.

Materials, verified

• CS336 Assignment 3 (scaling): Scaling Laws — CS336 Staff (2025), version 1.0.0 · the handout this page walks.
• stanford-cs336/assignment3-scaling @ spring2025 — the pinned tree; main has rolled forward to the 2026 restatement.
• data/isoflops_curves.json — 72 synthetic runs across 9 compute budgets; the sole input to problem 1.
• cs336_scaling/model.py — reference architecture for the API's training runs.
• README.md · pyproject.toml · CHANGELOG.md — setup, dependencies, and the note that the API is an XGBoost regressor.
• Training Compute-Optimal Large Language Models — Hoffmann et al., DeepMind (2022) · the Chinchilla paper; approach 2 is the method problem 1 reimplements and problem 2 extends.
• Scaling Laws for Neural Language Models — Kaplan et al., OpenAI (2020) · cited as an alternative framework you may borrow from; also the origin of the C ≈ 6ND convention the whole assignment runs on.
• Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer — Yang et al. (2022) · the handout's third reference; the theory behind tuning learning rates at small scale and transferring them upward.
• SlimPajama — Cerebras (2023) · the training corpus; the Hugging Face dataset page now requires a signed-in account.
• scipy.optimize.curve_fit — SciPy docs · the fitting routine the handout suggests, with the bounds and p0 arguments you will want if you insist on fitting in linear space.
• uv documentation — Astral · the environment manager the README assumes.
• DeepSeek LLM: Scaling Open-Source Language Models with Longtermism — DeepSeek-AI (2024) · field map extra. The clearest public worked example of exactly this exercise, including scaling laws fit for batch size and learning rate rather than only model size — the closest thing to a model answer for the hyperparameter half of problem 2. Dissected in L11.
• Chinchilla Scaling: A Replication Attempt — Besiroglu et al. (2024) · field map extra. Refits Chinchilla's approach 3 from the published figures and finds the reported confidence intervals implausibly tight — a useful corrective to writing "my law fits well" without an error bar. Companion write-up at Epoch AI.
• Resolving Discrepancies in Compute-Optimal Scaling of Language Models — Porian et al. (2024) · field map extra. Explains why Kaplan and Chinchilla disagreed — last-layer FLOPs, warm-up, and unturned optimizers — which is a direct list of the design decisions that can silently bias your own fit.
• Language models scale reliably with over-training and on downstream tasks — Gadre et al. (2024) · field map extra. What happens when you deliberately leave the compute-optimal point, which is what almost every shipped model does — the practical footnote to this assignment's premise.