Source
Zhao et al. (2022), LIT-083 — ARXIV-2304.11277.
What this adds beyond the mixed-precision practices
The FP16/BF16 recipe in SOTA-014 and SOTA-016 is about the weights and the arithmetic. Under sharding there is a second consumer: every parameter all-gather and gradient reduce-scatter moves bytes over the network, and halving the dtype halves them. On a job whose step time is the collectives — which is the job that reached for FSDP in the first place (SOTA-116) — that is often the larger effect.
Both halves are in LIT-083’s §4.4. Keeping a low-precision copy beside the
full-precision one normally adds memory, but FSDP holds only the local shard
in full precision and gathers the unsharded FlatParameter in low
precision, so the peak term for the largest unit falls from full-precision
to low-precision bytes; and it can run every collective in low precision,
which the paper says saves communication volume. The parameter-side
technique is the standard one, which the paper credits to Micikevicius et
al.; what originates here is its form under sharding. In §5.2 the 11B T5
model reached significantly higher TFLOPS with BF16 on.
The parameter worth naming
The reduction dtype is a separate choice from the compute dtype, and it is where accuracy is actually at risk: gradients are summed across the group, so a low-precision reduction accumulates error proportional to the group size. Keeping the reduce in FP32 while the gather stays in half precision is the usual compromise, and it costs back half the saving on the reduce alone.
That trade is invisible from the title, which is the argument for the practice naming the dtypes it means rather than “mixed precision” as a mode.