fix: WanAttentionBlock gradient checkpointing + FSDP2 dtype mismatch

[Jack47]

Summary

• Fix CheckpointError when using torch.utils.checkpoint with FSDP2 mixed precision on WanAttentionBlock
• Replace torch.amp.autocast('cuda', dtype=torch.float32) inside the block with explicit .float() + .to(dtype=compute_dtype), making it autocast-agnostic
• Add gradient checkpointing support to WanModel via _gc_autocast() context_fn

Problem

Under FSDP2 with MixedPrecisionPolicy(param_dtype=bf16), gradient checkpointing recompute runs without the outer bf16 autocast that FSDP2 provides during forward. The inner fp32 autocast in WanAttentionBlock.forward() then produces fp32 tensors during recompute vs bf16 during forward, triggering:

torch.utils.checkpoint.CheckpointError: Checkpoint state_dict has mismatched tensor dtypes

Fix

Replace autocast context managers with explicit dtype casts:

# Before (autocast-dependent, breaks under GC recompute):
with torch.amp.autocast('cuda', dtype=torch.float32):
    e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2)

# After (autocast-agnostic, works with GC):
compute_dtype = x.dtype
e = (self.modulation.float().unsqueeze(0) + e).chunk(6, dim=2)

The fix is numerically equivalent: bitwise identical in fp32, <1e-2 tolerance under bf16 autocast.

Impact

Enables gradient checkpointing for Wan2.2 training, reducing peak memory by 63-74%:

Config (5B)	GC off	GC on	Reduction
37f mbs=2	109 GB	40.7 GB	63%
121f (5s) mbs=2	OOM	51.1 GB	—
121f (5s) mbs=4	OOM	69.8 GB	—

Usage

model = WanModel(...)
model.gradient_checkpointing = True  # enable GC

Fixes #353

[ivanpoh11]

hh