1 of 34

HW5 Recitation

Fall 2026 10-405/605/805

2 of 34

Context: Hyperparameter Tuning

Hyperparameter tuning is the process of optimizing the hyperparameters of a machine learning model to improve its performance

Random Search

Grid Search

Adaptive Selection

3 of 34

Context: LLM sizes

Would you run grid search in this case?

4 of 34

Context: GPT2

In 2019, OpenAI released GPT 2, trained on a dataset containing 8 million web-pages with the objective of predicting the next word of a given sentence.

​

​

​

​

​

​

​

​

​

​

​

HW5 version (preliminary): Between GPT-2 Small and Medium (~180M)

5 of 34

Aside: GPT2 Training Data

Source: Language Models are Unsupervised Multitask Learners (2019)

6 of 34

Context: HellaSwag Evaluation

HellaSwag is a benchmark designed to test if large language models (LLMs) possess commonsense reasoning abilities.

Humans easily achieve over 95% accuracy on these questions

OpenAI’s GPT-2 model only reached around 30% accuracy

7 of 34

Goals of HW5

  • Find the “max learning rate” that achieves minimal validation loss and maximal HellaSwag accuracy for a fixed training objective
  • Pretrain from scratch a (prospective) 181M GPT-2 model on 4 billion tokens of training data from the FineWeb dataset
  • Side Note: Can anyone guess how we came up with these numbers?

​

8 of 34

Goals of HW5

  • Find the “max learning rate” that achieves minimal validation loss and maximal HellaSwag accuracy for a fixed training objective
  • Pretrain from scratch a (prospective) 181M GPT-2 model on 4 billion tokens of training data from the FineWeb dataset
  • Side Note: Can anyone guess how we came up with these numbers?

​

9 of 34

Goals of HW5

  • Find the “max learning rate” that achieves minimal validation loss and maximal HellaSwag accuracy for a fixed training objective
  • Pretrain from scratch a (prospective) 181M GPT-2 model on 4 billion tokens of training data from the FineWeb dataset
  • Side Note: Can anyone guess how we came up with these numbers?
  • Answer: “Training Compute Optimal Large Language Models”

​

​

10 of 34

Context: Max learning rate

In this graph, the max learning rate is the peak value that the learning rate reaches during training:

​

  • Linear Warm-Up Phase: The learning rate starts at zero and increases linearly until reaching the max learning rate. This gradual rise stabilizes training early on.

  • Cosine Decay Phase After reaching the max learning rate, the learning rate decreases smoothly towards zero, helping the model converge gradually.

​

​

11 of 34

Context: Effective batch size

- Micro-Batch Size: The number of samples processed in a single forward-backward pass on one GPU.

Example: micro batch size of 64 means 64 samples per GPU per step.

​

- Effective Batch Size The total batch size affecting model updates, calculated as:

Effective Batch Size = Micro-Batch Size X Gradient Accumulation Steps

​

- Gradient Accumulation If memory is limited, gradients are accumulated over several micro-batches before updating the model weights.

​

Example: With 8 GPUs and 64 micro batch size, the total effective batch size is 8x64x(maximum sequence length) tokens. If using 1 GPU, gradient accumulation (8 steps) is used to reach the same effective batch size.

​

Key takeaway: Micro-batch size affects memory usage, while effective batch size controls the training update.

​

​

12 of 34

Recommended approach

(1) Scaling down part of the training process,

(2) Tuning hyperparameters on this smaller setup, and

(3) Transferring the findings to estimate optimal hyperparameters for the full-scale model.

​

​

​

​

13 of 34

Properties that might be scaled down

  1. The number of training iterations
  2. The training dataset size
  3. The model size (total number of parameters)
  4. The model depth (number of layers)
  5. The number of attention heads
  6. The embedding dimension

14 of 34

HW5 Logistics

Part 1. Experimental overview and experiments. Deliverables:

