diff --git a/examples/fastvideo/predict_t2v_minimax_h3.py b/examples/fastvideo/predict_t2v_minimax_h3.py index c13866e..f0d412b 100644 --- a/examples/fastvideo/predict_t2v_minimax_h3.py +++ b/examples/fastvideo/predict_t2v_minimax_h3.py @@ -40,9 +40,10 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode, # GPU at full load, so keep an offloading mode here. GPU_memory_mode = "model_group_offload" # Multi GPUs config -# NOTE: the VSA sparse kernel has no sequence-parallel path, so a FastVideo-H3 student runs -# single-GPU only; keep ulysses_degree = ring_degree = 1 (enable_multi_gpus_inference() raises -# otherwise, rather than silently falling back to the dense multi-GPU processor). +# The FastVideo-H3 VSA student supports sequence-parallel inference through ulysses head-parallel +# only (the VSA kernel has no ring-attention path): put every GPU into ulysses_degree, keep +# ring_degree = 1, and ulysses_degree must divide the 56 attention heads (2 / 4 / 7 / 8 / ...). +# The defaults stay single-GPU because the ~19B stack still needs an offloading mode below. ulysses_degree = 1 ring_degree = 1 # Use FSDP to save more GPU memory in multi gpus. Unused at ulysses_degree = ring_degree = 1. @@ -247,8 +248,14 @@ if use_qfloat8: wants_offload = GPU_memory_mode in ("model_cpu_offload", "model_group_offload", "sequential_cpu_offload") if ulysses_degree > 1 or ring_degree > 1: + # The VSA sparse kernel only supports ulysses head-parallel (no ring-attention path); a VSA + # student under ring_degree>1 would silently fold ring into a larger ulysses degree and give + # none of ring's per-rank memory saving, so reject it instead of mis-attending. + if ring_degree > 1: + raise ValueError( + "FastVideo-H3 VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree " + "instead (ulysses_degree must divide the 56 attention heads).") from functools import partial - # NOTE: raises for the VSA student (no sequence-parallel path); keep ulysses_degree=ring_degree=1. transformer.enable_multi_gpus_inference() if fsdp_dit: fp32_modules = [m for m in transformer.modules() diff --git a/examples/fastvideo/predict_t2v_wan2.1.py b/examples/fastvideo/predict_t2v_wan2.1.py index 6a0d872..1d37815 100644 --- a/examples/fastvideo/predict_t2v_wan2.1.py +++ b/examples/fastvideo/predict_t2v_wan2.1.py @@ -239,6 +239,12 @@ pipeline = WanPipeline( scheduler=scheduler, ) if ulysses_degree > 1 or ring_degree > 1: + # The VSA kernel has no ring-attention path; a sparse student under ring_degree>1 would + # silently mis-attend. student_topk >= 1.0 is the dense fallback, which ring does support. + if ring_degree > 1 and student_topk < 1.0: + raise ValueError( + "FastVideo VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree " + "instead (num_heads must divide ulysses_degree).") from functools import partial transformer.enable_multi_gpus_inference() if fsdp_dit: diff --git a/examples/sparkdiffusion/predict_t2v_wan2.1.py b/examples/sparkdiffusion/predict_t2v_wan2.1.py index 29c323c..47214bd 100644 --- a/examples/sparkdiffusion/predict_t2v_wan2.1.py +++ b/examples/sparkdiffusion/predict_t2v_wan2.1.py @@ -295,6 +295,7 @@ def save_results(): else: video_path = os.path.join(save_path, prefix + ".mp4") save_videos_grid(sample, video_path, fps=fps) + print(f"Saved video to: {video_path}") if ulysses_degree * ring_degree > 1: import torch.distributed as dist diff --git a/examples/turbodiffusion/predict_i2v_wan2.2.py b/examples/turbodiffusion/predict_i2v_wan2.2.py index 984290c..e3dbac1 100644 --- a/examples/turbodiffusion/predict_i2v_wan2.2.py +++ b/examples/turbodiffusion/predict_i2v_wan2.2.py @@ -213,6 +213,12 @@ pipeline = Wan2_2I2VPipeline( scheduler=scheduler, ) if ulysses_degree > 1 or ring_degree > 1: + # The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a + # sparse expert under ring_degree>1 would silently mis-attend. + if ring_degree > 1: + raise ValueError( + "TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into " + "ulysses_degree instead (num_heads must divide ulysses_degree).") from functools import partial transformer.enable_multi_gpus_inference() transformer_2.enable_multi_gpus_inference() @@ -289,7 +295,6 @@ def save_results(): image = (image * 255).numpy().astype(np.uint8) image = Image.fromarray(image) image.save(video_path) - print(f"Saved image to: {video_path}") else: video_path = os.path.join(save_path, prefix + ".mp4") save_videos_grid(sample, video_path, fps=fps) diff --git a/examples/turbodiffusion/predict_t2v_wan2.1.py b/examples/turbodiffusion/predict_t2v_wan2.1.py index c8a1615..61f5a5c 100644 --- a/examples/turbodiffusion/predict_t2v_wan2.1.py +++ b/examples/turbodiffusion/predict_t2v_wan2.1.py @@ -67,7 +67,7 @@ sampler_name = "Flow" # Used when the sampler is in "Flow_Unipc", "Flow_DPM++". # If you want to generate a 480p video, it is recommended to set the shift value to 3.0. # If you want to generate a 720p video, it is recommended to set the shift value to 5.0. -shift = 3 +shift = 5 # Load pretrained model if need # Download from https://huggingface.co/TurboDiffusion/TurboWan2.1-T2V-1.3B-480P/ @@ -172,6 +172,12 @@ pipeline = WanPipeline( scheduler=scheduler, ) if ulysses_degree > 1 or ring_degree > 1: + # The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a + # sparse student under ring_degree>1 would silently mis-attend. + if ring_degree > 1: + raise ValueError( + "TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into " + "ulysses_degree instead (num_heads must divide ulysses_degree).") from functools import partial transformer.enable_multi_gpus_inference() if fsdp_dit: @@ -234,7 +240,6 @@ def save_results(): image = (image * 255).numpy().astype(np.uint8) image = Image.fromarray(image) image.save(video_path) - print(f"Saved image to: {video_path}") else: video_path = os.path.join(save_path, prefix + ".mp4") save_videos_grid(sample, video_path, fps=fps) diff --git a/scripts/taehv/README_TRAIN.md b/scripts/taehv/README_TRAIN.md index 22ac776..1993526 100644 --- a/scripts/taehv/README_TRAIN.md +++ b/scripts/taehv/README_TRAIN.md @@ -180,25 +180,32 @@ NCCL_DEBUG=INFO accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \ --config_path="config/wan2.2/wan_civitai_5b.yaml" \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --tae_path=$TAE_PATH \ - --train_data_dir=$DATASET_NAME \ - --train_data_meta=$DATASET_META_NAME \ - --video_sample_size=512 \ - --video_sample_stride=1 \ - --video_sample_n_frames=33 \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --tae_path="$TAE_PATH" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 \ + --video_sample_strides 1 2 3 \ + --video_sample_n_frames=25 \ + --validation_n_frames=49 \ --train_batch_size=1 \ --gradient_accumulation_steps=1 \ - --dataloader_num_workers=4 \ - --num_train_epochs=100 \ - --checkpointing_steps=500 \ + --dataloader_num_workers=8 \ + --max_train_steps=10000 \ + --checkpointing_steps=1000 \ + --validation_steps=1000 \ + --validation_paths "asset/inpaint_video.mp4" \ --learning_rate=1e-04 \ --lr_scheduler="constant_with_warmup" \ - --lr_warmup_steps=100 \ + --lr_warmup_steps=200 \ --latent_loss_weight=1.0 \ + --use_ema \ + --ema_decay=0.999 \ --seed=42 \ --output_dir="output_dir_taehv_w2.2" \ --mixed_precision="bf16" \ + --adam_beta1=0.9 \ + --adam_beta2=0.99 \ --adam_weight_decay=1e-2 \ --adam_epsilon=1e-08 \ --vae_mini_batch=1 \ @@ -208,6 +215,8 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \ --resume_from_checkpoint=latest ``` +These settings are also provided in `scripts/taehv/train_taehv.sh`. Edit its training arguments and run `bash scripts/taehv/train_taehv.sh`. + **Wan2.1 / Wan2.2 14B (taew2_1) Example**: Same as above, with these changes: @@ -246,6 +255,7 @@ export TAE_PATH="taew2_1.safetensors" | `--gradient_accumulation_steps` | Gradient accumulation steps | 1 | | `--dataloader_num_workers` | Number of DataLoader workers | 4 | | `--num_train_epochs` | Number of training epochs | 100 | +| `--max_train_steps` | Total training steps; overrides `--num_train_epochs` when set | 10000 | | `--checkpointing_steps` | Save a checkpoint every N steps | 500 | | `--checkpoints_total_limit` | Max number of checkpoints to store | `None` | | `--learning_rate` | Initial learning rate | 1e-4 | @@ -265,7 +275,7 @@ export TAE_PATH="taew2_1.safetensors" | `--trainable_modules_low_learning_rate` | Trainable modules with lr/2 | `[]` | | `--resume_from_checkpoint` | Resume training from checkpoint, use `"latest"` to auto-select | `latest` | | `--validation_steps` / `--validation_epochs` | Run validation every N steps / epochs | 2000 / 5 | -| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/1.mp4"` | +| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/inpaint_video.mp4"` | | `--validation_n_frames` | Validation clip length (4k+1); defaults to `--video_sample_n_frames` | `None` | **Sample Size Configuration Guide**: @@ -282,11 +292,11 @@ You can configure validation parameters to periodically reconstruct test videos |-----------|-------------|-------------------| | `--validation_steps` | Run validation every N steps | 2000 | | `--validation_epochs` | Run validation every N epochs | 5 | -| `--validation_paths` | Validation video paths | `"asset/1.mp4"` | +| `--validation_paths` | Validation video paths | `"asset/inpaint_video.mp4"` | | `--validation_n_frames` | Validation clip length, **4k+1**; omit to reuse `--video_sample_n_frames` | 81 | ```bash - --validation_paths "asset/1.mp4" \ + --validation_paths "asset/inpaint_video.mp4" \ --validation_steps=2000 \ --validation_epochs=5 \ --validation_n_frames=81 @@ -304,7 +314,7 @@ You can configure validation parameters to periodically reconstruct test videos - **Decoder-only distillation**: Add `--freeze_tae_encoder` to only improve decoding quality. - **Memory**: The teacher full VAE is the main memory consumer. Use `--low_vram` to keep it on CPU between encoding steps, or reduce `--video_sample_n_frames` / `--video_sample_size`. - **Sequential TAE**: `--use_taehv_sequential` trades speed for O(1) activation memory w.r.t. video length. -- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as `train_taehv_super.sh` does) when short runs need the average to catch up faster than the 0.9999 default. +- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as both `train_taehv.sh` and `train_taehv_super.sh` do) when short runs need the average to catch up faster than the 0.9999 default. ### 3.6 Multi-Node Distributed Training diff --git a/scripts/taehv/README_TRAIN_zh-CN.md b/scripts/taehv/README_TRAIN_zh-CN.md index 15cdbbd..068f74f 100644 --- a/scripts/taehv/README_TRAIN_zh-CN.md +++ b/scripts/taehv/README_TRAIN_zh-CN.md @@ -180,25 +180,32 @@ NCCL_DEBUG=INFO accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \ --config_path="config/wan2.2/wan_civitai_5b.yaml" \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --tae_path=$TAE_PATH \ - --train_data_dir=$DATASET_NAME \ - --train_data_meta=$DATASET_META_NAME \ - --video_sample_size=512 \ - --video_sample_stride=1 \ - --video_sample_n_frames=33 \ + --pretrained_model_name_or_path="$MODEL_NAME" \ + --tae_path="$TAE_PATH" \ + --train_data_dir="$DATASET_NAME" \ + --train_data_meta="$DATASET_META_NAME" \ + --fix_sample_size 720 1280 \ + --video_sample_strides 1 2 3 \ + --video_sample_n_frames=25 \ + --validation_n_frames=49 \ --train_batch_size=1 \ --gradient_accumulation_steps=1 \ - --dataloader_num_workers=4 \ - --num_train_epochs=100 \ - --checkpointing_steps=500 \ + --dataloader_num_workers=8 \ + --max_train_steps=10000 \ + --checkpointing_steps=1000 \ + --validation_steps=1000 \ + --validation_paths "asset/inpaint_video.mp4" \ --learning_rate=1e-04 \ --lr_scheduler="constant_with_warmup" \ - --lr_warmup_steps=100 \ + --lr_warmup_steps=200 \ --latent_loss_weight=1.0 \ + --use_ema \ + --ema_decay=0.999 \ --seed=42 \ --output_dir="output_dir_taehv_w2.2" \ --mixed_precision="bf16" \ + --adam_beta1=0.9 \ + --adam_beta2=0.99 \ --adam_weight_decay=1e-2 \ --adam_epsilon=1e-08 \ --vae_mini_batch=1 \ @@ -208,6 +215,8 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \ --resume_from_checkpoint=latest ``` +相同参数已写入 `scripts/taehv/train_taehv.sh`，也可修改其中的训练参数后，运行 `bash scripts/taehv/train_taehv.sh`。 + **Wan2.1 / Wan2.2 14B（taew2_1）训练示例**： 与上面相同，仅做如下修改： @@ -246,6 +255,7 @@ export TAE_PATH="taew2_1.safetensors" | `--gradient_accumulation_steps` | 梯度累积步数 | 1 | | `--dataloader_num_workers` | DataLoader 子进程数 | 4 | | `--num_train_epochs` | 训练 epoch 数 | 100 | +| `--max_train_steps` | 总训练步数；设置后覆盖 `--num_train_epochs` | 10000 | | `--checkpointing_steps` | 每 N 步保存 checkpoint | 500 | | `--checkpoints_total_limit` | 最多保留的 checkpoint 数量 | `None` | | `--learning_rate` | 初始学习率 | 1e-4 | @@ -265,7 +275,7 @@ export TAE_PATH="taew2_1.safetensors" | `--trainable_modules_low_learning_rate` | 以 lr/2 训练的模块 | `[]` | | `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | `latest` | | `--validation_steps` / `--validation_epochs` | 每 N 步 / 每 N 个 epoch 执行一次验证 | 2000 / 5 | -| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/1.mp4"` | +| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/inpaint_video.mp4"` | | `--validation_n_frames` | 验证片段帧数（4k+1），不传则复用 `--video_sample_n_frames` | `None` | **Sample Size 配置指南**： @@ -282,11 +292,11 @@ export TAE_PATH="taew2_1.safetensors" |------|------|--------| | `--validation_steps` | 每 N 步执行一次验证 | 2000 | | `--validation_epochs` | 每 N 个 epoch 执行一次验证 | 5 | -| `--validation_paths` | 验证视频路径 | `"asset/1.mp4"` | +| `--validation_paths` | 验证视频路径 | `"asset/inpaint_video.mp4"` | | `--validation_n_frames` | 验证片段帧数，**必须为 4k+1**；不传则复用 `--video_sample_n_frames` | 81 | ```bash - --validation_paths "asset/1.mp4" \ + --validation_paths "asset/inpaint_video.mp4" \ --validation_steps=2000 \ --validation_epochs=5 \ --validation_n_frames=81 @@ -304,7 +314,7 @@ export TAE_PATH="taew2_1.safetensors" - **Decoder-only 蒸馏**：加上 `--freeze_tae_encoder` 只提升解码质量。 - **显存**：teacher 完整 VAE 是主要的显存消耗者。使用 `--low_vram` 可以将其放在 CPU、仅在编码时搬上 GPU；也可以降低 `--video_sample_n_frames` / `--video_sample_size`。 - **串行 TAE**：`--use_taehv_sequential` 以速度换显存，激活显存与视频长度无关（O(1)）。 -- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。 +- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv.sh` 与 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。 ### 3.6 多机分布式训练 diff --git a/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md b/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md index 035fc4c..e244d0e 100644 --- a/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md +++ b/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md @@ -34,6 +34,7 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist - [3.5 Training with FSDP](#35-training-with-fsdp) - [3.6 Other Backends](#36-other-backends) - [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training) + - [3.8 Optional One-Forcing GAN, EMA, and Precision Controls](#38-optional-one-forcing-gan-ema-and-precision-controls) - [4. Inference Testing](#4-inference-testing) - [4.1 Inference Parameters](#41-inference-parameters) - [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference) @@ -107,6 +108,8 @@ DMD uses a **TextDataset** (`--train_mode="normal"`) — the generator creates i `train_distill.sh` defaults to `datasets/internal_datasets/metadata.json`. For a real run, swap in the **VidProM extended** prompts the paper uses. +Enabling any GAN/R1/R2 weight (see **3.8**) switches the loader to **ImageVideoDataset**: captioned real videos are then required, cropped to the critic's fake input frame count. + --- ## 3. Distillation Training @@ -174,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -204,6 +208,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -220,6 +226,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`. |-----------|-------------|---------------| | `--config_path` | Component subpaths / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | Base model directory (1.3B) | `$MODEL_NAME` | +| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher directory; defaults to `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | | `--train_data_meta` | Prompt JSON (only the `"text"` field is read) | `$DATASET_META_NAME` | | `--ode_transformer_path` | Stage 1 ODE checkpoint for generator/critic init | `$STAGE1_CKPT` | | `--fix_sample_size` | Fixed `H W` of the training clip | `480 832` | @@ -230,9 +237,9 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`. | `--dataloader_num_workers` | DataLoader workers | 8 | | `--num_train_epochs` | Number of training epochs (overridden by `--max_train_steps`) | 100 | | `--max_train_steps` | Hard step cap; the paper uses 600 | 600 | -| `--checkpointing_steps` | Save checkpoint every N steps | 50 | -| `--learning_rate` | Generator learning rate | `2e-06` | -| `--learning_rate_critic` | Critic learning rate | `2e-06` | +| `--checkpointing_steps` | Save checkpoint every N steps | 200 | +| `--learning_rate` | Generator learning rate | `2.0e-06` | +| `--learning_rate_critic` | Critic learning rate | `4.0e-07` | | `--lr_scheduler` | LR scheduler type | `constant_with_warmup` | | `--lr_warmup_steps` | LR warmup steps | 100 | | `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.5 | @@ -322,15 +329,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -352,6 +360,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -400,15 +410,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -430,10 +441,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` +`--use_ema` is dropped here: `train_distill.py` raises `NotImplementedError` for EMA under DeepSpeed Zero-3. + #### 3.6.2 Training without DeepSpeed and FSDP **This approach is not recommended because there is no memory-saving backend, which can easily cause out-of-memory errors**. We only provide the training shell for reference. @@ -464,15 +478,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -494,6 +509,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -538,15 +555,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -568,6 +586,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -602,6 +622,44 @@ NCCL_DEBUG=INFO - **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage) +### 3.8 Optional One-Forcing GAN, EMA, and Precision Controls + +These options port the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms already used by the causal-forcing DMD trainer. All GAN/R1/R2 weights default to **0** and stay off in the launcher. `--use_ema`, `--fsdp_reduce_fp32` and `--score_timestep_clamp` default to **off** in argparse, but `train_distill.sh` enables all three; the DeepSpeed Zero-3 command in **3.6.1** drops `--use_ema`, which Zero-3 rejects. + +The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`. + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `--use_ema` | EMA over the generator; checkpoints save a `transformer_ema` subfolder and validation samples from the EMA weights | off (enabled in `train_distill.sh`) | +| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` | +| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` | +| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` | +| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` | +| `--gan_block_num_heads` | Head count; must divide the fake-score hidden dimension | `12` | +| `--relativistic_discriminator` | Use paired real/fake logit differences | off | +| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` | +| `--r1_sigma` | R1 perturbation standard deviation | `0.01` | +| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` | +| `--r2_sigma` | R2 perturbation standard deviation | `0.01` | +| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in `train_distill.sh`) | +| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in `train_distill.sh`) | + +For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`. + +R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise. The non-relativistic critic runs fake and real as two sequential 1x-batch forwards with a backward after each (the memory-optimized split), the relativistic critic runs one packed 2x-batch forward, and R1/R2 each re-forward their own 2x batch after the main backward(s). + +To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_distill.py`. Requirements: +- Captioned real videos, not a prompt-only JSON. The real latents are cropped to the critic's fake input frame count, so the clip must cover at least that many latent frames. +- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch. +- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings. +- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement. + +Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule. + +Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags. + +Under `--use_ema`, the shadow stays in FP32 and mirrors the live model's mixed-precision policy; FSDP EMA requires FSDP1 and DeepSpeed Zero-3 is rejected. + --- ## 4. Inference Testing @@ -752,6 +810,7 @@ Any-order editing needs the full KV cache: both entries set `local_attn_size = - ## 5. Additional Resources - **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509 +- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing - **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md) - **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `sample_flex_level0` (the level-0 four-arm mixture, drawn by **both** trainers), `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes` - **Model / pipeline**: `videox_fun/models/wan_transformer3d_flex_forcing.py`, `videox_fun/pipeline/pipeline_wan_flex_forcing.py` diff --git a/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md b/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md index 6efb4ce..d6f3f66 100644 --- a/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md +++ b/scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md @@ -34,6 +34,7 @@ - [3.5 使用 FSDP 训练](#35-使用-fsdp-训练) - [3.6 其他后端](#36-其他后端) - [3.7 多机分布式训练](#37-多机分布式训练) + - [3.8 可选的 One-Forcing GAN、EMA 与精度开关](#38-可选的-one-forcing-ganema-与精度开关) - [四、推理测试](#四推理测试) - [4.1 推理参数解析](#41-推理参数解析) - [4.2 文生视频（T2V）推理](#42-文生视频t2v推理) @@ -107,6 +108,8 @@ DMD 使用 **TextDataset**（`--train_mode="normal"`）——生成器通过 rol `train_distill.sh` 默认指向 `datasets/internal_datasets/metadata.json`。真实训练请换成论文使用的 **VidProM extended** prompt 集。 +开启任何 GAN/R1/R2 权重（见 **3.8**）会把数据加载切到 **ImageVideoDataset**：此时需要带标注的真实视频，并会裁剪到 critic fake 输入的帧数。 + --- ## 三、蒸馏训练 @@ -174,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -204,6 +208,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -220,6 +226,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒 |-----------|-------------|---------------| | `--config_path` | 组件子路径 / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` | | `--pretrained_model_name_or_path` | 底座模型目录（1.3B） | `$MODEL_NAME` | +| `--real_score_pretrained_model_name_or_path` | 14B 真实分数教师目录；默认回落 `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | | `--train_data_meta` | prompt JSON（只读 `"text"` 字段） | `$DATASET_META_NAME` | | `--ode_transformer_path` | 用于生成器 / critic 初始化的阶段 1 ODE checkpoint | `$STAGE1_CKPT` | | `--fix_sample_size` | 固定训练片段的 `H W` | `480 832` | @@ -230,9 +237,9 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒 | `--dataloader_num_workers` | DataLoader 工作进程数 | 8 | | `--num_train_epochs` | 训练 epoch 数（被 `--max_train_steps` 覆盖） | 100 | | `--max_train_steps` | 硬性步数上限；论文用 600 | 600 | -| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 50 | -| `--learning_rate` | 生成器学习率 | `2e-06` | -| `--learning_rate_critic` | critic 学习率 | `2e-06` | +| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 200 | +| `--learning_rate` | 生成器学习率 | `2.0e-06` | +| `--learning_rate_critic` | critic 学习率 | `4.0e-07` | | `--lr_scheduler` | 学习率调度器类型 | `constant_with_warmup` | | `--lr_warmup_steps` | 学习率预热步数 | 100 | | `--real_guidance_scale` | 真实分数（14B 教师）的 CFG scale | 4.5 | @@ -333,15 +340,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -363,6 +371,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -405,15 +415,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -435,10 +446,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` +此处去掉 `--use_ema`：DeepSpeed Zero-3 下 `train_distill.py` 会对 EMA 抛出 `NotImplementedError`。 + DeepSpeed-Zero-3 训练完成之后，得到 `checkpoint-xxx` 状态的保存权重，我们需要将权重文件合并为单个，可以使用提供的`scripts/zero_to_bf16.py`脚本。 ```shell @@ -473,15 +487,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -503,6 +518,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -548,15 +565,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -578,6 +596,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" ``` @@ -604,6 +624,44 @@ export RANK=1 > - Use `NCCL_DEBUG=INFO` to troubleshoot issues in multi-machine training. > - If there is no RDMA environment, set `export NCCL_IB_DISABLE=1` and `export NCCL_P2P_DISABLE=1` to avoid communication issues. +### 3.8 可选的 One-Forcing GAN、EMA 与精度开关 + +这些开关把 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 的机制（因果 forcing 的 DMD trainer 已在用）移植过来。所有 GAN/R1/R2 权重默认是 **0**，launcher 中也保持关闭；`--use_ema`、`--fsdp_reduce_fp32` 与 `--score_timestep_clamp` 在 argparse 中默认**关闭**，但 `train_distill.sh` 把三者都打开了；**3.6.1** 的 DeepSpeed Zero-3 命令则去掉了 Zero-3 不支持的 `--use_ema`。 + +判别器与 **fake-score 骨干共享**，每个被选中的特征层有一个可学习的 register token。判别器头由 critic 优化器优化，而不是生成器优化器。生成器损失为 `DMD + 加权 GAN`；critic 损失为 `去噪 + 加权 GAN + R1 + R2`。 + +| 参数 | 说明 | 默认 | +|-----------|-------------|---------| +| `--use_ema` | 对生成器做 EMA；checkpoint 额外保存 `transformer_ema` 子目录，验证样片采样自 EMA 权重 | 关（launcher 已启用） | +| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` | +| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` | +| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` | +| `--gan_block_ffn_dim` | 判别器头 cross-attention FFN 维度 | `2048` | +| `--gan_block_num_heads` | 判别器头 head 数；必须整除 fake-score 隐维度 | `12` | +| `--relativistic_discriminator` | 使用成对的 real/fake logit 差值 | 关 | +| `--r1_weight` | 对含噪真实 latent 的 APT 式有限差分惩罚 | `0.0` | +| `--r1_sigma` | R1 扰动标准差 | `0.01` | +| `--r2_weight` | 对含噪生成 latent 的归一化有限差分惩罚 | `0.0` | +| `--r2_sigma` | R2 扰动标准差 | `0.01` | +| `--fsdp_reduce_fp32` | FSDP 梯度归约用 FP32；参数收集保持低精度 | 关（launcher 已启用） | +| `--score_timestep_clamp` | 把 DMD、生成器 GAN 与 critic 的时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关（launcher 已启用） | + +对 logit `D_fake` 与 `D_real`，普通 GAN 损失是生成器 `softplus(-D_fake)`、判别器 `softplus(-D_real) + softplus(D_fake)`；相对论版本改用 `softplus(-(D_fake - D_real))` 与 `softplus(-(D_real - D_fake))`。 + +R1 沿用来源的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 是 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。两者都是有限差分惩罚，不是输入梯度二重反向的正则。real/fake 样本共享时间步与噪声；扰动使用独立噪声。非相对论 critic 把 fake 与 real 拆成两次 1x batch 前向、各自一次 backward（省显存的拆分路径），相对论 critic 跑一次打包的 2x batch 前向，R1/R2 则在主 backward 之后各自再前向自己的 2x batch。 + +启用 GAN 需要给 `train_distill.py` 传入正的生成器/判别器权重，可选再加 R1/R2 权重。要求： +- 带标注的真实视频，不能是只有 prompt 的 JSON。真实 latent 会裁剪到 critic fake 输入的帧数，因此片段至少要覆盖那么多 latent 帧。 +- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。 +- 特征层索引必须对所选 fake-score 模型有效。默认的 `21 29` 层与 `12` head 适配 Wan2.1-T2V-1.3B；换骨干可能需要不同配置。 +- 只开 `gan_g_weight` 不会给判别器提供分类学习目标。请有意识地选择判别器权重；非零值不保证质量提升。 + +可训练的 fake-score 与判别器头 master 权重保持 **FP32**，前向使用配置的 autocast/FSDP 精度。logit、softplus 损失与有限差分归约都在 FP32 下计算。`--fsdp_reduce_fp32` 只改梯度通信精度，不改前向 dtype 或模型结构。时间步 clamp 只改分数训练的时间步，不改生成器 rollout 或验证调度。 + +恢复 GAN checkpoint 时请使用相同的 GAN/判别器头配置。判别器头随 fake-score 一起保存（含 FP32 的 head 张量），加载前重建。它不属于导出的生成器，因此推理不需要判别器或新的 GAN 旗标。 + +开启 `--use_ema` 时，shadow 保持 FP32 并与 live 模型的混合精度策略对齐；FSDP EMA 仅支持 FSDP1，DeepSpeed Zero-3 会被拒绝。 + --- ## 四、推理测试 @@ -741,4 +799,5 @@ num_frame_per_block = 7 # 干净上下文的提交粒度；None = 整 ## 五、更多资源 - **Flex-Forcing 论文**：https://arxiv.org/abs/2607.03509 +- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing - **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_flex_forcing/train_distill.py b/scripts/wan2.1_flex_forcing/train_distill.py index 5ca5b13..4b2e2ad 100644 --- a/scripts/wan2.1_flex_forcing/train_distill.py +++ b/scripts/wan2.1_flex_forcing/train_distill.py @@ -17,6 +17,7 @@ import argparse import contextlib +import dataclasses import gc import json import logging @@ -76,7 +77,8 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, WanTransformer3DModel, WanTransformer3DModel_FlexForcing, - WanTransformer3DModel_SelfForcing) + WanTransformer3DModel_SelfForcing, + add_wan_gan_discriminator) from videox_fun.pipeline import (WanFlexForcingPipeline, WanI2VPipeline, WanPipeline, WanSelfForcingPipeline) from videox_fun.utils.discrete_sampler import DiscreteSampling @@ -84,6 +86,7 @@ from videox_fun.utils.flex_chunking import (build_pyramid_partitions, chunk_boundaries, sample_flex_level0, validate_nested_partitions) +from videox_fun.utils.fsdp_ema import FSDPEMA from videox_fun.utils.tqdm_bar import PauseAwareTqdm from videox_fun.utils.utils import (calculate_dimensions, get_image_latent, get_image_to_video_latent, @@ -171,6 +174,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi return pred_for_score, score_num_frames, True +def repeat_batch(value, repeats): + """Repeat a batched tensor or per-sample list along the batch dimension.""" + if value is None or repeats == 1: + return value + if torch.is_tensor(value): + return torch.cat([value] * repeats) + return list(value) * repeats + + def filter_kwargs(cls, kwargs): import inspect sig = inspect.signature(cls.__init__) @@ -743,6 +755,15 @@ def parse_args(): " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" ), ) + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.") + parser.add_argument( + "--fsdp_reduce_fp32", action="store_true", + help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.", + ) + parser.add_argument( + "--score_timestep_clamp", action="store_true", + help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.", + ) parser.add_argument( "--non_ema_revision", type=str, @@ -965,6 +986,16 @@ def parse_args(): default=None, help=("If you want to load the ode-trained weight into generator transformer3d, input its path."), ) + parser.add_argument( + "--real_score_pretrained_model_name_or_path", + type=str, + default=None, + help=( + "Path to a non-causal pretrained model used as the DMD real_score teacher. " + "For CF Stage 3 this should point at Wan2.1-T2V-14B. If unset, falls back to " + "`--pretrained_model_name_or_path` (1.3B) — useful only for sanity testing." + ), + ) parser.add_argument( "--vae_path", type=str, @@ -1170,6 +1201,47 @@ def parse_args(): help="Probability of applying teacher forcing per step (1.0 = always)" ) + parser.add_argument( + "--gan_g_weight", type=float, default=0.0, + help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.", + ) + parser.add_argument( + "--gan_d_weight", type=float, default=0.0, + help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).", + ) + parser.add_argument( + "--gan_feature_layers", nargs="+", type=int, default=[21, 29], + help="fake_score block indices feeding the discriminator head, one register token per layer.", + ) + parser.add_argument( + "--gan_block_ffn_dim", type=int, default=2048, + help="FFN dimension of the discriminator cross-attention blocks.", + ) + parser.add_argument( + "--gan_block_num_heads", type=int, default=12, + help="Number of heads of the discriminator cross-attention blocks.", + ) + parser.add_argument( + "--relativistic_discriminator", action="store_true", + help="Use relativistic GAN losses based on real/fake logit differences.", + ) + parser.add_argument( + "--r1_weight", type=float, default=0.0, + help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).", + ) + parser.add_argument( + "--r1_sigma", type=float, default=0.01, + help="Perturbation standard deviation for the R1 penalty.", + ) + parser.add_argument( + "--r2_weight", type=float, default=0.0, + help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)." + ) + parser.add_argument( + "--r2_sigma", type=float, default=0.01, + help="Perturbation standard deviation for the R2 penalty.", + ) + args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) if env_local_rank != -1 and env_local_rank != args.local_rank: @@ -1184,6 +1256,14 @@ def parse_args(): def main(): args = parse_args() + use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)) + if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)): + raise ValueError("GAN and regularization weights must be finite and non-negative.") + if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1): + raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.") + if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0) + or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)): + raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.") if args.report_to == "wandb" and args.hub_token is not None: raise ValueError( @@ -1247,6 +1327,9 @@ def main(): else: fsdp_stage = 0 print(f"Using FSDP stage: {fsdp_stage}") + if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None: + fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32) + print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}") args.use_fsdp = True if fsdp_stage == 3: @@ -1374,8 +1457,11 @@ def main(): transformer_additional_kwargs=generator_transformer_kwargs, low_cpu_mem_usage=True, ).to(weight_dtype) + # CF Stage 3 uses a 14B non-causal teacher as real_score; pass an explicit + # `--real_score_pretrained_model_name_or_path` to point at Wan2.1-T2V-14B. + real_score_base = args.real_score_pretrained_model_name_or_path or args.pretrained_model_name_or_path real_score_transformer3d = WanTransformer3DModel.from_pretrained( - os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + os.path.join(real_score_base, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), low_cpu_mem_usage=True, ).to(weight_dtype) @@ -1476,6 +1562,41 @@ def main(): param.requires_grad = True break + if use_gan: + if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers) + or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0 + or fake_score_transformer3d.dim % args.gan_block_num_heads != 0): + raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.") + # Keep trainable discriminator master weights in fp32; forward uses autocast. + add_wan_gan_discriminator( + fake_score_transformer3d, + args.gan_feature_layers, + args.gan_block_ffn_dim, + args.gan_block_num_heads, + torch.float32, + seed=args.seed or 0, + ) + + # Create EMA for the generator_transformer3d. + if args.use_ema: + if zero_stage == 3: + raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.") + + if fsdp_plugin is not None: + if getattr(fsdp_plugin, "fsdp_version", 1) != 1: + raise NotImplementedError("FSDPEMA 仅支持 FSDP1。") + ema_module = generator_transformer_cls.from_config(generator_transformer3d.config) + ema_generator_transformer3d = FSDPEMA( + ema_module, source=generator_transformer3d, + accelerator=accelerator, fsdp_plugin=fsdp_plugin, + ) + else: + ema_generator_transformer3d = EMAModel( + generator_transformer3d.parameters(), + model_cls=generator_transformer_cls, + model_config=generator_transformer3d.config, + ) + # `accelerate` 0.16.0 will have better support for customized saving if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format @@ -1486,7 +1607,10 @@ def main(): from safetensors.torch import save_file safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors") - accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()} + accelerate_state_dict = { + k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype) + for k, v in accelerate_state_dict.items() + } save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"}) with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: @@ -1503,7 +1627,14 @@ def main(): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): if accelerator.is_main_process: - models[0].save_pretrained(os.path.join(output_dir, "transformer")) + model = accelerator.unwrap_model(models[0]) + save_dir = os.path.join(output_dir, "transformer") + if hasattr(model, "gan_head"): + from safetensors.torch import save_file + model.save_config(save_dir) + save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"}) + else: + model.save_pretrained(save_dir) if not args.use_deepspeed: weights.pop() @@ -1529,14 +1660,19 @@ def main(): else WanTransformer3DModel ) - # load diffusers style into model - load_model = model_cls.from_pretrained( - input_dir, subfolder="transformer" - ) - unwrapped.register_to_config(**load_model.config) + # GAN heads are recreated from the same CLI configuration before resume. + if hasattr(unwrapped, "gan_head"): + from safetensors.torch import load_file + unwrapped.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True) + else: + # load diffusers style into model + load_model = model_cls.from_pretrained( + input_dir, subfolder="transformer" + ) + unwrapped.register_to_config(**load_model.config) - unwrapped.load_state_dict(load_model.state_dict()) - del load_model + unwrapped.load_state_dict(load_model.state_dict()) + del load_model pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") if os.path.exists(pkl_path): @@ -1550,6 +1686,32 @@ def main(): accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook) accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook) + if args.use_ema: + def save_ema_hook(models, weights, output_dir): + if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process: + ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema")) + + def load_ema_hook(models, input_dir): + ema_path = os.path.join(input_dir, "transformer_ema") + if isinstance(ema_generator_transformer3d, FSDPEMA): + ema_generator_transformer3d.load_pretrained(ema_path) + else: + _, ema_kwargs = generator_transformer_cls.load_config(ema_path, return_unused_kwargs=True) + load_model = generator_transformer_cls.from_pretrained( + input_dir, subfolder="transformer_ema", + transformer_additional_kwargs=generator_transformer_kwargs, + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + load_model = EMAModel(load_model.parameters(), model_cls=generator_transformer_cls, model_config=load_model.config) + load_model.load_state_dict(ema_kwargs) + ema_generator_transformer3d.load_state_dict(load_model.state_dict()) + ema_generator_transformer3d.to(accelerator.device) + del load_model + + accelerator.register_save_state_pre_hook(save_ema_hook) + accelerator.register_load_state_pre_hook(load_ema_hook) + if args.gradient_checkpointing: generator_transformer3d.enable_gradient_checkpointing() fake_score_transformer3d.enable_gradient_checkpointing() @@ -1646,6 +1808,12 @@ def main(): print(f"Set {name} to lr : {args.learning_rate_critic / 2}") break + if use_gan: + fake_trainable_params_optim[0]['params'].extend( + param for name, param in fake_score_transformer3d.named_parameters() + if name.startswith("gan_head.") and name not in in_already + ) + if args.use_came: optimizer = optimizer_cls( trainable_params_optim, @@ -1687,7 +1855,7 @@ def main(): args.random_hw_adapt = False # Get the dataset - if args.train_mode != "normal" or args.use_teacher_forcing: + if args.train_mode != "normal" or args.use_teacher_forcing or use_gan: train_dataset = ImageVideoDataset( args.train_data_meta, args.train_data_dir, video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, @@ -1715,7 +1883,7 @@ def main(): return length_to_frame_num - if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing: + if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan: aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} batch_sampler_generator = torch.Generator().manual_seed(args.seed) batch_sampler = AspectRatioBatchImageVideoSampler( @@ -2058,6 +2226,9 @@ def main(): if args.train_mode != "normal": clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if args.use_ema: + ema_generator_transformer3d.to(accelerator.device) + # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) if overrode_max_train_steps: @@ -2148,6 +2319,8 @@ def main(): vae_stream_2 = None idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) + min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps) + max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps) def randomize_denoising_step_indices( denoising_step_indices_list, @@ -2208,6 +2381,8 @@ def main(): # must be averaged by contribution count, not by gradient_accumulation_steps. train_gen_log_count = 0 train_denoising_loss = 0.0 + train_gan_g_loss = 0.0 + train_gan_d_loss = 0.0 batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) for step, batch in enumerate(train_dataloader): # Data batch sanity check @@ -2340,7 +2515,7 @@ def main(): pixel_values_bs = pixel_values_bs.sample() new_pixel_values.append(pixel_values_bs) return torch.cat(new_pixel_values, dim = 0) - if args.use_teacher_forcing: + if args.use_teacher_forcing or use_gan: clean_latents = _batch_encode_vae(pixel_values) else: clean_latents = None @@ -2372,7 +2547,7 @@ def main(): clip_context.append(_clip_context) clip_context = torch.cat(clip_context) - if args.use_teacher_forcing: + if args.use_teacher_forcing or use_gan: target_shape = clean_latents.size() else: target_shape = mask_latents.size() @@ -3120,6 +3295,8 @@ def main(): # Common code for both KV cache and block mask training indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + if args.score_timestep_clamp: + generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep) generator_denoised_input = add_noise( generator_pred, torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), @@ -3225,17 +3402,53 @@ def main(): train_dmd_loss += avg_dmd_loss.item() train_gen_log_count += 1 + # --- One-Forcing Generator GAN Loss --- + generator_loss = dmd_loss + if args.gan_g_weight > 0: + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + if args.score_timestep_clamp: + gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep) + gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng) + gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)] + if args.relativistic_discriminator: + real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype) + gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep)) + num_gan_inputs = len(gan_inputs) + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, gan_logits = fake_score_transformer3d( + x=torch.cat(gan_inputs).to(dtype=weight_dtype), + context=repeat_batch(prompt_embeds, num_gan_inputs), + t=gan_timestep.repeat(num_gan_inputs), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None, + gan_classify=True, + ) + gan_logits = gan_logits.float().chunk(num_gan_inputs) + if args.relativistic_discriminator: + gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean() + else: + gan_g_loss = F.softplus(-gan_logits[0]).mean() + gan_g_loss = gan_g_loss * args.gan_g_weight + train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item() + generator_loss = generator_loss + gan_g_loss + if args.low_vram: real_score_transformer3d = real_score_transformer3d.to("cpu") - fake_score_transformer3d = fake_score_transformer3d.to("cpu") + if args.gan_g_weight == 0: + fake_score_transformer3d = fake_score_transformer3d.to("cpu") torch.cuda.empty_cache() - accelerator.backward(dmd_loss) + accelerator.backward(generator_loss) if accelerator.sync_gradients: accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) optimizer.step() lr_scheduler.step() optimizer.zero_grad() + if args.gan_g_weight > 0: + # Discard discriminator gradients from the generator GAN backward. + critic_optimizer.zero_grad() if args.low_vram: fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) @@ -3626,6 +3839,8 @@ def main(): indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + if args.score_timestep_clamp: + critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep) critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng) fake_score_denoised_input = add_noise( @@ -3634,15 +3849,25 @@ def main(): critic_timestep ) - with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): - fake_score_denoised_output = fake_score_transformer3d( - x=fake_score_denoised_input, - context=prompt_embeds, - t=critic_timestep, - seq_len=seq_len, - y=inpaint_latents if args.train_mode != "normal" else None, - clip_fea=clip_context if args.train_mode != "normal" else None, - ) + # Generated and real latents share critic_noise / critic_timestep. Prepare the + # real input and the R1/R2 perturbations once; they are consumed by the critic + # forward below and by the R1/R2 penalties after the main backward(s). + real_noisy = None + r1_perturbation = None + r2_perturbation = None + if use_gan: + real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype) + real_noisy = add_noise(real_latents, critic_noise, critic_timestep) + del real_latents + if args.r1_weight > 0: + r1_perturbation = args.r1_sigma * torch.randn( + real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng + ) + if args.r2_weight > 0: + r2_perturbation = args.r2_sigma * torch.randn( + fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype, + device=fake_score_denoised_input.device, generator=torch_rng + ) def custom_mse_loss(noise_pred, target, weighting=None, threshold=50): noise_pred = noise_pred.float() @@ -3656,11 +3881,136 @@ def main(): final_loss = masked_loss.mean() return final_loss - denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) - avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() - train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps - - accelerator_fake_score_transformer3d.backward(denoising_loss) + # --- One-Forcing Discriminator GAN Loss --- + if use_gan and not args.relativistic_discriminator: + # Memory-optimized critic: the non-relativistic gan_d_loss = + # softplus(-real) + softplus(fake) is separable, so run fake and real as two + # sequential 1x-batch forwards and backward the fake pass first. Backwarding + # fake frees its gradient-checkpointing block inputs before the real forward, + # dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to + # 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s) + # accumulate into fake_score.grad, giving the identical update as the packed + # single-forward critic. + # -- Fake pass (1x batch): denoising output + fake discriminator logit -- + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d( + x=fake_score_denoised_input, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + gan_classify=True, + ) + denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) + avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() + train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps + + fake_logit = fake_gan_logits.float() + fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight + critic_loss = denoising_loss + fake_d_loss + accelerator_fake_score_transformer3d.backward(critic_loss) + # Fake graph freed: the transformer blocks' stored inputs are released here. + fake_d_detached = fake_d_loss.detach() + del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output + + # -- Real pass (1x batch): discriminator logit only -- + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, real_gan_logits = fake_score_transformer3d( + x=real_noisy, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + gan_classify=True, + ) + real_logit = real_gan_logits.float() + real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight + accelerator_fake_score_transformer3d.backward(real_d_loss) + # Real graph freed. Combine the two detached halves for logging. + gan_d_loss = fake_d_detached + real_d_loss.detach() + train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item() + del fake_d_detached, real_gan_logits, real_logit, real_d_loss + else: + # Combined path: 2x-batch forward for relativistic GAN (both logits coupled by + # softplus(-(real - fake))), or 1x-batch forward when GAN is disabled. + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + if use_gan: + fake_score_denoised_output, gan_logits = fake_score_transformer3d( + x=torch.cat([fake_score_denoised_input, real_noisy]), + context=repeat_batch(prompt_embeds, 2), + t=critic_timestep.repeat(2), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None, + gan_classify=True, + ) + fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0] + gan_logits = gan_logits.float().chunk(2) + else: + fake_score_denoised_output = fake_score_transformer3d( + x=fake_score_denoised_input, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) + avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() + train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps + + critic_loss = denoising_loss + if use_gan: + fake_logit, real_logit = gan_logits[0], gan_logits[1] + if args.relativistic_discriminator: + gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean() + else: + gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean() + gan_d_loss = gan_d_loss * args.gan_d_weight + train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item() + critic_loss = critic_loss + gan_d_loss + + accelerator_fake_score_transformer3d.backward(critic_loss) + del critic_loss + + # R1/R2 gradient penalties: each re-forwards its own 2x batch after the main + # backward(s) and finite-differences the perturbed logit against the clean one, + # matching the packed-formula value while keeping the peak at 2x batch. + if use_gan and args.r1_weight > 0: + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, r1_logits = fake_score_transformer3d( + x=torch.cat([real_noisy + r1_perturbation, real_noisy]), + context=repeat_batch(prompt_embeds, 2), + t=critic_timestep.repeat(2), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None, + gan_classify=True, + ) + r1_logits = r1_logits.float().chunk(2) + r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean() + accelerator_fake_score_transformer3d.backward(r1_loss) + del r1_loss + + if use_gan and args.r2_weight > 0: + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, r2_logits = fake_score_transformer3d( + x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]), + context=repeat_batch(prompt_embeds, 2), + t=critic_timestep.repeat(2), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None, + gan_classify=True, + ) + r2_logits = r2_logits.float().chunk(2) + r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean() + accelerator_fake_score_transformer3d.backward(r2_loss) + del r2_loss + if accelerator_fake_score_transformer3d.sync_gradients: accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm) critic_optimizer.step() @@ -3674,12 +4024,19 @@ def main(): # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: + if args.use_ema: + ema_generator_transformer3d.step(generator_transformer3d.parameters()) progress_bar.update(1) global_step += 1 - accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step) + logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)} + if use_gan: + logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss}) + accelerator.log(logs, step=global_step) train_dmd_loss = 0.0 train_gen_log_count = 0 train_denoising_loss = 0.0 + train_gan_g_loss = 0.0 + train_gan_d_loss = 0.0 if global_step % args.checkpointing_steps == 0: if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: @@ -3722,6 +4079,10 @@ def main(): if args.validation_prompts is not None and global_step % args.validation_steps == 0: with progress_bar.paused(): + if args.use_ema: + # Store the generator parameters temporarily and load the EMA parameters to perform inference. + ema_generator_transformer3d.store(generator_transformer3d.parameters()) + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) log_validation( vae, text_encoder, @@ -3734,6 +4095,9 @@ def main(): weight_dtype, global_step, ) + if args.use_ema: + # Switch back to the original generator parameters. + ema_generator_transformer3d.restore(generator_transformer3d.parameters()) logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} progress_bar.set_postfix(**logs) @@ -3743,6 +4107,10 @@ def main(): if args.validation_prompts is not None and epoch % args.validation_epochs == 0: with progress_bar.paused(): + if args.use_ema: + # Store the generator parameters temporarily and load the EMA parameters to perform inference. + ema_generator_transformer3d.store(generator_transformer3d.parameters()) + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) log_validation( vae, text_encoder, @@ -3755,6 +4123,9 @@ def main(): weight_dtype, global_step, ) + if args.use_ema: + # Switch back to the original generator parameters. + ema_generator_transformer3d.restore(generator_transformer3d.parameters()) # Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever # something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto diff --git a/scripts/wan2.1_flex_forcing/train_distill.sh b/scripts/wan2.1_flex_forcing/train_distill.sh index 5643e75..6c672fa 100644 --- a/scripts/wan2.1_flex_forcing/train_distill.sh +++ b/scripts/wan2.1_flex_forcing/train_distill.sh @@ -28,15 +28,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --gradient_accumulation_steps=1 \ --dataloader_num_workers=8 \ --num_train_epochs=100 \ - --checkpointing_steps=50 \ - --learning_rate=2e-06 \ - --learning_rate_critic=2e-06 \ + --checkpointing_steps=200 \ + --learning_rate=2.0e-06 \ + --learning_rate_critic=4.0e-07 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ --output_dir="output_dir_wan2.1_flex_forcing_distill" \ --gradient_checkpointing \ --mixed_precision="bf16" \ + --fsdp_reduce_fp32 \ --adam_weight_decay=3e-2 \ --adam_epsilon=1e-10 \ --vae_mini_batch=1 \ @@ -58,5 +59,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --randomize_step_indices \ --flow_euler_rollout \ --use_kv_cache_training \ + --use_ema \ + --score_timestep_clamp \ --num_frame_per_block=3 \ --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors" diff --git a/videox_fun/dist/__init__.py b/videox_fun/dist/__init__.py index 88b4d51..70b1c41 100755 --- a/videox_fun/dist/__init__.py +++ b/videox_fun/dist/__init__.py @@ -21,10 +21,13 @@ from .longcatvideo_xfuser import (usp_attn_longcatvideo_avatar_forward, usp_rope_longcatvideo_forward) from .ltx2_xfuser import (LTX2MultiGPUsAttnProcessor, LTX2PerturbedMultiGPUsAttnProcessor) -from .minimax_h3_xfuser import MiniMaxH3MultiGPUsAttnProcessor +from .minimax_h3_xfuser import (MiniMaxH3MultiGPUsAttnProcessor, + MiniMaxH3MultiGPUsVSAAttnProcessor) from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0 from .qwenimage21_xfuser import QwenImage21MultiGPUsAttnProcessor -from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward, +from .wan_xfuser import (usp_attn_forward, usp_attn_fastvideo_forward, + usp_attn_s2v_forward, usp_attn_self_forcing_forward, - usp_attn_sparkwan_forward) + usp_attn_sparkwan_forward, + usp_attn_turbowan_forward) from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor diff --git a/videox_fun/dist/fuser.py b/videox_fun/dist/fuser.py index d76bb7a..a148f2c 100755 --- a/videox_fun/dist/fuser.py +++ b/videox_fun/dist/fuser.py @@ -41,7 +41,7 @@ def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance if get_sp_group is None: raise RuntimeError("xfuser is not installed.") dist.init_process_group("nccl") - print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=% rank=%d world_size=%d' % ( + print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=%d rank=%d world_size=%d' % ( ulysses_degree, ring_degree, classifier_free_guidance_degree, dist.get_rank(), dist.get_world_size())) assert dist.get_world_size() == ring_degree * ulysses_degree * classifier_free_guidance_degree, \ @@ -53,6 +53,14 @@ def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance ulysses_degree=ulysses_degree) # device = torch.device("cuda:%d" % dist.get_rank()) device = torch.device(f"cuda:{get_world_group().local_rank}") + # Bind this process to its own GPU. ``init_process_group`` above runs before xfuser's + # ``init_distributed_environment``, whose ``set_device`` is guarded by + # ``if not torch.distributed.is_initialized()`` and so never fires; without this call the + # current device stays 0 on every rank while the model lives on ``cuda:{local_rank}``. + # PyTorch ops dispatch on ``tensor.device`` and are unaffected, but a Triton kernel launches + # on ``torch.cuda.current_device()`` and rejects the rank's own pointers as inaccessible, + # so any Triton sparse attention (e.g. VSA) would fail on rank > 0. + torch.cuda.set_device(device) print('rank=%d device=%s' % (get_world_group().rank, str(device))) else: device = "cuda" diff --git a/videox_fun/dist/minimax_h3_xfuser.py b/videox_fun/dist/minimax_h3_xfuser.py index 3e7b1e0..2ec7f6c 100644 --- a/videox_fun/dist/minimax_h3_xfuser.py +++ b/videox_fun/dist/minimax_h3_xfuser.py @@ -1,5 +1,7 @@ import torch +import torch.nn.functional as F +from .fuser import get_sp_group, ulysses_all_to_all from .fuser import xFuserLongContextAttention @@ -80,3 +82,99 @@ class MiniMaxH3MultiGPUsAttnProcessor: hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) return hidden_states + + +class MiniMaxH3MultiGPUsVSAAttnProcessor: + """Sequence-parallel attention processor for the FastVideo VSA student on MiniMax-H3. + + Mirrors ``MiniMaxH3MultiGPUsAttnProcessor``'s row-split -> head-split Ulysses scheme, but swaps + the dense ``xFuserLongContextAttention`` core for the VSA block-sparse op (``attn.attn_op``) so + the trained coarse-to-fine selection and the ``to_gate_compress`` compensation branch are kept + instead of being silently dropped by a dense fallback. + + Two VSA-specific points under head-parallelism: + * ``to_gate_compress`` is a ``Linear(dim, heads*head_dim)`` that MIXES heads, so it cannot run + in the head-shard layout. It is strictly per-token and SP shards rows while keeping the full + hidden dim, so the gate is evaluated here on this rank's ROW shard (yielding all heads), then + all-to-all'd alongside q/k/v -- element-wise identical to computing it on the full sequence + and slicing heads, and it lands on the same contiguous head block the all-to-all hands back. + * the parent forward pads the packed sequence up to a multiple of the group size, so the + gathered sequence carries ``<= sp_world_size - 1`` zero rows at its tail. VSA builds its tile + geometry from ``vsa_segments``, whose lengths sum to the PADLESS ``valid_length``; the padding + tail is therefore sliced off before the op and the output padded back afterwards. Unlike the + dense processor (which lets the pad rows dilute the softmax by ~1e-5), this keeps VSA + bitwise-equal to the single-GPU run -- the pad rows never enter the kernel. + + Requires a padless packed sequence (``attention_mask is None``) and a pure-Ulysses split + (``num_heads`` divisible by the group size, checked by the model's ``enable_multi_gpus_inference``). + A dense student (``topk >= 1.0``) carries no ``attn_op`` and must keep the dense processor. + """ + + def __call__( + self, + attn, + hidden_states: torch.Tensor, + rotary_emb=None, + attention_mask=None, + valid_length=None, + vsa_segments=None, + ) -> torch.Tensor: + from ..models.minimax_h3_transformer3d import \ + apply_minimax_h3_rotary_emb + + if attention_mask is not None: + raise ValueError( + "MiniMaxH3MultiGPUsVSAAttnProcessor does not support a masked (padded) packed sequence.") + if not hasattr(attn, "attn_op"): + raise RuntimeError( + "The VSA sequence-parallel processor needs attn.attn_op; a dense student (topk>=1.0) " + "must bind MiniMaxH3MultiGPUsAttnProcessor instead.") + + query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1)) + key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1)) + value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + # The rotary coordinates are this rank's rows, so apply them before the all-to-all moves rows. + if rotary_emb is not None: + query = apply_minimax_h3_rotary_emb(query, *rotary_emb) + key = apply_minimax_h3_rotary_emb(key, *rotary_emb) + + # Cross-head compress gate on this rank's ROW shard (full hidden dim present): [B, chunk, H, D]. + gate = attn.attn_op.to_gate_compress(hidden_states).unflatten(-1, (attn.heads, -1)) + + half_dtypes = (torch.float16, torch.bfloat16) + + def half(x): + return x if x.dtype in half_dtypes else x.to(torch.bfloat16) + + # row-split -> head-split over the whole (padded) packed sequence, gate alongside q/k/v. + query = ulysses_all_to_all(half(query), scatter_dim=2, gather_dim=1) + key = ulysses_all_to_all(half(key), scatter_dim=2, gather_dim=1) + value = ulysses_all_to_all(half(value), scatter_dim=2, gather_dim=1) + gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1) + + # Drop the SP padding tail so VSA sees exactly the padless length its segments describe. + padded_length = query.shape[1] + trim = valid_length is not None and valid_length < padded_length + if trim: + query = query[:, :valid_length] + key = key[:, :valid_length] + value = value[:, :valid_length] + gate = gate[:, :valid_length] + + out = attn.attn_op(query, key, value, None, segments=vsa_segments, gate=gate) + + # Pad the output back to the gathered length before the inverse all-to-all; the zero rows are + # dropped by the model's final all_gather and never touch a real row's attention. + if trim: + out = F.pad(out, (0, 0, 0, 0, 0, padded_length - valid_length)) + + # head-split -> row-split, back to this rank's rows. + out = ulysses_all_to_all(out.to(hidden_states.dtype), scatter_dim=1, gather_dim=2) + out = out.flatten(2, 3).type_as(query) + out = attn.to_out[0](out) + out = attn.to_out[1](out) + return out diff --git a/videox_fun/dist/wan_xfuser.py b/videox_fun/dist/wan_xfuser.py index 8815ecb..59dff42 100755 --- a/videox_fun/dist/wan_xfuser.py +++ b/videox_fun/dist/wan_xfuser.py @@ -224,7 +224,10 @@ def usp_attn_sparkwan_forward(self, q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1) k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1) v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1) - x = self.attn_op(q, k, v) + # heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to + # the single-GPU result despite this rank holding only the H/P head block. + h_per = n // sp_size + x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per) x = ulysses_all_to_all(x, scatter_dim=1, gather_dim=2) x = x.flatten(2) x = self.o(x) @@ -267,6 +270,140 @@ def usp_attn_sparkwan_forward(self, out = self.o(out) return out +def usp_attn_turbowan_forward(self, + x, + seq_lens, + grid_sizes, + freqs, + dtype=torch.bfloat16, + t=0): + """Sequence-parallel self-attention for the TurboDiffusion SLA student. + + ``TurboWanTransformer3DModel`` swaps ``WanSelfAttention`` for a local variant whose kernel + is ``self.attn_op`` (``FakeLocalAttention`` -> ``SparseLinearAttention``, block-sparse exact + softmax + a ``proj_l`` feature-map linear branch). The stock ``usp_attn_forward`` binds + ``xFuserLongContextAttention``'s dense full-softmax and only touches the base-class + ``q/k/v/norm_q/norm_k/o`` attributes, so patching it onto this model would silently drop the + SLA kernel -- every needed attribute is inherited, hence no AttributeError, no shape error, + just dense numbers. This variant keeps ``self.attn_op`` live under Ulysses head-parallelism. + + Equivalence: ``SparseLinearAttention.forward`` takes and returns ``[B, L, H, D]``, transposes + to BHLD internally and does every reduction over L/D only (``get_block_map`` pooling and + top-k, the sparse kernel, the ``k^T v`` linear branch), and ``proj_l`` is a per-head + ``head_dim -> head_dim`` projection -- no cross-head mixing and no head-indexed parameter, + so an all-to-all from the sequence shard ``[B, S/P, H, D]`` to the head shard + ``[B, S, H/P, D]`` reproduces the single-GPU result exactly (top-k is a per-query key-block + ratio). RoPE is applied BEFORE the all-to-all because ``rope_apply`` slices the frequency + table by sequence-parallel rank. + + Requires a pure-Ulysses split (``ring_degree`` must be 1: the SLA kernel exposes no partial + output / logsumexp interface, so ring's online rescale-concat cannot be wired to it) and + ``num_heads`` divisible by the SP degree. + """ + # No sparse kernel on this instance: keep the stock dense SP path. + if not hasattr(self, "attn_op"): + return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t) + + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + half_dtypes = (torch.float16, torch.bfloat16) + + def half(t_): + return t_ if t_.dtype in half_dtypes else t_.to(dtype) + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d) + k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d) + v = self.v(x.to(dtype)).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + q, k = rope_apply_qk(q, k, grid_sizes, freqs) + + # seq-shard -> head-shard, SLA over the FULL sequence on a subset of heads, then back. + q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1) + k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1) + v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1) + # heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to the + # single-GPU result despite this rank holding only the H/P head block. + sp_size = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + h_per = n // sp_size + x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per) + x = ulysses_all_to_all(x.to(dtype), scatter_dim=1, gather_dim=2) + + # output + x = x.flatten(2) + x = self.o(x) + return x + +def usp_attn_fastvideo_forward(self, + x, + seq_lens, + grid_sizes, + freqs, + dtype=torch.bfloat16, + t=0): + """Sequence-parallel self-attention for the FastVideo VSA student. + + Same Ulysses head-parallel scheme as ``usp_attn_sparkwan_forward``, with one VSA-specific + wrinkle: ``VSASparseAttention.to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)`` + that MIXES heads, so it cannot run inside the head-shard layout the way the rest of the VSA + kernel can (``sparse`` / ``compressed`` only reduce over L/D and are head-independent). + It is, however, strictly per-token -- and sequence parallelism shards the sequence while + leaving the hidden dim intact, so every rank still holds the full ``C`` for its own tokens. + The gate is therefore evaluated here on the SEQUENCE shard (yielding all ``H`` heads), then + all-to-all'd alongside q/k/v, which is element-wise identical to computing it on the full + sequence and slicing heads -- and lands on exactly the contiguous head block + ``[sp_rank*H/P, (sp_rank+1)*H/P)`` the all-to-all hands this rank. + + The gathered sequence is ``s * sp_size``, which the parent forward pads to a multiple of the + SP degree, so it can exceed ``grid_sizes.prod()``; ``_vsa_dispatch`` routes that case to its + per-sample branch (attend over the valid grid tokens, pad the output back), keeping the tile + geometry built from the true grid. + + Requires a pure-Ulysses split (``ring_degree`` must be 1) and ``num_heads`` divisible by the + SP degree. ``topk >= 1.0`` (the dense fallback) defers to ``usp_attn_forward``. + """ + # Dense student (topk >= 1.0) or a plain Wan self-attention: keep the stock SP path. + if getattr(self, "topk", 1.0) >= 1.0 or not hasattr(self, "attn_op"): + return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t) + + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + half_dtypes = (torch.float16, torch.bfloat16) + + def half(t_): + return t_ if t_.dtype in half_dtypes else t_.to(dtype) + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d) + k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d) + v = self.v(x.to(dtype)).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + q, k = rope_apply_qk(q, k, grid_sizes, freqs) + + # Cross-head compress gate, evaluated on this rank's sequence shard where the full hidden + # dim is available: [B, S/P, C] -> [B, S/P, H, D] (all heads). + gate = self.attn_op.to_gate_compress(x.to(dtype)).view(b, s, n, d) + + # seq-shard -> head-shard for q/k/v and the gate alike. + q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1) + k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1) + v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1) + gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1) + + # VSA over the FULL sequence on a head shard; x is unused once gate is supplied. + out = self._vsa_dispatch(q, k, v, None, seq_lens, grid_sizes, gate=gate) + + # head-shard -> seq-shard, back to this rank's sequence chunk. + out = ulysses_all_to_all(out.to(dtype), scatter_dim=1, gather_dim=2) + out = out.flatten(2) + out = self.o(out) + return out + @amp.autocast(enabled=False) @torch.compiler.disable() def s2v_rope_apply(x, grid_sizes, freqs): diff --git a/videox_fun/models/attention_kernel.py b/videox_fun/models/attention_kernel.py index de720a2..7e7d34e 100644 --- a/videox_fun/models/attention_kernel.py +++ b/videox_fun/models/attention_kernel.py @@ -51,17 +51,6 @@ import triton.language as tl # loads, a fused tl.dot(p, v, acc) accumulate, and a per-query-tile variable kv-block count (vs a # uniform-MAXK loop over a dummy-padded tile). # -# This is deliberately NOT the fastest forward available. That same sm90 CUDA kernel (ThunderKittens) -# IS buildable from source on this box -- only its prebuilt wheels are pinned -- and measures -# 4.66 ms/call on the workload above, i.e. 1.35x over this Triton kernel (step-level VSA-vs-dense -# 1.019x -> 1.115x at 480P, 1.233x -> 1.381x at 720P). It stays out of the production path because -# it adds a build dependency and fails SILENTLY three ways: on a non-sm90a device pass it compiles -# to an empty stub returning uninitialised memory; it reinterpret_casts `variable_block_sizes` to -# int32 with no dtype check, so an int64 tensor yields all-NaN; and its lse is [B,H,L,1] in the -# natural-log domain scaled by -1/qk_scale rather than our [B,H,L] base-2. The build recipe, the -# runtime gate it would need, and the measurement set are archived in -# `_fastvideo_sm90_kernel/README.md`. -# # Deviations from upstream, listed so a diff against the sdist still audits this block: # 1. its `import torch / triton / triton.language`, and its SPARSE-ADDITION `import math`, are # hoisted to this module's top import block -- same names, no behavioural change, and the @@ -495,12 +484,10 @@ class _sparse_linear_attention_varlen(torch.autograd.Function): @staticmethod def forward(ctx, q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes, kb_id, qk_scale): assert q_tile.is_contiguous() and k_tile.is_contiguous() and v_tile.is_contiguous() - B, H, L_BUF, D = q_tile.shape - # Forward runs the vendored official Triton block-sparse kernel, which consumes the - # official (q2k_index, q2k_num) pair map_to_index produced: the first q2k_num[row] entries - # of each row are that query tile's kv-tile ids and the -1 pads are never read. Its M is - # [B,H,L_BUF] base-2 lse (m_i + log2 l_i), byte-compatible with what our backward expects, - # so the save_for_backward contract is unchanged. + # The forward consumes the official (q2k_index, q2k_num) pair map_to_index produced: + # the first q2k_num[row] entries of each row are that query tile's kv-tile ids and the + # -1 pads are never read. It produces M as [B,H,L_BUF] base-2 lse (m_i + log2 l_i), + # which is exactly what the varlen backward kernels below consume. out, lse = triton_block_sparse_attn_forward( q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes) ctx.save_for_backward(q_tile, k_tile, v_tile, q2k_idx, q2k_num, @@ -625,7 +612,13 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) K = pooled_score.shape[-1] - topk = min(K, int(topk_ratio * K)) + # Clamp to at least one key block. A truncating int() selects ZERO blocks whenever + # topk_ratio * K < 1 (short sequence and/or aggressive ratio -- e.g. RoLa's default 0.03 + # below 2176 tokens, SLA's 0.1 below 640), and the sparse kernel then attends to nothing and + # emits all-NaN. Keeps the truncating count everywhere else, so no existing configuration + # changes its selection; matches the lower bound the VSA paths already apply + # (VSASparseAttention.sparse / _h3_compute_topk). + topk = max(1, min(K, int(topk_ratio * K))) lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8) diff --git a/videox_fun/models/attention_utils.py b/videox_fun/models/attention_utils.py index 0330e2f..8068790 100644 --- a/videox_fun/models/attention_utils.py +++ b/videox_fun/models/attention_utils.py @@ -345,6 +345,32 @@ def attention( return out +def _head_shard_seq_sum(x, heads_total, head_start): + """``torch.sum(x, dim=-2, keepdim=True)`` that is bitwise-invariant to Ulysses head-sharding. + + ``x`` is ``[B, h_local, L, D]``. The pipeline runs under autocast and ``torch.sum`` is in + PyTorch's float32 promotion list, so this normalizer accumulates in fp32 over ``L`` and + returns fp32; the CUDA ReduceConfig picks its accumulation tree from the output count + ``B*H*D``. When head-parallelism shrinks ``H`` (e.g. 12 -> 6 -> 3) the tree changes and the + fp32 sum moves in its last bits (~1e-7), which the linear branch's denominator passes through + undamped -- a bf16 output would absorb it, fp32 does not. Padding the head dim back to + ``heads_total`` makes every rank reduce over the full ``B*heads_total*D`` outputs, i.e. the + single-GPU tree, then this shard's own head block is sliced back out, so multi-GPU reproduces + the single-GPU result bitwise. + + Uninitialized padding is safe: each output element reads only its own head's rows, so the + garbage heads (sliced away) never touch the real ones. ``heads_total=None`` or an already-full + head count falls back to the plain reduction, leaving the single-GPU path unchanged. + """ + h_local = x.shape[1] + if heads_total is None or h_local == heads_total: + return torch.sum(x, dim=-2, keepdim=True) + buf = torch.empty(x.shape[0], heads_total, x.shape[2], x.shape[3], + device=x.device, dtype=x.dtype) + buf[:, head_start:head_start + h_local].copy_(x) + return torch.sum(buf, dim=-2, keepdim=True)[:, head_start:head_start + h_local] + + class SparseLinearAttention(nn.Module): # Modified from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/kernel.py def __init__(self, head_dim, topk, feature_map='softmax', BLKQ=64, BLKK=64, use_bf16=True, tie_feature_map_qk=True): @@ -391,13 +417,18 @@ class SparseLinearAttention(nn.Module): nn.init.zeros_(self.proj_l.weight) nn.init.zeros_(self.proj_l.bias) - def forward(self, q, k, v, return_sparsity=False): + def forward(self, q, k, v, return_sparsity=False, heads_total=None, head_start=0): R''' Args: q: queries [B, L, H, D] k: keys [B, L, H, D]. v: values [B, L, H, D]. return_sparsity: whether to return the actual sparsity. + heads_total: full head count before Ulysses head-sharding, or None when unsharded. + Kept so the linear branch's fp32 normalizer sum stays bitwise-equal to the + single-GPU result under sequence parallelism (see _head_shard_seq_sum). + head_start: this shard's offset into the full head range; used only when heads_total + is not None. ''' dtype = q.dtype @@ -417,7 +448,7 @@ class SparseLinearAttention(nn.Module): def calc_linear(q, k, v): kvsum = k.transpose(-1, -2) @ v - ksum = torch.sum(k, dim=-2, keepdim=True) + ksum = _head_shard_seq_sum(k, heads_total, head_start) return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True)) o_l = calc_linear(q, k, v) @@ -943,7 +974,7 @@ class VSASparseAttention(nn.Module): oc = oc_blocks.repeat_interleave(self.BLKQ, dim=2)[:, :, :L] return oc - def forward_h3(self, q, k, v, x, segments, return_sparsity=False): + def forward_h3(self, q, k, v, x, segments, return_sparsity=False, gate=None): R'''Segment-aware H3 VSA -- the official FastVideo segment/tile mask semantics. Args: @@ -951,6 +982,12 @@ class VSASparseAttention(nn.Module): x: [B, L, C] attention-input hidden states (drives the compress gate). segments: packed layout -- dense row counts (``int``) and video region ``(t, h, w)`` token grid, in packed order. Built by the model forward. + gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; when given, ``x`` is + ignored. ``to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)`` and thus + mixes heads, so under Ulysses head-parallelism the caller must evaluate it on + its own sequence shard (where the FULL hidden dim is still present) and + all-to-all the result -- passing the head shard in here keeps this module + head-local. ``H`` is then the local head count. Returns: [B, L, H, D] attention output (before the output projection). @@ -985,7 +1022,8 @@ class VSASparseAttention(nn.Module): qb = _tile(q, self.dtype, 'q') kb = _tile(k, self.dtype, 'k') vb = _tile(v, self.dtype, 'v') - gate = self.to_gate_compress(x).view(B, L, H, D) + gate = self.to_gate_compress(x) if gate is None else gate + gate = gate.view(B, L, H, D) gb = _tile(gate, gate.dtype, 'g') pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles) # [B,H,n_tiles,D] fp32 @@ -1017,12 +1055,15 @@ class VSASparseAttention(nn.Module): return out, self.topk return out - def forward(self, q, k, v, x, return_sparsity=False, segments=None): + def forward(self, q, k, v, x, return_sparsity=False, segments=None, gate=None): R''' Args: q, k, v: [B, L, H, D] POST-RoPE queries/keys/values. x: [B, L, C] attention-input hidden states (drives the compress gate). return_sparsity: also return the sparse-branch keep ratio. + gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; see ``forward_h3`` + (required under Ulysses head-parallelism, where ``x`` is unavailable in the + head-shard layout because the gate Linear mixes heads). Returns: [B, L, H, D] attention output (before the output projection). ''' @@ -1031,7 +1072,8 @@ class VSASparseAttention(nn.Module): # kernel (forward + backward), so training and inference share one selection # semantics. Falls back to the flat 1-D block map only when the caller cannot # supply a clean segment layout (segments is None). - return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity) + return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity, + gate=gate) B, L, H, D = q.shape qb = q.transpose(1, 2).contiguous() # [B,H,L,D] kb = k.transpose(1, 2).contiguous() @@ -1042,7 +1084,8 @@ class VSASparseAttention(nn.Module): # Shorter than one block: exact dense attention (sparsity gives no benefit). o_s = F.scaled_dot_product_attention( qb.float(), kb.float(), vb.float()).to(qb.dtype) - gate = self.to_gate_compress(x).view(B, L, H, D).transpose(1, 2) # [B,H,L,D] + gate = self.to_gate_compress(x) if gate is None else gate + gate = gate.view(B, L, H, D).transpose(1, 2) # [B,H,L,D] o_c = self.compressed(qb, kb, vb) out = o_s + gate.to(o_s.dtype) * o_c.to(o_s.dtype) out = out.transpose(1, 2) # [B,L,H,D] diff --git a/videox_fun/models/fastvideo_minimax_h3_transformer3d.py b/videox_fun/models/fastvideo_minimax_h3_transformer3d.py index ffc573f..1cbf782 100644 --- a/videox_fun/models/fastvideo_minimax_h3_transformer3d.py +++ b/videox_fun/models/fastvideo_minimax_h3_transformer3d.py @@ -240,16 +240,23 @@ class FastVideoMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel): return tuple(segments) def enable_multi_gpus_inference(self,): - """VSA has no sequence-parallel path in this repo. + """Enable sequence-parallel inference for the VSA sparse student. - The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse forward; - there is no ulysses head-parallel or ring-attention binding for it (unlike the base - ``MiniMaxH3MultiGPUsAttnProcessor``). Fail loudly rather than silently falling back to a - DENSE multi-GPU processor (which would bypass VSA and drop the gate params). + Reuses the parent's packed-sequence SP assembly -- it sizes ``sp_world_size`` / + ``all_gather``, pads the sequence to a group multiple, hands every rank a contiguous row + slice and checks ``num_heads`` is divisible by the group -- then re-binds the VSA + sequence-parallel processor onto every block that carries an ``attn_op``. The parent + binding (``MiniMaxH3MultiGPUsAttnProcessor``) is a DENSE core; leaving it on a sparse + block would silently bypass VSA and drop the ``to_gate_compress`` params, so the sparse + blocks are swapped to ``MiniMaxH3MultiGPUsVSAAttnProcessor``, which turns the row split + into a head split around the VSA op. A dense student (``topk >= 1.0``) has no ``attn_op`` + and keeps the parent's dense processor untouched. """ - raise NotImplementedError( - "FastVideoMiniMaxH3Transformer3DModel (VSA) supports single-GPU inference only; the " - "VSA kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.") + super().enable_multi_gpus_inference() + from ..dist import MiniMaxH3MultiGPUsVSAAttnProcessor + for block in self.transformer_blocks: + if hasattr(block.attn, "attn_op"): + block.attn.set_processor(MiniMaxH3MultiGPUsVSAAttnProcessor()) def sparse_trainable_params(self): """Names of the extra VSA-branch parameters, for logging / freezing.""" diff --git a/videox_fun/models/fastvideo_wan_transformer3d.py b/videox_fun/models/fastvideo_wan_transformer3d.py index 8a69e83..8ec944f 100644 --- a/videox_fun/models/fastvideo_wan_transformer3d.py +++ b/videox_fun/models/fastvideo_wan_transformer3d.py @@ -67,26 +67,45 @@ class FastVideoWanSelfAttention(WanSelfAttention): q, k = rope_apply_qk(q, k, grid_sizes, freqs) + out = self._vsa_dispatch(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype), + seq_lens, grid_sizes) + x = out.to(dtype).flatten(2) + return self.o(x) + + def _vsa_dispatch(self, q, k, v, x, seq_lens, grid_sizes, gate=None): + """Run the VSA kernel over POST-RoPE ``q/k/v`` ([B, L, H, D]) -> [B, L, H, D]. + + ``x`` drives ``attn_op.to_gate_compress`` unless a pre-computed ``gate`` ([B, L, H, D]) + is handed in: that Linear mixes heads, so the sequence-parallel path + (``usp_attn_fastvideo_forward``) evaluates it on its own sequence shard and passes the + head-sharded result, with ``x`` left as None. + + The segment-aware tile geometry needs the sequence length to equal the true video token + count. Under sequence parallelism the gathered length is padded to a multiple of the SP + degree and can exceed it, so that case takes the per-sample branch below, which attends + over the valid grid tokens only and pads the output back to ``L``. + """ + s = q.shape[1] g0 = grid_sizes[0] if (bool((grid_sizes == g0).all()) and s == int(g0.prod()) and bool((seq_lens == s).all())): segments = ((int(g0[0]), int(g0[1]), int(g0[2])),) - out = self.attn_op(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype), - segments=segments) - else: - outputs = [] - for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())): - grid = tuple(int(axis) for axis in grid) - length = int(length) - if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s: - raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。") - oi = self.attn_op(q[i:i + 1, :length].to(dtype), k[i:i + 1, :length].to(dtype), - v[i:i + 1, :length].to(dtype), x[i:i + 1, :length].to(dtype), - segments=(grid,)) - outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length))) - out = torch.cat(outputs, dim=0) - x = out.to(dtype).flatten(2) - return self.o(x) + return self.attn_op(q, k, v, x, segments=segments, gate=gate) + + outputs = [] + for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())): + grid = tuple(int(axis) for axis in grid) + length = int(length) + if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s: + raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。") + gate_i = None if gate is None else gate[i:i + 1, :length] + # x is only consumed by to_gate_compress, so it is dropped once gate is supplied. + x_i = None if gate is not None else x[i:i + 1, :length] + oi = self.attn_op(q[i:i + 1, :length], k[i:i + 1, :length], + v[i:i + 1, :length], x_i, + segments=(grid,), gate=gate_i) + outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length))) + return torch.cat(outputs, dim=0) class FastVideoWanTransformer3DModel(WanTransformer3DModel): @@ -187,17 +206,45 @@ class FastVideoWanTransformer3DModel(WanTransformer3DModel): self._apply_sparse_self_attention(topk) def enable_multi_gpus_inference(self,): - """VSA has no sequence-parallel path in this repo. + """Bind the VSA-aware USP forward instead of the parent's dense one. - The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse - forward; there is no ulysses head-parallel or ring-attention binding for it (unlike - the RoLa/SLA ``usp_attn_*`` forwards). The 1.3B student fits comfortably on one GPU, - so run inference single-device. Fail loudly rather than silently falling back to the - parent's DENSE ``usp_attn_forward`` (which would bypass VSA and drop the gate params). + ``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every + ``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via + ``xFuserLongContextAttention``), which never reaches ``FastVideoWanSelfAttention.attn_op`` + and would silently drop the VSA kernel plus the trained ``to_gate_compress`` gate. + Rebind to ``usp_attn_fastvideo_forward``: Ulysses head-parallel, with the cross-head + compress gate evaluated on each rank's sequence shard (where the full hidden dim still + lives) and all-to-all'd next to q/k/v. + + Constraint: the VSA kernel has no ring-attention path (no partial-output/logsumexp + interface to rescale across KV chunks), so this is pure Ulysses -- ``ring_degree`` must + be 1 and ``num_heads`` must divide the SP degree. """ - raise NotImplementedError( - "FastVideoWanTransformer3DModel (VSA) supports single-GPU inference only; the VSA " - "kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.") + import types + + from ..dist import (get_sequence_parallel_rank, + get_sequence_parallel_world_size, get_sp_group, + usp_attn_fastvideo_forward) + self.sp_world_size = get_sequence_parallel_world_size() + self.sp_world_rank = get_sequence_parallel_rank() + self.all_gather = get_sp_group().all_gather + + # Ulysses head-split needs num_heads divisible by the SP degree. + n_heads = self.blocks[0].self_attn.num_heads + if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0: + raise ValueError( + f"VSA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the " + f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree " + f"accordingly and keep ring_degree=1 (the VSA kernel has no ring path).") + + for block in self.blocks: + block.self_attn.forward = types.MethodType( + usp_attn_fastvideo_forward, block.self_attn) + + if hasattr(self, 'vace_blocks'): + for block in self.vace_blocks: + block.self_attn.forward = types.MethodType( + usp_attn_fastvideo_forward, block.self_attn) def sparse_trainable_params(self): """Names of the extra VSA-branch parameters, for logging / freezing.""" diff --git a/videox_fun/models/minimax_h3_conversion.py b/videox_fun/models/minimax_h3_conversion.py index 7dd8f85..47a73d9 100644 --- a/videox_fun/models/minimax_h3_conversion.py +++ b/videox_fun/models/minimax_h3_conversion.py @@ -111,10 +111,13 @@ class MiniMaxH3MixedPrecisionLoaderMixin: kwargs["device_map"] = device_map if max_memory is not None: kwargs["max_memory"] = max_memory - model = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs) + result = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs) + # ModelMixin returns (model, loading_info) when output_loading_info=True. + # Restore precision on the model while preserving that return contract. + model = result[0] if isinstance(result, tuple) else result if torch_dtype not in (None, torch.float32) and cls._keep_in_fp32_modules: restore_fp32_modules(model, pretrained_model_name_or_path, kwargs.get("subfolder", None)) - return model + return result @classmethod def _from_pretrained_with_missing_keys(cls, pretrained_model_name_or_path, torch_dtype, **kwargs): diff --git a/videox_fun/models/minimax_h3_transformer3d.py b/videox_fun/models/minimax_h3_transformer3d.py index f3becb3..875a093 100644 --- a/videox_fun/models/minimax_h3_transformer3d.py +++ b/videox_fun/models/minimax_h3_transformer3d.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import Any, Dict, Optional, Tuple, Union @@ -264,7 +264,11 @@ class MiniMaxH3Attention(nn.Module): rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, valid_length: Optional[int] = None, + vsa_segments=None, ) -> torch.Tensor: + if vsa_segments is not None: + return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length, + vsa_segments=vsa_segments) return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length) @@ -374,6 +378,7 @@ class MiniMaxH3TransformerBlock(nn.Module): rotary_emb: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor] = None, valid_length: Optional[int] = None, + vsa_segments=None, ) -> torch.Tensor: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb) @@ -382,7 +387,8 @@ class MiniMaxH3TransformerBlock(nn.Module): norm_hidden_states = norm_hidden_states * ( 1.0 + _index_select_rows(scale_msa, adaln_indices) ) + _index_select_rows(shift_msa, adaln_indices) - attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length) + attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length, + vsa_segments=vsa_segments) hidden_states = residual + _index_select_rows(gate_msa, adaln_indices) * attn_output residual = hidden_states @@ -546,6 +552,9 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin self.gradient_checkpointing = False self.gradient_checkpointing_save_on_cpu = False + # Optional non-reentrant checkpoint `context_fn` for the main blocks; the streaming KV cache installs one so + # backward recomputes see the cache view of the original forward. + self.gradient_checkpointing_context_fn = None self.sp_world_size = 1 self.sp_world_rank = 0 self.all_gather = None @@ -747,10 +756,14 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin adaln_indices = adaln_indices[rows] rotary_emb = (rotary_emb[0][rows], rotary_emb[1][rows]) + vsa_segments = attention_kwargs.get("vsa_segments") if attention_kwargs else None + block_kwargs = {"vsa_segments": vsa_segments} if vsa_segments is not None else {} for block in self.transformer_blocks: with torch.autograd.graph.save_on_cpu() if self.gradient_checkpointing_save_on_cpu else nullcontext(): if torch.is_grad_enabled() and self.gradient_checkpointing: ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + if self.gradient_checkpointing_context_fn is not None: + ckpt_kwargs["context_fn"] = self.gradient_checkpointing_context_fn hidden_states = torch.utils.checkpoint.checkpoint( block, hidden_states, @@ -759,10 +772,12 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin rotary_emb, attention_mask, valid_length, + **block_kwargs, **ckpt_kwargs, ) else: - hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask, valid_length) + hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask, + valid_length, **block_kwargs) if self.sp_world_size > 1: hidden_states = self.all_gather(hidden_states.contiguous(), dim=1)[:, :sequence_length] @@ -822,11 +837,20 @@ class MiniMaxH3StreamingKVCache: Args: num_layers (`int`): The number of main transformer blocks (50 for the released H3). dtype (`torch.dtype`, defaults to `torch.bfloat16`): The dtype the clean K/V is stored at. + offload_to_cpu (`bool`, defaults to `False`): Whether the clean K/V is parked on the CPU. + commit_attend_history (`bool`, defaults to `False`): Whether a clean commit pass runs the + live attention pattern (media rows attend `[text K/V; persistent history; current chunk + K/V]`, text rows self-attend) instead of the legacy standalone pass over + `[text | current chunk]`. Enable it when the history must condition the committed K/V, + as in the Self-Forcing clean refresh pass. """ - def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16): + def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16, offload_to_cpu: bool = False, + commit_attend_history: bool = False): self.layer_names = tuple(f"transformer_blocks.{index}.attn" for index in range(num_layers)) self.dtype = dtype + self.offload_to_cpu = offload_to_cpu + self.commit_attend_history = commit_attend_history self._history: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {} self._staged: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {} self._commit_active = False @@ -838,6 +862,10 @@ class MiniMaxH3StreamingKVCache: # Per-forward attention mode: None ("plain"), "live" or "commit". self.mode: Optional[str] = None self.token_tags: Optional[torch.Tensor] = None + # Activation-checkpoint replay: while set, `history_pair` exposes only this many leading history rows (the + # history a checkpointed forward saw); `_epoch` changes whenever rows are dropped, invalidating pending replays. + self._history_view_rows: Optional[int] = None + self._epoch = 0 # ------------------------------------------------------------------ # Per-forward setup consumed by the streaming processor. @@ -902,11 +930,51 @@ class MiniMaxH3StreamingKVCache: self._staged_row_count = int(key.shape[0]) elif int(key.shape[0]) != self._staged_row_count: raise RuntimeError(f"{layer_name}: staged row count differs across layers") - self._staged[layer_name] = (key, value) + device = torch.device("cpu") if self.offload_to_cpu else key.device + self._staged[layer_name] = ( + key.detach().to(device=device, dtype=self.dtype).contiguous(), + value.detach().to(device=device, dtype=self.dtype).contiguous(), + ) def history_pair(self, layer_name: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: r"""Return the persistent `(key, value)` of one layer, or `None` when empty.""" - return self._history.get(layer_name) + pair = self._history.get(layer_name) + if pair is None or self._history_view_rows is None: + return pair + if self._history_view_rows == 0: + return None + return pair[0][: self._history_view_rows], pair[1][: self._history_view_rows] + + # ------------------------------------------------------------------ + # Activation checkpointing. + # ------------------------------------------------------------------ + def checkpoint_context(self): + r""" + `context_fn` for non-reentrant `torch.utils.checkpoint` over a main block. + + Called when the checkpointed forward runs, it snapshots the cache view that forward sees — the mode, the + token tags and the committed history length — and returns `(forward_context, recompute_context)`; the + recompute context restores that view while backward replays the block, after `end_forward()` cleared the + mode and later commits appended more history. The history is append-only (`commit` concatenates), so the + snapshot is a row count, not a tensor copy: the replay reads a prefix view of the current history. + """ + history_rows = self.history_tokens if self._history_view_rows is None else self._history_view_rows + snapshot = (self.mode, self.token_tags, history_rows, self._epoch) + return nullcontext(), self._replay_view(*snapshot) + + @contextmanager + def _replay_view(self, mode, token_tags, history_rows, epoch): + if epoch != self._epoch: + raise RuntimeError( + "the streaming KV history was cleared or trimmed between a checkpointed forward and its backward " + "recompute; clear the cache only after backward" + ) + saved = (self.mode, self.token_tags, self._history_view_rows) + self.mode, self.token_tags, self._history_view_rows = mode, token_tags, history_rows + try: + yield + finally: + self.mode, self.token_tags, self._history_view_rows = saved # ------------------------------------------------------------------ # Transaction control. @@ -931,6 +999,7 @@ class MiniMaxH3StreamingKVCache: torch.cat((previous[1], current[1]), dim=0), ) self._history[layer_name] = combined + del current, previous, combined self._commit_token_counts.append(len(self._commit_tags)) self._commit_token_tags.append(self._commit_tags) self._commit_active = False @@ -1005,6 +1074,7 @@ class MiniMaxH3StreamingKVCache: def clear(self) -> None: r"""Clear persistent and staged KV and reset the logical chunk count.""" self._history.clear() + self._epoch += 1 self._commit_token_counts.clear() self._commit_token_tags.clear() self._staged.clear() @@ -1014,6 +1084,7 @@ class MiniMaxH3StreamingKVCache: self.end_forward() def _retain_commit_rows(self, selection) -> None: + self._epoch += 1 offsets = [0] for count in self._commit_token_counts: offsets.append(offsets[-1] + count) @@ -1068,9 +1139,11 @@ class MiniMaxH3StreamingKVCacheAttnProcessor: [`MiniMaxH3AttnProcessor`]; * *live*: text rows self-attend, while media rows attend to `[text K/V; persistent clean history; current chunk K/V]`; - * *commit*: a plain pass whose output is discarded, staging the media rows' K/V of - every layer into the cache (the K/V of a row depends only on that row's input, so - the commit pass's attention output is irrelevant). + * *commit*: the pass whose output is discarded, staging the media rows' K/V of every + layer into the cache. With `commit_attend_history` off (legacy) it is a plain pass over + `[text | current chunk]`; with it on, the commit pass runs the live attention pattern so + the staged K/V is conditioned on the persistent history, matching the Self-Forcing clean + refresh pass. Live and commit passes must be padless (`attention_mask is None`); the streaming packed sequences are single documents by construction. @@ -1107,19 +1180,19 @@ class MiniMaxH3StreamingKVCacheAttnProcessor: query = apply_minimax_h3_rotary_emb(query, *rotary_emb) key = apply_minimax_h3_rotary_emb(key, *rotary_emb) - if cache.mode == "commit": - # The committed rows' K/V depend only on their own inputs; the attention output - # of a commit pass is discarded by the caller. + if cache.mode == "commit" and not cache.commit_attend_history: + # Legacy commit: a plain pass over [text | current chunk]; the attention output is + # discarded by the caller and only the staged media K/V rows matter. output = attention(query, key, value, causal=False) rows = cache.commit_row_indices.to(key.device) staged_key = key.index_select(1, rows).squeeze(0) staged_value = value.index_select(1, rows).squeeze(0) cache.stage( self.layer_name, - staged_key.detach().to(cache.dtype).contiguous(), - staged_value.detach().to(cache.dtype).contiguous(), + staged_key, + staged_value, ) - else: # live + else: # live, or a commit with commit_attend_history: media rows attend the history. tags = cache.token_tags text_rows = tags == MINIMAX_H3_STREAMING_TEXT_TAG media_rows = ~text_rows @@ -1153,6 +1226,16 @@ class MiniMaxH3StreamingKVCacheAttnProcessor: values.transpose(1, 2), ).transpose(1, 2) + if cache.mode == "commit" and cache.commit_attend_history: + # The history-conditioned commit pass output is still discarded; only the staged K/V + # rows matter, now computed with the history in context. + rows = cache.commit_row_indices.to(key.device) + cache.stage( + self.layer_name, + key.index_select(1, rows).squeeze(0), + value.index_select(1, rows).squeeze(0), + ) + output = output.flatten(2, 3).type_as(query) output = attn.to_out[0](output) output = attn.to_out[1](output) @@ -1201,21 +1284,26 @@ def install_minimax_h3_streaming_kv_cache( transformer: "MiniMaxH3Transformer3DModel", *, dtype: torch.dtype = torch.bfloat16, + offload_to_cpu: bool = False, + commit_attend_history: bool = False, ) -> MiniMaxH3StreamingKVCache: r""" Swap every main block's attention processor for the streaming one and return the shared cache. The token refiner keeps its standard processor: the streaming contract only concerns the - main block stack over the media sequence. + main block stack over the media sequence. `commit_attend_history` is forwarded to + [`MiniMaxH3StreamingKVCache`]. """ cache = MiniMaxH3StreamingKVCache( - num_layers=len(transformer.transformer_blocks), dtype=dtype + num_layers=len(transformer.transformer_blocks), dtype=dtype, offload_to_cpu=offload_to_cpu, + commit_attend_history=commit_attend_history, ) for index, block in enumerate(transformer.transformer_blocks): block.attn.set_processor( MiniMaxH3StreamingKVCacheAttnProcessor(cache, cache.layer_names[index]) ) + transformer.gradient_checkpointing_context_fn = cache.checkpoint_context return cache @@ -1223,3 +1311,4 @@ def uninstall_minimax_h3_streaming_kv_cache(transformer: "MiniMaxH3Transformer3D r"""Restore the standard attention processor on every main block.""" for block in transformer.transformer_blocks: block.attn.set_processor(MiniMaxH3AttnProcessor()) + transformer.gradient_checkpointing_context_fn = None diff --git a/videox_fun/models/turbowan_transformer3d.py b/videox_fun/models/turbowan_transformer3d.py index 0703efa..ad96331 100644 --- a/videox_fun/models/turbowan_transformer3d.py +++ b/videox_fun/models/turbowan_transformer3d.py @@ -51,8 +51,8 @@ class FakeLocalAttention(nn.Module): super().__init__() self.local_attn = SparseLinearAttention(head_dim, topk) - def forward(self, q, k, v): - x = self.local_attn(q, k, v) + def forward(self, q, k, v, heads_total=None, head_start=0): + x = self.local_attn(q, k, v, heads_total=heads_total, head_start=head_start) return x @@ -310,3 +310,45 @@ class TurboWanTransformer3DModel(WanTransformer3DModel): if hasattr(self, "img_emb"): del self.img_emb + + def enable_multi_gpus_inference(self,): + """Bind the SLA-aware USP forward instead of the parent's dense one. + + ``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every + ``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via + ``xFuserLongContextAttention``). That path only touches base-class attributes + (``q/k/v/norm_q/norm_k/o/window_size/num_heads/head_dim``), all of which this module's + ``WanSelfAttention`` inherits, so the inherited patch used to degrade the SLA student to + dense attention SILENTLY -- no AttributeError, no shape error, just ``self.attn_op`` + (``FakeLocalAttention`` -> ``SparseLinearAttention``) never being reached. Rebind to + ``usp_attn_turbowan_forward``: Ulysses head-parallel all-to-all around the SLA kernel, + which is head-independent (all reductions over L/D, ``proj_l`` is per-head), so the + multi-GPU result matches single-GPU. + + Constraint: the SLA kernel exposes no partial-output/logsumexp interface, so ring + attention cannot be wired to it -- pure Ulysses only, ``ring_degree`` must be 1 and + ``num_heads`` must divide the SP degree. + """ + from ..dist import (get_sequence_parallel_rank, + get_sequence_parallel_world_size, get_sp_group, + usp_attn_turbowan_forward) + self.sp_world_size = get_sequence_parallel_world_size() + self.sp_world_rank = get_sequence_parallel_rank() + self.all_gather = get_sp_group().all_gather + + # Ulysses head-split needs num_heads divisible by the SP degree. + n_heads = self.blocks[0].self_attn.num_heads + if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0: + raise ValueError( + f"SLA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the " + f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree " + f"accordingly and keep ring_degree=1 (the SLA kernel has no ring path).") + + for block in self.blocks: + block.self_attn.forward = types.MethodType( + usp_attn_turbowan_forward, block.self_attn) + + if hasattr(self, 'vace_blocks'): + for block in self.vace_blocks: + block.self_attn.forward = types.MethodType( + usp_attn_turbowan_forward, block.self_attn) diff --git a/videox_fun/pipeline/pipeline_wan.py b/videox_fun/pipeline/pipeline_wan.py index 0dc024b..a65ab3b 100755 --- a/videox_fun/pipeline/pipeline_wan.py +++ b/videox_fun/pipeline/pipeline_wan.py @@ -1,5 +1,6 @@ import inspect import math +import time from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -21,6 +22,17 @@ from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000): + """Official FlashHead timestep schedule with shift transform.""" + if num_inference_steps == 4: + timesteps = [1000, 750, 500, 250] + else: + timesteps = np.linspace(num_timesteps, 1, num_inference_steps, dtype=np.float32).tolist() + timesteps = torch.tensor(timesteps + [0.0], dtype=torch.float32, device=device) + t = timesteps / num_timesteps + return shift * t / (1 + (shift - 1) * t) * num_timesteps + + EXAMPLE_DOC_STRING = """ Examples: ```python @@ -409,6 +421,7 @@ class WanPipeline(DiffusionPipeline): max_sequence_length: int = 512, comfyui_progressbar: bool = False, shift: int = 5, + stochastic_sampling: bool = False, ) -> Union[WanPipelineOutput, Tuple]: """ Function invoked when calling the pipeline for generation. @@ -471,7 +484,9 @@ class WanPipeline(DiffusionPipeline): in_prompt_embeds = prompt_embeds # 4. Prepare timesteps - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): + if stochastic_sampling: + timesteps = stochastic_sampling_timesteps(num_inference_steps, shift, device) + elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1) elif isinstance(self.scheduler, FlowUniPCMultistepScheduler): self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift) @@ -513,8 +528,9 @@ class WanPipeline(DiffusionPipeline): # 7. Denoising loop num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self.transformer.num_inference_steps = num_inference_steps - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): + denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps + with self.progress_bar(total=len(denoise_timesteps)) as progress_bar: + for i, t in enumerate(denoise_timesteps): self.transformer.current_steps = i if self.interrupt: @@ -528,6 +544,8 @@ class WanPipeline(DiffusionPipeline): timestep = t.expand(latent_model_input.shape[0]) # predict noise model_output + torch.cuda.synchronize() + forward_start_time = time.time() with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device): noise_pred = self.transformer( x=latent_model_input, @@ -535,6 +553,9 @@ class WanPipeline(DiffusionPipeline): t=timestep, seq_len=seq_len, ) + torch.cuda.synchronize() + forward_time = time.time() - forward_start_time + print(f"Step {i}/{len(timesteps)} - Transformer forward time: {forward_time:.4f}s") # perform guidance if do_classifier_free_guidance: @@ -542,7 +563,15 @@ class WanPipeline(DiffusionPipeline): noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + if stochastic_sampling: + t_i = (timesteps[i] / 1000).to(weight_dtype) + t_i_1 = (timesteps[i + 1] / 1000).to(weight_dtype) + x_0 = latents - noise_pred * t_i + latents = (1 - t_i_1) * x_0 + t_i_1 * torch.randn( + latents.shape, dtype=latents.dtype, device=device, generator=generator + ) + else: + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] if callback_on_step_end is not None: callback_kwargs = {} @@ -554,7 +583,7 @@ class WanPipeline(DiffusionPipeline): prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if i == len(denoise_timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if comfyui_progressbar: pbar.update(1)