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>