perf(fsdp): accelerate pinned sleep and wake

[Rockdu]

What

• Move complete padded FSDP2 shard storage directly between CUDA and pinned host memory.
• Rebuild the DTensor wrapper through FSDP's native to_sharded_dtensor() helper and atomically swap it into the existing Parameter identity.
• Offload optimizer state and diffusion EMA shadows to pinned memory, and discard gradients after the optimizer step.
• Add a two-GPU regression covering uneven shards, padding, DTensor layout, repeated sleep/wake, sleep-time state dict/EMA access, and training after wakeup.

Why

The previous sleep path moved the model to pageable CPU memory. This made wakeup synchronous and slow. Pinning after model.cpu() is also insufficient because it performs an extra full pageable-to-pinned copy and only exposes the logical DTensor shard, which can omit FSDP's physical padding.

This change transfers FSDPParam._sharded_param_data directly, preserves the logical DTensor view over the complete padded shard, and enables non-blocking host-to-device wakeup. It does not mutate DTensor's private _local_tensor; the wrapper and physical storage therefore keep consistent device metadata while preserving the original parameter object identity.

Validation

Settings:

Item	Value
Model	Qwen-Image 20.8B
GPUs	4x H200
Training precision	FP32 master, BF16 forward
Adaptation	LoRA rank 64
Per-rank model storage	19.42 GiB
Per-rank optimizer storage	0.79 GiB

Results use the slowest rank:

Transition	Pageable baseline	Pinned physical storage	Speedup
First sleep + wake	15.31 s	13.24 s	1.16x
Steady-state sleep	9.95 s	0.84 s	11.9x
Steady-state wake	5.37 s	0.64 s	8.35x
Steady-state sleep + wake	15.31 s	1.48 s	10.4x

Sleep releases CUDA allocations for the measured model and optimizer storage to 0 GiB.

Tests:

• pytest -q tests/fast-gpu/backends/fsdp_utils/test_sleep_wakeup.py
• pytest -q tests/fast/backends/fsdp_utils/test_model_backend.py tests/fast/backends/fsdp_utils/test_scaler_bf16_passthrough.py
• Ruff and Black checks for all changed files

Files

• miles/backends/fsdp_utils/actor.py: pinned physical-shard and optimizer sleep/wake.
• miles/backends/fsdp_utils/ema.py: EMA device and pinned-memory migration.
• tests/fast-gpu/backends/fsdp_utils/_sleep_wakeup_worker.py: distributed lifecycle regression worker.
• tests/fast-gpu/backends/fsdp_utils/test_sleep_wakeup.py: GPU CI entry point.

Checklist

• pre-commit run --all-files passes — not run
• Added/updated tests for new behaviour
• pytest -x is green — targeted FSDP tests pass; full suite not run
• If launch flags changed, python3 train.py --help still parses — N/A, no launch flags changed
• If a public flag was added, it appears in the CLI reference docs — N/A, no public flag added
• If an example was added, it has a real walkthrough — N/A, no example added