Week 9 — Optimizers and Large-Scale Pretraining

Optimization · Pretraining · Muon

Advanced

III — Training the model

Overview

The report does not use one optimizer uniformly: AdamW for normalization/non-matrix parameters, Muon for many matrix parameters, head-wise Muon for Q/K, and Sinkhorn-balanced updates for embeddings and prediction heads. The regime: 45T multimodal tokens, sparse attention from the start at 64K, extended to 1M at 34T, final 7:1 text:multimodal mix.

What You Will Learn

• Explain why matrix, norm and embedding parameters want different optimizers.
• Derive AdamW's decoupled decay and Muon's matrix-preconditioned update.
• Implement Sinkhorn row/column normalization.
• Read a large-scale pretraining recipe.

Core Concepts

AdamW decoupled decay

\theta\leftarrow \theta-\eta\,(\hat m/(\sqrt{\hat v}+\epsilon))-\eta\lambda\theta

Sinkhorn balancing

Alternate row and column normalization to balance an embedding/prediction update matrix.

The regime

45T multimodal tokens; sparse attention from the start at 64K; extended to 1M at 34T tokens; final corpus is 7:1 text:multimodal.

Prerequisites

Work these pages on this site before the lecture.

• Adam
• AdamW
• Warmup + decay LR
• Adaptive learning rates
• Jacobian & Hessian

Lecture notes

One optimizer for every parameter is a convenience, not a law. The report’s split:

• AdamW — norms, biases, non-matrix params (decoupled decay).
• Muon — many matrices (orthogonalized / Newton-Schulz-style preconditioning).
• Head-wise Muon — Q/K.
• Sinkhorn — embedding and lm-head updates (alternate row/column normalize).

Regime to remember: 45T multimodal tokens, sparse attention from the start at 64K, extend to 1M at 34T, finish 7:1 text:multimodal.

Required readings

• DeepSeek-V4.1-Flash report § training
• Muon (Keller Jordan)
• AdamW

Related on this site

• Assignment A3
• Adaptive rates

Lab / Implementation

Implement Sinkhorn-style row/column normalization and inspect update RMS; compare Adam-style, row-normalized and row+column-normalized updates on a large embedding matrix (Assignment 3).

Not on this site (paper reading required)

• Muon
• head-wise Muon
• the 45T/64K→1M/7:1 regime (DeepSeek §3–4)

Mastery Check

• DERIVE: AdamW decay and the Sinkhorn row/column update
• IMPLEMENT: AdamW + Sinkhorn-balanced embedding update
• BENCHMARK: update RMS; Adam vs row vs row+col
• DEBUG: weight decay applied inside the gradient
• EXPLAIN: different optimizers per parameter class

Part of CS/AI 684 — Efficient Long-Context Agent Systems.