Add xpk recipes for training of WAN 2.2 T2V 14b for Ironwood (TPU 7x)

[Toshi-31]

Summary

Adds the XPK pretraining recipe and documentation for Wan 2.2 T2V 14B (720x1280x81) on Google Cloud TPU v7x (Ironwood) in training/ironwood/wan2.2-t2v-14b/bf16-tpu7x-2x4x4/xpk/ and updates the summary table in training/ironwood/README.md.

This recipe incorporates multi-dimensional FSDP+CP parallelism, custom Pallas splash attention kernels with tuned tiles (block_q=2816, block_kv_compute_in=4096 forward; block_kv_dkv=4096, dq_reduction_steps=3 fused backward), chunked Ulysses context parallelism (ulysses_attention_chunks=2), and selective attention output rematerialization, achieving 12.15 s / step (402.67 TFLOP/s/device, 40.43% MFU vs the 996 TFLOP/s/device BF16 peak) on tpu7x-2x4x4.

Companion PR is open in AI-Hypercomputer/maxdiffusion: AI-Hypercomputer/maxdiffusion#487.

---

Optimization Progression (Wan 2.2 T2V 14B on TPU v7x)

MFU is reported against the TPU v7x BF16 peak of 996 TFLOP/s per device. Model FLOPs: 4,892.47 TFLOP/device/step; step time is the mean of steady-state steps 2–10.

Optimization Stage	Step Time	TFLOP/s / Device	MFU (% of 996 TFLOP/s Peak)
Unoptimized Baseline	39.70 s	123.24	12.37%
Selective Activation Remat & 1D FSDP Mesh	27.63 s	177.07	17.78%
Overlapped Ring Attention & 2048 Compute Tiles	21.82 s	224.22	22.51%
Unsegmented Attention & All-Reduce Elimination	19.37 s	252.58	25.36%
Forward Rematerialization Elimination & Chunked Ulysses	15.02 s	325.73	32.70%
Custom Splash Attention Backward Kernel	13.57 s	360.46	36.19%
Attention Tile Tuning + pl.CostEstimate	12.15 s	402.67	40.43%

• Hardware: tpu7x-2x4x4 (64 TensorCore devices, 32 chips, 8 hosts)
• Workload: Wan 2.2 T2V 14B (720x1280x81, batch size per device = 0.25, ici_fsdp=16, ici_cp=4, ici_tp=1)
• GCS Output Directory: gs://wan2-2-training/toshipahadia-wan22-aishared-outputs/tp-wan22-v7a-e4tilescost-0930-1400/
• XProf Profile: xprof.corp.google.com/trace_viewer/toshipahadia-4840309546544091074

[gemini-code-assist]

Code Review

This pull request introduces a new training recipe for running the Wan 2.2 Text-to-Video 14B model on Ironwood GKE clusters using XPK, including documentation and a run script. Feedback on the run script suggests critical improvements: changing the Hugging Face cache directory from '/dev/shm' to '/tmp/hf_hub_cache' to prevent out-of-space errors on GKE, adding 'set -eo pipefail' for robust error handling, validating required environment variables before execution, and using '--no-deps --no-index' with 'pip install' to ensure compatibility with private or air-gapped clusters.