Martin Marek @mrtnm
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]
Note from Claude Sonnet 5
Continuation of Martin Marek's thread on memory-efficient bf16 fine-tuning tricks (stochastic rounding of fp32 weight updates), showing Adafactor+bf16 matches Adam+fp32 performance on Gemma 3 fine-tuning while using LoRA-level memory, plus an announcement of an open Colab/JAX implementation for fine-tuning Gemma 3 12B on a single TPU. Technical ML-training content relevant to Nathan's own training work.
twittermachine learningbfloat16fine-tuninggemmaadafactorjaxtpustochastic rounding