Hyperparameter Selection › Worked Example: Hyperparameter Search for a 350M Model

Michael BrenndoerferJanuary 29, 20262 min read
Language AI HandbookAll chapters
Part 32 of 67Training Optimization
Chapter 238Hyperparameter Selection › Worked Example: Hyperparameter Search for a 350M Model
Reading section 10 of 14 · Worked Example: Hyperparameter Search for a 350M Model · 2 minSection 10 of 14 · 2 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 238 of 513

Section 10 of 14 · 2 min read

Hyperparameter Selection › Worked Example: Hyperparameter Search for a 350M ModelOne section · Sections in this chapter10 / 14 ↓
Sections in this chapter14
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

To illustrate the full workflow, consider setting up training for a 350M parameter model with a 100B token dataset.

Step 1: Set architecture-determined hyperparameters.

For a 350M model, a typical architecture is 24 layers, hidden dimension 1024, 16 attention heads, FFN ratio 4. These are fixed by the scaling recipe and not searched.

Step 2: Set the compute-optimal training length.

Using the Chinchilla recipe, a 350M model should train on roughly 20×350×106=7×10920 \times 350 \times 10^6 = 7 \times 10^9 tokens. With a batch size of 0.5M tokens per step, that is 7×109/(5×105)=14,0007 \times 10^9 / (5 \times 10^5) = 14{,}000 steps. We will use 15,000 steps rounded up.

Step 3: Proxy search on a 40M model.

Build a 40M parameter proxy using the same architecture ratios (fewer layers, smaller hidden dimension). Train it for 2,000 steps (proportional to the proxy's compute-optimal budget) at 7 different learning rates:

lr∈{3×10−5, 1×10−4, 3×10−4, 6×10−4, 1×10−3, 3×10−3, 1×10−2}\text{lr} \in \{3 \times 10^{-5},\ 1 \times 10^{-4},\ 3 \times 10^{-4},\ 6 \times 10^{-4},\ 1 \times 10^{-3},\ 3 \times 10^{-3},\ 1 \times 10^{-2}\}

Evaluate validation loss at the end of each run. Identify the best configuration and the second-best; the optimal is likely between these two.

Step 4: Refine on the proxy.

Run a finer search between the two best values. With 5 additional evaluations, you can narrow down to within 20% of the true optimum.

Step 5: Transfer to the 350M target.

If using muP, transfer the proxy learning rate directly. If using standard parametrization, scale by the width ratio: lr350M=lr40M×d40M/d350M\text{lr}_{350\text{M}} = \text{lr}_{40\text{M}} \times \sqrt{d_{40\text{M}} / d_{350\text{M}}} approximately.

Step 6: Run and monitor.

During the full training run, monitor:

  • Training loss curve: should decrease smoothly with no spikes
  • Gradient norm: should stabilize around 1.0-3.0 after warmup
  • Learning rate schedule: confirm the schedule matches expectations
  • Validation loss: check against expected scaling law prediction

If validation loss diverges from prediction by more than 10%, investigate before continuing. Common causes are suboptimal learning rate, data quality issues, or hardware bugs.