optimizer

Dimitris Papailiopo... (@DimitrisPa... ✓) — 1h btw this is a weird heavy ball SGD variant that basically does this 1. Load previous checkpoint weights. 2. Reset optimizer state / momentum buffers. 3. Train N steps. 4. For first M steps: warm LR from 0.1x -> 1.0x. 5. Hold LR flat until ~50% of the cycle. 6. Linearly decay LR to zero. 7. Save checkpoint. 8. Repeat the optimizer is exactly this buf = mu * buf + grad p *= 1 - lr * wd p -= lr * buf / 524288 @CevherLIONS @_arohan_ does this have a name? I'd call it wave SGD lol

Getting small batch sizes to work in bfloat16 precision can be challenging. In our recent paper on batch size, we ran all experiments in float32, but memory-constrained settings demand lower precision. Here are two tricks that we used to enable bf16 training at small batch sizes: [Chart: "Pretraining 30M model, weights dtype" — FineWeb Edu loss (y-axis, 3.6–5.0) vs Batch size (x-axis, log scale 1–1024). Three lines: BF16 (closest) [gray dashed, spikes badly around batch size 64], BF16 (stochastic) [orange dashed, tracks closely with FP32], FP32 [blue, baseline]. BF16 (closest) diverges sharply upward around batch size 64 while stochastic rounding stays close to FP32 across the whole range.] 4:16 AM · Aug 28, 2025 · 11.7K Views 💬 4 🔁 18 ❤ 156 🔖 110 ⤴ Martin Marek @mrtnm · 20h (1) We recommend using decay rates like b2=0.9999 for small batch sizes. However, bf16 only has ~2.4 decimal points of precision. Since Adafactor's state is so tiny compared to the model size, we can store it in float32 without meaningfully affecting the overall memory footprint.