DeepSpeed was secretly upcasting to FP32 — here's why your FSDP loss diverged

Running the same training pipeline with DeepSpeed and PyTorch FSDP should produce similar results. But when Hugging Face engineers tried it with Mistral-7B on four A100 GPUs, the loss curves told a very different story. DeepSpeed’s loss converged cleanly. FSDP’s loss flatlined. The culprit wasn’t the sharding strategy or some subtle distributed training bug. It was precision.

DeepSpeed’s Zero3 optimizer has always maintained master weights in FP32, silently upcasting from bfloat16 before the optimizer ever touches the parameters. FSDP, by contrast, flattens parameters and lets them stay in whatever precision you loaded them in. So the PyTorch optimizers were operating directly on bfloat16 weights — and at certain learning rates, that’s a recipe for stagnation. The team initially suspected learning rate scaling issues tied to GPU count, tried a 4x bump, and saw improvement. But the real explanation was sitting in DeepSpeed’s codebase, in a function called _create_fp32_partitions.

The fix landed in Hugging Face Accelerate 0.30.0. FSDP now has two modes: a mixed-precision mode that mirrors DeepSpeed’s upcasting behavior, and a low-precision mode for memory-constrained setups where every gigabyte counts. That flexibility is actually an advantage over DeepSpeed, which forces the upcast whether you want it or not. On a small GPU count, DeepSpeed’s 2x memory hit from FP32 master weights can sting.

Throughput benchmarks with IBM’s Granite 7B on four A100s show the two frameworks performing comparably once precision is aligned — ballpark MFU and tokens/sec/GPU figures are in the same range. The team also published a migration guide covering sharding equivalence, model loading, weight prefetching, and wrapping strategies. If you’ve ever wondered why your FSDP run felt off compared to a DeepSpeed baseline, the answer might be simpler than you think: check your dtypes.