Fix CPU-offload tensor placement in the FSDP2 full-state loader

[hjinnkim]

Fix CPU-offload tensor placement in the FSDP2 full-state loader

Summary

• Problem: the SimpleTuner loader override accepted neither Accelerate
  1.14.0's explicit cpu_offload call nor legacy three-argument calls safely
  under an active offload policy.
• Impact: Accelerator.prepare() could reject the keyword, or loaded tensors
  could be placed on the wrong device.
• Fix: accept a tri-state placement input and apply it consistently to every
  loaded tensor.

Changes

• patch_fsdp2_state_dict_loader().replacement(): add
  cpu_offload: Optional[bool] = None so current and legacy call forms work.
• Placement resolution: explicit True selects CPU, explicit False selects the
  accelerator device, and None infers placement from the active
  CPUOffloadPolicy.
• Sharded and unsharded load paths: apply the resolved device on main and
  non-main ranks before load_state_dict(assign=True).
• RAM-efficient loading: keep loader memory optimization independent from the
  FSDP2 CPU-offload policy.

Result

• Accelerate 1.14.0 calls and legacy three-argument calls both work.
• Main/non-main and sharded/unsharded paths follow one placement contract.
• Example: with an active CPUOffloadPolicy, an omitted value now infers CPU;
  explicit True selects CPU and explicit False keeps tensors on the
  accelerator. Before the fix, the explicit keyword could raise during
  Accelerator.prepare().
• Verified: 12 targeted cases passed. Two-GPU NCCL validation covered
  RAM-efficient on/off, CPU offload on/off, and mixed DTensor/unsharded state;
  all enabled cases placed tensors on CPU and completed a finite optimizer step.

Environment

• Python: 3.13.14
• PyTorch: 2.13.0+cu130
• Accelerate: 1.14.0
• Diffusers: 0.39.0
• Transformers: 5.13.1
• GPU: NVIDIA B200, driver 590.48.01

[Copilot]

Pull request overview

This PR updates SimpleTuner’s override of Accelerate’s FSDP2 full-state-dict loader to accept Accelerate 1.14.0’s cpu_offload keyword (while keeping legacy call compatibility) and to move loaded shards to CPU when CPU offload is in effect.

Changes:

• Extends the loader override signature with cpu_offload: Optional[bool] = None to support both keyword and legacy call forms.
• Infers cpu_offload from the active FSDP CPU offload policy when the argument is omitted.
• Ensures loaded shards are moved to CPU when cpu_offload resolves to True in both main and non-main rank paths.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.