adafactor

Martin Marek @mrtnm · 20h (2) Instead of directly updating model weights in bf16, we compute updated weights in fp32, then stochastically round to bf16 for storage. This means we can accumulate many small gradient steps without introducing bias. 💬 1 🔁 ❤ 13 📊 446 ⤴ Martin Marek @mrtnm · 20h After applying these two tricks to our fine-tuning experiment, Adafactor with bf16 weights still matches the baseline performance of Adam with fp32 weights but crucially its memory footprint is similar to LoRA (with bf16 weights). [Chart: "Gemma 3 (4B) fine-tuning" — MATH score (y-axis, 17%-19%) across four bar conditions: LoRA BS=1 bf16 (~16.9%), Adafactor BS=1 bf16 (~18.4%), Adam BS=1 fp32 (~18.6%), Adam BS=16 fp32 (~18.2%), with error bars.] 💬 1 🔁 ❤ 11 📊 496 ⤴ Martin Marek @mrtnm · 20h We updated our codebase with a Colab notebook to finetune Gemma 3 (12B) using a TPU v6e-1 with just 32 GB of memory. We implemented everything from scratch in JAX, including sampling! We also updated our paper to be more explicit about [cut off]

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.

[Top, cut off tweet]: ...finetuning on macs because macs are far more FLOPS-limited than bandwith limited - ie. there isn't as much gain from bs=4 throughput wise. Thus bs=1 isn't as horrible a tradeoff as it would be on 8xH100. 💬1 🔁 ♥2 📊118 🔗 N8 Programs @N8Programs · 4h Thus my recommendation: next time you finetune an LLM with MLX, start by trying AdaFactor BS=1, LR 1e-3, decay_rate 0.997, scale_parameter: true. Can be added to mlx-lm with literally a single line of code. Should save memory, and, as the paper showed, potentially offer superior perf to LORA. 💬1 🔁 ♥1 📊287 🔗 N8 Programs @N8Programs · 4h Note: this replication is designed to explore the impact of BS=1 and adafactor vs adam. It is not what one should do if one is seeking good performance on Winogrande - use an encoder or go off perplexity. Training on verified gpt-4.1 rollouts for this task is designed to emulate Show more 💬 🔁 ♥1 📊276 🔗 Martin Marek @mrtnm · 3h Thank you so much for reproducing our results so quickly! 💬 🔁1 ♥1 📊259 🔗