(a) Completed experiential design.

(b) Completed cost analysis.

(c) Completed hyperparameter tuning table.

(d) Completed empirical questions.

(e) Final selection of “max learning rate” hyperparameter.

​

Budget: $20+ (make sure to leave $20 for part 2!!!)

You may use extra if you have more in your AWS account

15 of 34

Starter Strategies

​

16 of 34

Scaling training iterations

Grid/random sample learning rates over a small range of values

17 of 34

Heuristic Model Scaling

Make a plot of optimal learning rate as a function of model size, and extrapolate the results.

​

Find the optimal learning rate for a set of models that are much smaller than the one you hope to test (GPT2-Tiny has ~30M params)

​

​

18 of 34

Example Models:

​

Tiniest ~3M

  • 1 layer
  • 64 channels
  • 1 attention head

Tinier ~7M

  • 2 layers
  • 128 channels
  • 2 attention heads

Tiny ~ 11M

  • 3 layers
  • 196 channels
  • 3 attention heads

Strategy overview:

​

  • Use random/grid search to train each model 10 times with different learning rates.

​

  • Create a plot of model size vs. optimal learning rate, and extrapolate to the target model size of 181M parameters.

​

  • Select the learning rate projected to be optimal for a model of 181M parameters for your final training run.

19 of 34

Questions you’ll probably run into

  • If I train each model 10 times, that’s 30 experiments. How long can each experiment be if I only have 20$ of credit?

​

  • What are some initial guesses for what optimal lr is for each model size?

​

  • How can I run 30 experiments efficiently? (Bash/Shell scripts!)

​

  • What if I run 10 experiments and notice that the optimal lr is probably higher/lower than the range of lrs that I tested?

​

  • Will I need to iterate on my experiments?

​

20 of 34

Hyperband

Recall from lecture

Train n = 9 models in the first round

Train each model to epoch 9 in the first round

21 of 34

Hyperband

Recall from lecture

Train n = 3 models in the second round

Train each model to epoch 27

22 of 34

Hyperband

Recall from lecture

Train n = 1 models in the final round

Train the model to epoch 81

23 of 34

Hyperband

Now consider the full Hyperband table. Repeat the same procedure, but s+1 times with slightly different values of n and r

Variables n and r have the same meaning as last slide

In this bracket, train 5 models for 81 epochs

24 of 34

Hyperband

Here is the Hyperband table we gave you in the homework.

​

Hints: Flags -n, -y, and -e may be helpful to continue training a model

25 of 34

Hyperband

Here is the pseudocode.

You will likely need to implement “get_hyperparameter_configuration”, “run_then_return_val_loss”, and “top_k” yourself. You may be able to use an off-the-shelf version of Hyperband.

​

26 of 34

Scaling Batch Size

It turns out that scaling to the optimal batch size and fitting learning rate as function of batch size might be a good strategy

​

We are still fleshing out the details on this strategy, look to the HW5 write-up for more details

Source: https://arxiv.org/pdf/2505.13738

27 of 34

Homework 5 Part 2

28 of 34

Objective: Use your hyperparameters to pretrain an LLM from scratch

Part 2 Deliverables:

(a) Final HellaSwag accuracy

(b) Final validation loss

(c) Pre-trained model weights

​

Budget: $15-20

29 of 34

Background: Embarrassingly Parallel Training

​

Typical multi-GPU training strategies require complex synchronization schemes

30 of 34

Background: Embarrassingly Parallel Training

31 of 34

Background: Branch-Train-Merge

32 of 34

Background: Branch-Train-Merge

We have trained the seed LM for you

33 of 34

Background: Branch-Train-Merge

We have trained the seed LM for you

Each of you will be training an expert LM on the domains you collected in HW4

34 of 34

Background: Branch-Train-Merge

​

With approx 100B tokens of data (2B tokens per student with 50 students) our class should be able to train a single massive language mode that is close to optimal