feat(fsdp): meta-init + rank0-streamed weight loading (--fsdp-load-mode stream)

[zhihengy]

Problem

The FSDP init path materializes the full pipeline on every rank (DiffusionPipeline.from_pretrained), moves the whole unsharded model to GPU, and only then calls fully_shard:

• peak GPU = full model, not shard: one Wan2.2-A14B expert in fp32 master dtype is ~54 GB before training starts — OOM on anything smaller than ~60 GB free
• every rank reads the checkpoint from storage (N× disk traffic; on multi-node, × nodes again)
• there is no way to opt into a meta-device init: diffusers' pipeline loader ignores/crashes under an ambient init context (NotImplementedError: Cannot copy out of meta tensor)

Fix: --fsdp-load-mode stream (new default)

Per component, invert "load-then-shard" into "shard-then-load":

1. Build on meta from model_index.json + component from_config under accelerate.init_empty_weights(include_buffers=False) — params on meta (zero memory), buffers/plain tensor attrs computed for real by __init__ (same context diffusers' own from_pretrained uses, so values match bit-exactly).
2. fully_shard while still on meta — sharding plan needs only shapes, costs nothing.
3. to_empty allocates only this rank's shards. Non-persistent buffers (Wan's rope.freqs_cos/sin — not in any checkpoint) are snapshotted and restored across the transition.
4. Stream weights: rank0 reads safetensors file-by-file (mmap; peak CPU ≈ one ~5 GB file) and set_model_state_dict(..., broadcast_from_rank0=True, strict=False) scatters each tensor into the shards over NVLink. A coverage check makes strict=False safe: any parameter not covered by the checkpoint raises.
5. LoRA: adapters are created on meta too (get_peft_model(..., low_cpu_mem_usage=True)), base weights land on the peft-wrapped FQNs via a key map (exact inverse of the stripping in diffusion_update_weight_utils), and A/B initialize after materialization via reset_lora_parameters on the sharded DTensors. Data-aware inits (pissa/olora/...) need real base weights at wrap time and are rejected with a pointer to legacy mode.
6. Family hook: new postprocess_model_after_materialize (default: run the existing before-fsdp hook) so qwen-image's rope-parity rebuild sees real CUDA tensors — under stream its device.type != "cuda" guard would otherwise silently no-op on meta and lose train/rollout bit-parity.

--fsdp-load-mode legacy keeps the previous behavior (custom/remote-code pipelines, data-aware LoRA inits).

Measurements (2×H200, GPFS)

Real code path, Wan2.2-T2V-A14B, both experts, fp32 master:

	legacy	stream
init wall-clock	285.3 s	71.3 s
peak GPU / rank	80.5 GB	54.0 GB (= the sharded params themselves)
checkpoint readers	every rank	rank0 only

Standalone cold-cache single-expert comparison (fair A/B on untouched file sets): 94.9 s → 65.6 s, peak 53.9 → 27.3 GB. With more ranks the gap widens: legacy peak stays ≈ full-component size while stream peak shrinks with world size.

Correctness

• Streamed weights verified bitwise-equal to from_pretrained on Wan2.1-T2V-1.3B via the real backend path: 825/825 params (gathered via full_tensor()) + both non-persistent rope buffers.
• New 2-GPU CI test (tests/fast-gpu/test_fsdp_stream_load_equivalence.py): self-contained tiny sharded Wan checkpoint; base + LoRA variants; asserts param/buffer bitwise equality and LoRA A-init/B-zero.
• New CPU unit tests (tests/fast/.../test_stream_load.py): meta build contract, non-diffusers component rejection, buffer survival across to_empty, key-map inversion, data-aware-init rejection. tests/fast/backends/fsdp_utils: 26 passed.
• Checkpoint resume unaffected: DCP load runs after init and overwrites; LoRA resume (lora_only, strict=False) actually requires base weights materialized from hf_checkpoint — satisfied by both modes.

Known limitations

• qwen-image + --fsdp-cpu-offload + stream: the rope rebuild's CUDA-device guard no-ops when params materialize on CPU (same class of issue existed pre-PR only if the model never hit CUDA; flagging for a follow-up).
• Remote-code pipelines (trust_remote_code) can't be class-resolved from model_index.json → clear error pointing at legacy mode.
• set_model_state_dict broadcast is per-tensor; further wins available (pinned staging, batched broadcast) but init is already 4× faster.

🤖 Generated with Claude Code