Step 4 of 28Skill: Model Training Fundamentals

← #2.1 Architecture Tweaks#2.3 Dataloading Tweaks →

TRAINING SIMULATION

Speedrun — #2.2 Muon Optimizer

running

◷train_gpt2.py▸Muon class◎learner⌁03-muon＄no GPU

baseline

GPT-2 d12

Starting architecture and training loop.

step 01

active

optimize

Speedrun step

Architecture diffs and optimizer changes.

steps 02–05

target

val loss 3.28

Cumulative loss timeline and finale.

step 06

Blocks

Quick summary

Orthogonalized momentum updates via Newton–Schulz iterations on gradient matrices.

Full explanation below the code →

fba-lab — train_gpt2.py · Muon classexecuting

// block: Muon class · lines 70–123$ study train_gpt2.py --block muon_classOrthogonalized momentum updates via Newton–Schulz iterations on gradient matrices. ✓

Explanation

Orthogonalized momentum updates via Newton–Schulz iterations on gradient matrices.

Think about

What does zeropower_via_newtonschulz5 compute?

// architecture

Live diagram

What changed vs 02-architectural

+ # Muon optimizer
+ # Reference: https://kellerjordan.github.io/posts/muon/
+ 
+ 
+ def zeropower_via_svd(G, steps=None):
+     U, S, V = G.svd()
+     return U @ V.T
+ 
+ 
+ @torch.compile
+ def zeropower_via_newtonschulz5(G, steps=10, eps=1e-7):
+     """
+     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+     quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+     of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+     zero even beyond the point where the iteration no longer converges all the way to one everywhere
+     on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+     where S' is diagonal with S_{ii}' \\sim Uniform(0.5, 1.5), which turns out not to hurt model
+     performance at all relative to UV^T, where USV^T = G is the SVD.
+     """
+     assert len(G.shape) == 2
+     a, b, c = (3.4445, -4.7750, 2.0315)
+     X = G.bfloat16() / (G.norm() + eps)  # ensure top singular value <= 1
+     if G.size(0) > G.size(1):
+         X = X.T
+     for _ in range(steps):
+         A = X @ X.T
+         B = A @ X
+         X = a * X + b * B + c * A @ B
+     if G.size(0) > G.size(1):
+         X = X.T
+     return X.to(G.dtype)
+ 
+ 
+ zeropower_backends = dict(svd=zeropower_via_svd, newtonschulz5=zeropower_via_newtonschulz5)
+ 
+ 
+ class Muon(torch.optim.Optimizer):
+     """
+     Muon: MomentUm Orthogonalized by Newton-schulz
+ 
+     Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+     processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+     matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has
+     the advantage that it can be stably run in bfloat16 on the GPU.
+ 
+     Some warnings:
+     - This optimizer assumes that all parameters passed in are 2D.
+     - It should not be used for the embedding layer, the final fully connected layer, or any {0,1}-D
+     parameters; those should all be optimized by a standard method (e.g., AdamW).
+     - To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions.
+     - We believe it is unlikely to work well for training with small batch size.
+     - We believe it may not work well for finetuning pretrained models, but we haven't tested this.
+     - We have not yet tried this optimizer for training scenarios larger than NanoGPT (124M).
+ 
+     Arguments:
+         lr: The learning rate used by the internal SGD.
+         momentum: The momentum used by the internal SGD.
+         nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
+         backend: The chosen backend for the orthogonalization step. (recommended: 'newtonschulz5')
+         backend_steps: The number of iteration steps to use in the backend, if it is iterative.
+     """
+ 
+     def __init__(self, params, lr=3e-4, momentum=0.95, nesterov=True, backend='newtonschulz5', backend_steps=5):
+         defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, backend=backend, backend_steps=backend_steps)
+         super().__init__(params, defaults)
+ 
+     def step(self):
+         for group in self.param_groups:
+             lr = group['lr']
+             momentum = group['momentum']
+             zeropower_backend = zeropower_backends[group['backend']]
+             for p in group['params']:
+                 g = p.grad
+                 if g is None:
+                     continue
+                 state = self.state[p]
+                 if 'momentum_buffer' not in state:
+                     state['momentum_buffer'] = torch.zeros_like(g)
+                 buf = state['momentum_buffer']
+                 buf.mul_(momentum).add_(g)
+                 if group['nesterov']:
+                     g = g.add(buf, alpha=momentum)
+                 if g.size(0) == 3 * g.size(1):  # split grouped QKV parameters
+                     g = torch.cat([zeropower_backend(g1, steps=group['backend_steps']) for g1 in g.split(g.size(1))])
+                     scale = g.size(1) ** 0.5
+                 else:
+                     g = zeropower_backend(g, steps=group['backend_steps'])
+                     scale = max(g.size(0), g.size(1)) ** 0.5  # scale to have update.square().mean() == 1
+                 p.data.add_(g, alpha=-lr * scale)
+ 
+ 
+ # -----------------------------------------------------------------------------
- @dataclass
- class GPTConfig:

What it bought

Learn this step, your way

WatchReadDeep Dive

🎬 Video lesson coming soon

This step's video hasn't been recorded yet. The text lesson and Rune's deep-dives cover the same ground in the meantime.

Lab

Speedrun — #2.2 Muon Optimizer

Live diagram

What changed vs 02-architectural

Muon optimizer + Newton–Schulz

AdamW + Muon split

Higher LR, no warmup

What it bought

Learn this step, your way

Speedrun — #2.2 Muon Optimizer

Live diagram

What changed vs 02-architectural

Muon optimizer + Newton–Schulz

AdamW + Muon split

Higher LR, no warmup

What it bought

Learn this step, your way