feat(diffusion): support Wan2.2 TI2V Flow-GRPO training

[wwwsctvcom]

Add Wan2.2-TI2V Flow-GRPO Training with FSDP and LoRA

Summary

This PR adds a general Wan2.2-TI2V-5B training recipe based on the existing Miles training framework.

Supported features:

• Flow-GRPO training
• FSDP distributed training
• LoRA fine-tuning
• SGLang TI2V rollout
• Image-conditioned video generation
• PickScore frame-based reward
• Optional evaluation
• Checkpoint saving and resuming through native training arguments

Main Changes

• Add the wan2_2_ti2v training pipeline configuration.

• Add TI2V masked timestep construction for video latent tokens.

• Add the Wan2.2-TI2V FSDP precision plan.

• Forward conditioning image paths from dataset metadata to SGLang rollout.

• Parse and preserve TI2V rollout metadata required by the training forward pass.

• Add a complete training script:

  scripts/run_diffusion_grpo_wan22_ti2v.py

• Add a regression test for TI2V mask normalization and timestep shape handling.

Dataset Format

The training data uses the existing Miles metadata format:

{
"input": "A camera moves around the subject.",
"metadata": {
"image_path": "/data/images/frame.jpg"
}
}

The image path must be available to all Ray workers.

Training Example

python3 scripts/run_diffusion_grpo_wan22_ti2v.py
--model /models/Wan2.2-TI2V-5B-Diffusers
--data-jsonl /data/train.jsonl
--num-gpus 2
--cuda-visible-devices 0,1

The script enables the following training components by default:

• FSDP training backend
• LoRA adapter training
• LoRA weight synchronization
• Flow-SDE rollout
• GRPO advantage estimation
• PickScore reward
• SGLang rollout
• Checkpoint saving

Evaluation Example

python3 scripts/run_diffusion_grpo_wan22_ti2v.py
--model /models/Wan2.2-TI2V-5B-Diffusers
--data-jsonl /data/train.jsonl
--eval-data-jsonl /data/validation.jsonl
--eval-interval 20
--num-gpus 2

Validation Results

Unit and configuration tests:

7 passed
git diff --check passed

The complete training recipe was also passed through the native train_diffusion.py argument validation.

A real two-GPU execution validated the following path:

• SGLang tensor parallelism: TP=2
• FSDP data parallelism: DP=2
• FSDP-wrapped Wan transformer blocks: 30
• LoRA trainable parameters: 10,076,160 / 5,009,863,872
• LoRA synchronization: 600 tensors across 300 layer prefixes
• GRPO reward calculation
• GRPO advantage calculation
• Flow-GRPO loss
• Backward pass
• Optimizer update
• Distributed checkpoint saving

The real execution also confirmed that:

• Standard CUDA IPC LoRA synchronization works.
• All 600 LoRA tensors were synchronized successfully.
• No layer prefixes were unmapped.
• Image-conditioned Wan2.2-TI2V rollout completed successfully.
• Generated videos were decoded successfully.
• FSDP initialization completed with DP=2.

Scope

This PR only contains the general Wan2.2-TI2V training support and its required tests.

It does not include:

• Environment-specific CPU weight-sync fallback code
• Smoke-only reward implementations
• Fixed local model paths
• Fixed local datasets
• Unrelated H3 or Wan2.1 validation code

The existing training paths for other model families are unchanged.

[Rockdu]

Hi, thanks for your PR, do you have a curve for verification?