Hyperparameter Selection › Implementing a Hyperparameter Search Workflow

Michael BrenndoerferJanuary 29, 20264 min read
Language AI HandbookAll chapters
Part 32 of 67Training Optimization
Chapter 238Hyperparameter Selection › Implementing a Hyperparameter Search Workflow
Reading section 9 of 14 · Implementing a Hyperparameter Search Workflow · 4 minSection 9 of 14 · 4 min

Book contents

Section 9 of 14 · 4 min read

Hyperparameter Selection › Implementing a Hyperparameter Search WorkflowOne section · Sections in this chapter9 / 14 ↓
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

A complete hyperparameter search workflow for a new LLM training run combines the elements above into a practical sequence:

In[18]:
Code
from dataclasses import dataclass
from typing import List

import numpy as np


@dataclass
class HPConfig:
    """Complete hyperparameter configuration for an LLM training run."""

    # Optimizer
    learning_rate: float
    beta1: float = 0.9
    beta2: float = 0.95
    epsilon: float = 1e-8
    weight_decay: float = 0.1

    # Schedule
    warmup_steps: int = 2000
    total_steps: int = 100000
    min_lr_ratio: float = 0.1  # min_lr = peak_lr * min_lr_ratio

    # Gradient clipping
    max_grad_norm: float = 1.0

    # Architecture-linked (not searched)
    model_size_billions: float = 1.3

    def validate(self) -> List[str]:
        """Check configuration for common mistakes."""
        warnings = []

        # LR sanity check
        size_params = self.model_size_billions * 1e9
        expected_lr = 1.2e-3 * (size_params**-0.07)  # rough power law
        lr_ratio = self.learning_rate / expected_lr
        if lr_ratio > 5:
            warnings.append(
                f"LR {self.learning_rate:.2e} is {lr_ratio:.1f}x above size-expected {expected_lr:.2e}"
            )
        elif lr_ratio < 0.2:
            warnings.append(
                f"LR {self.learning_rate:.2e} is {1 / lr_ratio:.1f}x below size-expected {expected_lr:.2e}"
            )

        # Warmup as fraction of total
        warmup_frac = self.warmup_steps / self.total_steps
        if warmup_frac < 0.005:
            warnings.append(
                f"Warmup fraction {warmup_frac:.3f} is very short; recommend at least 0.01"
            )
        elif warmup_frac > 0.10:
            warnings.append(
                f"Warmup fraction {warmup_frac:.3f} is unusually long; typical is 0.01-0.05"
            )

        # Weight decay range
        if self.weight_decay < 0.001:
            warnings.append(
                f"Weight decay {self.weight_decay} is very small; may cause overfitting"
            )
        elif self.weight_decay > 0.5:
            warnings.append(
                f"Weight decay {self.weight_decay} is very large; may hurt convergence"
            )

        return warnings


def cosine_lr_schedule(step: int, config: HPConfig) -> float:
    """Compute learning rate at a given step."""
    if step < config.warmup_steps:
        return config.learning_rate * (step / config.warmup_steps)

    progress = (step - config.warmup_steps) / (
        config.total_steps - config.warmup_steps
    )
    cosine_factor = 0.5 * (1 + np.cos(np.pi * progress))
    min_lr = config.learning_rate * config.min_lr_ratio
    return min_lr + (config.learning_rate - min_lr) * cosine_factor


# Example: configure a 1.3B model run
config_1b = HPConfig(
    learning_rate=3e-4,
    model_size_billions=1.3,
    warmup_steps=2000,
    total_steps=100000,
)

# Test LR schedule across training
steps = np.arange(0, config_1b.total_steps + 1, 100)
lr_values = [cosine_lr_schedule(s, config_1b) for s in steps]

warnings_1b = config_1b.validate()
Out[19]:
Console
HPConfig for 1.3B model:
  Learning rate:   3.00e-04
  Weight decay:    0.1
  Beta1, Beta2:    0.9, 0.95
  Warmup steps:    2000 (2.0% of total)
  Total steps:     100,000
  Min LR:          3.00e-05

Configuration: no warnings.

LR at step 0:                    0.00e+00
LR at end of warmup (2000): 3.00e-04
LR at 50% of training:           1.69e-04
LR at 90% of training:           3.69e-05
LR at final step:                3.00e-05

The configuration object makes hyperparameter relationships explicit and enables validation of common mistakes before training begins. The schedule shows the warmup phase rising linearly to the peak, followed by cosine decay to 10% of the peak value.

Key Parameters

The primary hyperparameters for an LLM training run and their recommended defaults are:

  • learning_rate: The most important parameter. Start with the size-based prior from Table lr-by-size; search within a factor of 3x around this prior.
  • beta2: Use 0.950.95 for large-scale training; 0.9990.999 for smaller models. This is robust within each setting.
  • weight_decay: Default 0.10.1 for pretraining. Lower (0.010.01) for data-rich scenarios; higher (0.30.3) for fine-tuning on small datasets.
  • warmup_steps: 1-4% of total steps, or a fixed 2000 steps for long runs. Very robust to exact value within this range.
  • max_grad_norm: 1.0. Monitor actual gradient norms during training; if they regularly exceed 2x the clip threshold, investigate instability.
  • min_lr_ratio: 0.1 (cosine decays to 10% of peak). Some recipes decay to near-zero; this is acceptable but slightly more aggressive.