accelerate
2f4c053a - FSDP2: move placeholder weights to meta before the fp32 upcast (#4369)

Commit
17 days ago
FSDP2: move placeholder weights to meta before the fp32 upcast (#4369) With cpu_ram_efficient_loading and mixed precision, fsdp2_prepare_model upcasts trainable parameters to fp32 on every rank before sharding. The weights of non-main ranks are later replaced by the main process weights in fsdp2_load_full_state_dict, but the upcast still writes a full fp32 copy into host memory on each rank, about 4 bytes per trainable parameter per process. Capture the non-persistent buffers and move non-main ranks to meta before the upcast, which is then free on those ranks. With Qwen2.5-7B and 4 ranks, host memory after prepare drops from 28.4 GiB to 1.1 GiB on each non-main rank, the main rank is unchanged and the losses are identical. Co-authored-by: Sebastian Pagel <24877965+pagel-s@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Author
Parents
Loading