[LoRA] Fix lora merge weights (#579)

.buildkite/pipeline.yml

@@ -1,66 +1,138 @@
 env:
   IMAGE_VERSION: "py3.12-latest"
+  BUILDKITE_CLEAN_CHECKOUT: true
 
 steps:
-  - block: "Start Build"
-    blocked_state: "running"
-    prompt: "Approve build?"
+  - label: "pre-commit"
+    command: ".buildkite/scripts/pre_commit.sh"
+    agents:
+      queue: "default"
+
+  - wait
 
   - label: "Trigger Tests"
-    command: |
-      echo "Current working directory: $(pwd)"
-      echo "Current branch:"
-      git branch --show-current
-      echo "Full diff:"
-      git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD
     plugins:
       - monorepo-diff#v1.4.0:
           diff: "git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD"
           watch:
             - path: 
                 - "fastvideo/v1/models/encoders/**"
-                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/models/loader/**"
                 - "fastvideo/v1/tests/encoders/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: "Encoder Tests"
                 env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                   - TEST_TYPE=encoder
                 agents:
                   queue: "default"
             - path:
                 - "fastvideo/v1/models/vaes/**"
-                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/models/loader/**"
                 - "fastvideo/v1/tests/vaes/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: "VAE Tests"
                 env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                   - TEST_TYPE=vae
                 agents:
                   queue: "default"
             - path:
                 - "fastvideo/v1/models/dits/**"
-                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/models/loader/**"
                 - "fastvideo/v1/tests/transformers/**"
                 - "fastvideo/v1/layers/**"
                 - "fastvideo/v1/attention/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: "Transformer Tests"
                 env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                   - TEST_TYPE=transformer
                 agents:
                   queue: "default"
-            - path: "fastvideo/v1/**/*.py"
+            - path: 
+                - "fastvideo/v1/**/*.py"
               config:
-                command: "timeout 60m .buildkite/scripts/pr_test.sh"
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                 label: "SSIM Tests"
                 env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                   - TEST_TYPE=ssim
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Training Tests"
+                env:
+                  - TEST_TYPE=training
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Training Tests VSA"
+                env:
+                  - TEST_TYPE=training_vsa
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Inference Tests STA"
+                env:
+                  - TEST_TYPE=inference_sta
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests STA"
+                env:
+                  - TEST_TYPE=precision_sta
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests VSA"
+                env:
+                  - TEST_TYPE=precision_vsa
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -31,6 +31,10 @@ log "Setting up Modal authentication from Buildkite secrets..."
 MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
 MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)
 
+WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+
+WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+
 if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
     log "Retrieved Modal credentials from Buildkite secrets"
     python3 -m modal token set --token-id "$MODAL_TOKEN_ID" --token-secret "$MODAL_TOKEN_SECRET" --profile buildkite-ci --activate --verify
@@ -54,22 +58,44 @@ if [ -z "${TEST_TYPE:-}" ]; then
 fi
 log "Test type: $TEST_TYPE"
 
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$BUILDKITE_PULL_REQUEST IMAGE_VERSION=$IMAGE_VERSION"
+
 case "$TEST_TYPE" in
     "encoder")
         log "Running encoder tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
         ;;
     "vae")
         log "Running VAE tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
         ;;
     "transformer")
         log "Running transformer tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
         ;;
     "ssim")
         log "Running SSIM tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
+        ;;
+    "training")
+        log "Running training tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
+        ;;
+    "training_vsa")
+        log "Running training VSA tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
+        ;;
+    "inference_sta")
+        log "Running inference STA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
+        ;;
+    "precision_sta")
+        log "Running precision STA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
+        ;;
+    "precision_vsa")
+        log "Running precision VSA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
         ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"

.buildkite/scripts/pre_commit.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+set -uo pipefail
+
+log() {
+  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
+}
+
+log "=== Starting pre-commit checks ==="
+
+cd "$(dirname "$0")/../.."
+PROJECT_ROOT=$(pwd)
+log "Project root: $PROJECT_ROOT"
+
+if ! python3 -m pre_commit --version &> /dev/null; then
+    log "pre-commit not found, installing..."
+    python3 -m pip install --user pre-commit==4.0.1
+    
+    if ! python3 -m pre_commit --version &> /dev/null; then
+        log "Error: Failed to install pre-commit."
+        exit 1
+    fi
+fi
+
+log "Pre-commit version: $(python3 -m pre_commit --version)"
+
+log "Installing/updating pre-commit hooks..."
+python3 -m pre_commit install --install-hooks
+
+log "Running pre-commit checks on all files..."
+python3 -m pre_commit run --all-files
+PRE_COMMIT_EXIT_CODE=$?
+
+if [ $PRE_COMMIT_EXIT_CODE -eq 0 ]; then
+    log "Pre-commit checks completed successfully"
+else
+    log "Error: Pre-commit checks failed with exit code: $PRE_COMMIT_EXIT_CODE"
+fi
+
+log "=== Pre-commit checks completed with exit code: $PRE_COMMIT_EXIT_CODE ==="
+exit $PRE_COMMIT_EXIT_CODE 

.github/workflows/pr-test.yml

@@ -264,7 +264,7 @@ jobs:
     with:
       job_id: "training-test-VSA"
       gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
+      gpu_count: 2
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
@@ -284,7 +284,7 @@ jobs:
     with:
       job_id: "inference-test-STA"
       gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
+      gpu_count: 2
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"

README.md

@@ -15,6 +15,11 @@ With FastVideo's optimizations, you can achieve more than 3x inference improveme
 <img src=assets/perf.png width="90%"/>
 </div>
 
+## NEWS
+- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
+- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
+- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+
 ## Key Features
 
 FastVideo has the following features:
@@ -128,6 +133,15 @@ We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support thro
 If you use FastVideo for your research, please cite our paper:
 
 ```bibtex
+@misc{zhang2025vsafastervideodiffusion,
+      title={VSA: Faster Video Diffusion with Trainable Sparse Attention}, 
+      author={Peiyuan Zhang and Haofeng Huang and Yongqi Chen and Will Lin and Zhengzhong Liu and Ion Stoica and Eric Xing and Hao Zhang},
+      year={2025},
+      eprint={2505.13389},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2505.13389}, 
+}
 @misc{zhang2025fastvideogenerationsliding,
       title={Fast Video Generation with Sliding Tile Attention},
       author={Peiyuan Zhang and Yongqi Chen and Runlong Su and Hangliang Ding and Ion Stoica and Zhenghong Liu and Hao Zhang},

examples/training/finetune/wan_i2v_14b_480p/crush_smol/README.md

@@ -0,0 +1,10 @@
+This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+
+Execute the following commands from `FastVideo/` to run training:
+
+- Download crush-smol dataset:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
+- Preprocess the videos and captions into latents:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
+- Edit the following file and run finetuning:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`

examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh

@@ -22,4 +22,5 @@ torchrun --nproc_per_node=$GPU_NUM \
     --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
     --preprocess_task "i2v" 

examples/training/finetune/wan_t2v_1_3b/crush_smol/README.md

@@ -0,0 +1,10 @@
+This directory contain e2e examples scripts for finetuning Wan2.1 T2v. 
+
+Execute the following commands from `FastVideo/` to run training:
+
+- Download crush-smol dataset:
+`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
+- Preprocess the videos and captions into latents:
+`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
+- Edit the following file and run finetuning:
+`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`

examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh

@@ -22,4 +22,5 @@ torchrun --nproc_per_node=$GPU_NUM \
     --validation_dataset_file $VALIDATION_PATH \
     --samples_per_file 8 \
     --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
     --preprocess_task "t2v" 

fastvideo/v1/configs/models/dits/wanvideo.py

@@ -13,73 +13,6 @@ def is_blocks(n: str, m) -> bool:
 class WanVideoArchConfig(DiTArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
 
-    _videox_param_names_mapping: dict = field(
-        default_factory=lambda: {
-            r"^patch_embedding\.(.*)$":
-            r"patch_embedding.proj.\1",
-            r"^text_embedding\.0\.(.*)$":
-            r"condition_embedder.text_embedder.fc_in.\1",
-            r"^text_embedding\.2\.(.*)$":
-            r"condition_embedder.text_embedder.fc_out.\1",
-            r"^time_embedding\.0\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_in.\1",
-            r"^time_embedding\.2\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            r"^time_projection\.1\.(.*)$":
-            r"condition_embedder.time_modulation.linear.\1",
-            r"^img_emb\.proj\.0\.(.*)$":
-            r"condition_embedder.image_embedder.norm1.\1",
-            r"^img_emb\.proj\.1\.(.*)$":
-            r"condition_embedder.image_embedder.ff.fc_in.\1",
-            r"^img_emb\.proj\.3\.(.*)$":
-            r"condition_embedder.image_embedder.ff.fc_out.\1",
-            r"^img_emb\.proj\.4\.(.*)$":
-            r"condition_embedder.image_embedder.norm2.\1",
-            r"^head\.modulation":
-            r"scale_shift_table",
-            r"^head\.head\.(.*)$":
-            r"proj_out.\1",
-            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
-            r"blocks.\1.to_q.\2",
-            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
-            r"blocks.\1.to_k.\2",
-            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
-            r"blocks.\1.to_v.\2",
-            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
-            r"blocks.\1.to_out.\2",
-            r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
-            r"blocks.\1.norm_q.\2",
-            r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
-            r"blocks.\1.norm_k.\2",
-            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
-            r"blocks.\1.attn2.to_q.\2",
-            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
-            r"blocks.\1.attn2.to_k.\2",
-            r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
-            r"blocks.\1.attn2.add_k_proj.\2",
-            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
-            r"blocks.\1.attn2.to_v.\2",
-            r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
-            r"blocks.\1.attn2.add_v_proj.\2",
-            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
-            r"blocks.\1.attn2.to_out.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
-            r"blocks.\1.attn2.norm_q.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
-            r"blocks.\1.attn2.norm_k.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
-            r"blocks.\1.attn2.norm_added_k.\2",
-            r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
-            r"blocks.\1.ffn.fc_in.\2",
-            r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
-            r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.modulation":
-            r"blocks.\1.scale_shift_table",
-            r"^blocks\.(\d+)\.norm3\.(.*)$":
-            r"blocks.\1.self_attn_residual_norm.norm.\2",
-        }
-    )
-
     _param_names_mapping: dict = field(
         default_factory=lambda: {
             r"^patch_embedding\.(.*)$":

fastvideo/v1/configs/pipelines/base.py

@@ -52,7 +52,7 @@ class PipelineConfig:
 
     # VAE configuration
     vae_config: VAEConfig = field(default_factory=VAEConfig)
-    vae_precision: str = "fp16"
+    vae_precision: str = "fp32"
     vae_tiling: bool = True
     vae_sp: bool = True
 

fastvideo/v1/configs/pipelines/wan.py

@@ -50,7 +50,7 @@ class WanT2V480PConfig(PipelineConfig):
 
     # Precision for each component
     precision: str = "bf16"
-    vae_precision: str = "fp16"
+    vae_precision: str = "fp32"
     text_encoder_precisions: Tuple[str, ...] = field(
         default_factory=lambda: ("fp32", ))
 

fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json

@@ -4,7 +4,7 @@
   "use_cpu_offload": true,
   "disable_autocast": false,
   "precision": "bf16",
-  "vae_precision": "fp16",
+  "vae_precision": "fp32",
   "vae_tiling": false,
   "vae_sp": false,
   "vae_config": {

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

@@ -4,7 +4,7 @@
   "use_cpu_offload": true,
   "disable_autocast": false,
   "precision": "bf16",
-  "vae_precision": "fp16",
+  "vae_precision": "fp32",
   "vae_tiling": false,
   "vae_sp": false,
   "vae_config": {

fastvideo/v1/fastvideo_args.py

@@ -413,7 +413,7 @@ class TrainingArgs(FastVideoArgs):
     lr_scheduler: str = "constant"
     lr_warmup_steps: int = 0
     max_grad_norm: float = 0.0
-    gradient_checkpointing: bool = False
+    enable_gradient_checkpointing_type: Optional[str] = None
     selective_checkpointing: float = 0.0
     allow_tf32: bool = False
     mixed_precision: str = ""
@@ -612,9 +612,11 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--max-grad-norm",
                             type=float,
                             help="Maximum gradient norm")
-        parser.add_argument("--gradient-checkpointing",
-                            action=StoreBoolean,
-                            help="Whether to use gradient checkpointing")
+        parser.add_argument("--enable-gradient-checkpointing-type",
+                            type=str,
+                            choices=["full", "ops", "block_skip"],
+                            default=None,
+                            help="Gradient checkpointing type")
         parser.add_argument("--selective-checkpointing",
                             type=float,
                             help="Selective checkpointing threshold")

fastvideo/v1/layers/layernorm.py

@@ -38,6 +38,12 @@ class RMSNorm(CustomOp):
         if self.has_weight:
             self.weight = nn.Parameter(self.weight)
 
+    # if we do fully_shard(model.layer_norm), and we call layer_form.forward_native(input) instead of layer_norm(input),
+    # we need to call model.layer_norm.register_fsdp_forward_method(model, "forward_native") to make sure fsdp2 hooks are triggered
+    # for mixed precision and cpu offloading
+
+    # the even better way might be fully_shard(model.layer_norm, mp_policy=, cpu_offloading=), and call model.layer_norm(input). everything should work out of the box
+    # because fsdp2 hooks will be triggered with model.layer_norm.__call__
     def forward_native(
         self,
         x: torch.Tensor,

fastvideo/v1/layers/lora/linear.py

@@ -71,15 +71,16 @@ class BaseLayerWithLoRA(nn.Module):
                 f"cuda:{torch.cuda.current_device()}").full_tensor()
             data += (self.slice_lora_b_weights(self.lora_B)
                      @ self.slice_lora_a_weights(self.lora_A)).to(data)
-            self.base_layer.weight.data = distribute_tensor(
-                data, mesh, placements=placements).to(current_device)
+            self.base_layer.weight = nn.Parameter(
+                distribute_tensor(data, mesh,
+                                  placements=placements).to(current_device))
         else:
             current_device = self.base_layer.weight.data.device
-            data = self.base_layer.weight.data.to(
+            data = self.base_layer.weight.to(
                 f"cuda:{torch.cuda.current_device()}")
             data += \
                 (self.slice_lora_b_weights(self.lora_B) @ self.slice_lora_a_weights(self.lora_A)).to(data)
-            self.base_layer.weight.data = data.to(current_device)
+            self.base_layer.weight = nn.Parameter(data.to(current_device))
         self.merged = True
 
     @torch.no_grad()
@@ -106,8 +107,8 @@ class BaseLayerWithLoRA(nn.Module):
                 f"cuda:{torch.cuda.current_device()}").full_tensor()
             data -= self.slice_lora_b_weights(
                 self.lora_B) @ self.slice_lora_a_weights(self.lora_A)
-            self.base_layer.weight.data = distribute_tensor(
-                data, mesh, placements=placement).to(device)
+            self.base_layer.weight = nn.Parameter(
+                distribute_tensor(data, mesh, placements=placement).to(device))
         else:
             self.base_layer.weight.data -= \
                 self.slice_lora_b_weights(self.lora_B) @\

fastvideo/v1/models/dits/wanvideo.py

@@ -518,7 +518,6 @@ class WanTransformer3DModel(CachableDiT):
     _supported_attention_backends = WanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = WanVideoConfig()._param_names_mapping
-    _videox_param_names_mapping = WanVideoConfig()._videox_param_names_mapping
     _reverse_param_names_mapping = WanVideoConfig()._reverse_param_names_mapping
     _lora_param_names_mapping = WanVideoConfig()._lora_param_names_mapping
 

fastvideo/v1/models/loader/component_loader.py

@@ -379,19 +379,10 @@ class TransformerLoader(ComponentLoader):
         # Config from Diffusers supersedes fastvideo's model config
         dit_config = fastvideo_args.pipeline_config.dit_config
         dit_config.update_model_arch(config)
-        # dit_config.update_model_arch({
-        #     "num_attention_heads": 12,
-        #     "ffn_dim": 8960,
-        #     "in_channels": 36,
-        #     "num_layers": 30,
-        #     "image_dim": 1280,
-        #     "added_kv_proj_dim": 0
-        # })
 
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         # Find all safetensors files
-        # model_path="/mnt/weka/home/hao.zhang/wei/FastVideo"
         safetensors_list = glob.glob(
             os.path.join(str(model_path), "*.safetensors"))
         if not safetensors_list:
@@ -467,8 +458,6 @@ class SchedulerLoader(ComponentLoader):
         assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
 
         scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-        from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler
-        scheduler_cls = FlowMatchDiscreteScheduler
 
         scheduler = scheduler_cls(**config)
         if fastvideo_args.pipeline_config.flow_shift is not None:

fastvideo/v1/tests/encoders/test_clip_encoder.py

@@ -168,5 +168,5 @@ def test_clip_encoder():
                 f"Pooler outputs differ significantly: mean diff = {mean_diff_pooler.item()}"
             assert max_diff_hidden < 1e-1, \
                 f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}"
-            assert max_diff_pooler < 1e-2, \
+            assert max_diff_pooler < 2e-2, \
                 f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}"

fastvideo/v1/tests/inference/STA/test_STA_inference.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import pytest
 
 NUM_NODES = "1"
-NUM_GPUS_PER_NODE = "1"
+NUM_GPUS_PER_NODE = "2"
 
 # Set environment variables
 os.environ["FASTVIDEO_ATTENTION_CONFIG"] = "assets/mask_strategy_wan.json"
@@ -17,9 +17,9 @@ def test_inference():
     cmd = [
         "fastvideo", "generate",
         "--model-path", "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-        "--sp-size", "1",
-        "--tp-size", "1",
-        "--num-gpus", "1",
+        "--sp-size", "2",
+        "--tp-size", "2",
+        "--num-gpus", "2",
         "--height", "768",
         "--width", "1280",
         "--num-frames", "69",

fastvideo/v1/tests/modal/pr_test.py

@@ -4,95 +4,96 @@ app = modal.App()
 
 import os
 
-image_version = os.getenv("IMAGE_VERSION", "latest")
+image_version = os.getenv("IMAGE_VERSION")
 image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
 print(f"Using image: {image_tag}")
 
 image = (
     modal.Image.from_registry(image_tag, add_python="3.12")
+    .run_commands("rm -rf /FastVideo")
     .apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev")
     .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
     .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
-    .env({"PATH": "/root/.cargo/bin:$PATH"})
-    .run_commands("/bin/bash -c 'source $HOME/.local/bin/env && source /opt/venv/bin/activate && cd /FastVideo && uv pip install -e .[test]'")
+    .env({
+        "PATH": "/root/.cargo/bin:$PATH",
+        "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
+        "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
+        "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
+        "IMAGE_VERSION": os.environ.get("IMAGE_VERSION", ""),
+    })
 )
 
-@app.function(gpu="L40S:1", image=image, timeout=1800)
+def run_test(pytest_command: str):
+    """Helper function to run a test suite with custom pytest command"""
+    import subprocess
+    import sys
+    import os
+    
+    git_repo = os.environ.get("BUILDKITE_REPO")
+    git_commit = os.environ.get("BUILDKITE_COMMIT")
+    pr_number = os.environ.get("BUILDKITE_PULL_REQUEST")
+    
+    print(f"Cloning repository: {git_repo}")
+    print(f"Target commit: {git_commit}")
+    if pr_number:
+        print(f"PR number: {pr_number}")
+    
+    # For PRs (including forks), use GitHub's PR refs to get the correct commit
+    if pr_number and pr_number != "false":
+        checkout_command = f"git fetch --prune origin refs/pull/{pr_number}/head && git checkout FETCH_HEAD"
+        print(f"Using PR ref for checkout: {checkout_command}")
+    else:
+        checkout_command = f"git checkout {git_commit}"
+        print(f"Using direct commit checkout: {checkout_command}")
+    
+    command = f"""
+    source $HOME/.local/bin/env &&
+    source /opt/venv/bin/activate &&
+    git clone {git_repo} /FastVideo &&
+    cd /FastVideo &&
+    {checkout_command} &&
+    uv pip install -e .[test] &&
+    {pytest_command}
+    """
+    
+    result = subprocess.run([
+        "/bin/bash", "-c", command
+    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
+    
+    sys.exit(result.returncode)
+
+@app.function(gpu="L40S:1", image=image, timeout=900)
 def run_encoder_tests():
-    """Run encoder tests on L40S GPU"""
-    import subprocess
-    import sys
-    import os
-    
-    os.chdir("/FastVideo")
-    
-    command = """
-    source /opt/venv/bin/activate && 
-    pytest ./fastvideo/v1/tests/encoders -s
-    """
-    
-    result = subprocess.run([
-        "/bin/bash", "-c", command
-    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
-    
-    sys.exit(result.returncode)
+    run_test("pytest ./fastvideo/v1/tests/encoders -vs")
 
-@app.function(gpu="L40S:1", image=image, timeout=1800)
+@app.function(gpu="L40S:1", image=image, timeout=900)
 def run_vae_tests():
-    """Run VAE tests on L40S GPU"""
-    import subprocess
-    import sys
-    import os
-    
-    os.chdir("/FastVideo")
-    
-    command = """
-    source /opt/venv/bin/activate && 
-    pytest ./fastvideo/v1/tests/vaes -s
-    """
-    
-    result = subprocess.run([
-        "/bin/bash", "-c", command
-    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
-    
-    sys.exit(result.returncode)
+    run_test("pytest ./fastvideo/v1/tests/vaes -vs")
 
-@app.function(gpu="L40S:1", image=image, timeout=1800)
+@app.function(gpu="L40S:1", image=image, timeout=900)
 def run_transformer_tests():
-    """Run transformer tests on L40S GPU"""
-    import subprocess
-    import sys
-    import os
-    
-    os.chdir("/FastVideo")
-    
-    command = """
-    source /opt/venv/bin/activate && 
-    pytest ./fastvideo/v1/tests/transformers -s
-    """
-    
-    result = subprocess.run([
-        "/bin/bash", "-c", command
-    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
-    
-    sys.exit(result.returncode)
+    run_test("pytest ./fastvideo/v1/tests/transformers -vs")
 
-@app.function(gpu="L40S:2", image=image, timeout=3600)
+@app.function(gpu="L40S:2", image=image, timeout=1800)
 def run_ssim_tests():
-    """Run SSIM tests on 2x L40S GPUs"""
-    import subprocess
-    import sys
-    import os
-    
-    os.chdir("/FastVideo")
-    
-    command = """
-    source /opt/venv/bin/activate && 
-    pytest ./fastvideo/v1/tests/ssim -vs
-    """
-    
-    result = subprocess.run([
-        "/bin/bash", "-c", command
-    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
-    
-    sys.exit(result.returncode)
+    run_test("pytest ./fastvideo/v1/tests/ssim -vs")
+
+@app.function(gpu="L40S:4", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
+def run_training_tests():
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/Vanilla -srP")
+
+@app.function(gpu="H100:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
+def run_training_tests_VSA():
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/VSA -srP")
+
+@app.function(gpu="H100:2", image=image, timeout=900)
+def run_inference_tests_STA():
+    run_test("pytest ./fastvideo/v1/tests/inference/STA -srP")
+
+@app.function(gpu="H100:1", image=image, timeout=900)
+def run_precision_tests_STA():
+    run_test("python csrc/attn/tests/test_sta.py")
+
+@app.function(gpu="H100:1", image=image, timeout=900)
+def run_precision_tests_VSA():
+    run_test("python csrc/attn/tests/test_block_sparse.py")

fastvideo/v1/tests/ssim/README.md

@@ -8,7 +8,7 @@ all reference videos are were generated on commit `4aeabbc629e0edf91477e80e795e7
 
 ## Generation Details
 
-2 x NVIDIA A40 GPUs
+2 x NVIDIA L40S GPUs
 
 ## Generation Parameters
 

fastvideo/v1/tests/ssim/test_inference_similarity.py

@@ -337,5 +337,5 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     if not success:
         logger.error("Failed to write SSIM results to file")
 
-    min_acceptable_ssim = 0.94
+    min_acceptable_ssim = 0.95
     assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"

fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json

@@ -1 +1 @@
-{"grad_norm":0.478515625,"_runtime":95.727033597,"_wandb":{"runtime":95},"_step":5,"validation_videos_50_steps":{"videos":[{"_type":"video-file","sha256":"42a1c311521a9d460db788713be1cbf2db767494e02619b43be5bf3eed8381d8","size":158632,"path":"media/videos/validation_videos_50_steps_0_42a1c311521a9d460db7.mp4"},{"path":"media/videos/validation_videos_50_steps_0_818505095b4b5e8b7f51.mp4","_type":"video-file","sha256":"818505095b4b5e8b7f511012d45f04d151ce3344bc058fc0f3225a414a851e4a","size":147825},{"sha256":"fc334ba9ed5e66c8527ee3b408e3be2d76167fef03588bf2840f4a0792f2fe34","size":136933,"path":"media/videos/validation_videos_50_steps_0_fc334ba9ed5e66c8527e.mp4","_type":"video-file"},{"size":201797,"path":"media/videos/validation_videos_50_steps_0_ccd98f6f907635d266a7.mp4","_type":"video-file","sha256":"ccd98f6f907635d266a74783688e7ecf1dac752d79d72d69eab9ef0e3f7413eb"},{"_type":"video-file","sha256":"ca79f40a0aed38f676f12779b349ce40e9e3fb7f36c578f49a20854c70508fb4","size":147114,"path":"media/videos/validation_videos_50_steps_0_ca79f40a0aed38f676f1.mp4"},{"size":175104,"path":"media/videos/validation_videos_50_steps_0_32c9b33ff920c17e5881.mp4","_type":"video-file","sha256":"32c9b33ff920c17e588133d7a27aa400ff3dc529b01ed4f16ac4d6bb2afa0f00"},{"sha256":"2cf520bfb93401c914e93c87ef791c2f12a4e043b95dfdc98115c930e11dfe67","size":139655,"path":"media/videos/validation_videos_50_steps_0_2cf520bfb93401c914e9.mp4","_type":"video-file"},{"_type":"video-file","sha256":"1d73aba17ce582c7aef4af4d64079e3e9d3df205634eff453446bdaf2340b214","size":149028,"path":"media/videos/validation_videos_50_steps_0_1d73aba17ce582c7aef4.mp4"}],"captions":false,"_type":"videos","count":8},"train_loss":0.08922439813613892,"_timestamp":1.750202051751466e+09,"avg_step_time":0.7536672964692116,"step_time":0.4742048177868128,"learning_rate":1e-05,"vsa_sparsity":0.05}
+{"step_time":0.6983645600266755,"_wandb":{"runtime":107},"grad_norm":0.50390625,"avg_step_time":1.002151239803061,"_step":5,"validation_videos_50_steps":{"captions":false,"_type":"videos","count":8,"videos":[{"size":159131,"path":"media/videos/validation_videos_50_steps_0_dc447599dbe48350e9c9.mp4","_type":"video-file","sha256":"dc447599dbe48350e9c920f4971e1786bde580dd20d52b4aa147ae8d3dc564d6"},{"_type":"video-file","sha256":"4e283876ddfbf5a2cb6f8aca07a39f832b5f806fbefde8854bc73ed904ff20ee","size":160315,"path":"media/videos/validation_videos_50_steps_0_4e283876ddfbf5a2cb6f.mp4"},{"size":135225,"path":"media/videos/validation_videos_50_steps_0_78185c41e1935306e93c.mp4","_type":"video-file","sha256":"78185c41e1935306e93c2d416ee40b31d038abffb029cb5bfb11c2a634eb2fcf"},{"_type":"video-file","sha256":"27e9819d002d3f63c8918bbdc5bf2857b5effe0caf5d5b8374b9c590fc6432eb","size":197873,"path":"media/videos/validation_videos_50_steps_0_27e9819d002d3f63c891.mp4"},{"_type":"video-file","sha256":"46fe548e86144ca60a9396fcd15e8788d3a05c066c6819ccdd6a9041feaaec8f","size":170601,"path":"media/videos/validation_videos_50_steps_0_46fe548e86144ca60a93.mp4"},{"sha256":"91ec338774bec870b9c5c81be4330a8f0f2535124cb372e881c3703e6b65ed77","size":164462,"path":"media/videos/validation_videos_50_steps_0_91ec338774bec870b9c5.mp4","_type":"video-file"},{"_type":"video-file","sha256":"ee4e811080a619215fd7541b39203dfb69c2db4cdadbfdd7a234a2cff684f6f3","size":139435,"path":"media/videos/validation_videos_50_steps_0_ee4e811080a619215fd7.mp4"},{"sha256":"22e31e048ba5e5b9d6587d7306904d6edc6c618b8f77c3ceb05ba2d602309274","size":147072,"path":"media/videos/validation_videos_50_steps_0_22e31e048ba5e5b9d658.mp4","_type":"video-file"}]},"_timestamp":1.75118195270901e+09,"vsa_sparsity":0.05,"learning_rate":1e-05,"train_loss":0.19960195198655128,"_runtime":107.325113071}

fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py

@@ -15,7 +15,7 @@ wandb_name = "test_training_loss_VSA"
 reference_wandb_summary_file = "fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json"
 
 NUM_NODES = "1"
-NUM_GPUS_PER_NODE = "1"
+NUM_GPUS_PER_NODE = "2"
 
 os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
 
@@ -35,14 +35,14 @@ def run_worker():
         "--validation_preprocessed_path", "data/mini_dataset_i2v_VSA/validation_parquet_dataset",
         "--train_batch_size", "1",
         "--num_latent_t", "4",
-        "--num_gpus", "1",
-        "--sp_size", "1",
-        "--tp_size", "1",
+        "--num_gpus", "2",
+        "--sp_size", "2",
+        "--tp_size", "2",
         "--hsdp_replicate_dim", "1",
-        "--hsdp_shard_dim", "1",
+        "--hsdp_shard_dim", "2",
         "--train_sp_batch_size", "1",
         "--dataloader_num_workers", "4",
-        "--gradient_accumulation_steps", "1",
+        "--gradient_accumulation_steps", "2",
         "--max_train_steps", "5",
         "--learning_rate", "1e-5",
         "--mixed_precision", "bf16",
@@ -110,7 +110,7 @@ def test_distributed_training():
     fields_and_thresholds = {
         'avg_step_time': 1.0,
         'grad_norm': 0.1,
-        'step_time': 0.5,
+        'step_time': 1.0,
         'train_loss': 0.001
     }
 

fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json

@@ -0,0 +1 @@
+{"step_time":5.501357046999999,"grad_norm":0.384765625,"train_loss":0.07890288904309273,"avg_step_time":5.831571423200001}

fastvideo/v1/tests/training/Vanilla/test_training_loss.py

@@ -18,7 +18,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.utils import FlexibleArgumentParser
 
 wandb_name = "test_training_loss"
-reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json"
+a40_reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/a40_reference_wandb_summary.json"
+l40s_reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json"
 
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "4"
@@ -108,13 +109,21 @@ def test_distributed_training():
 
     summary_file = 'wandb/latest-run/files/wandb-summary.json'
 
+    device_name = torch.cuda.get_device_name()
+    if "A40" in device_name:
+        reference_wandb_summary_file = a40_reference_wandb_summary_file
+    elif "L40S" in device_name:
+        reference_wandb_summary_file = l40s_reference_wandb_summary_file
+    else:
+        raise ValueError(f"Unknown device: {device_name}")
+
     reference_wandb_summary = json.load(open(reference_wandb_summary_file))
     wandb_summary = json.load(open(summary_file))
 
     fields_and_thresholds = {
-        'avg_step_time': 1.0,
+        'avg_step_time': 6.0,
         'grad_norm': 0.3,
-        'step_time': 0.5,
+        'step_time': 6.0,
         'train_loss': 0.0025
     }
 

fastvideo/v1/tests/transformers/test_hunyuanvideo.py

@@ -92,7 +92,7 @@ def test_hunyuanvideo_distributed():
 
     # Move to GPU based on local rank (0 or 1 for 2 GPUs)
     device = torch.device(f"cuda:0")
-    model = model.to(device)
+    model = model
 
     batch_size = 1
     seq_len = 3

fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py

@@ -62,7 +62,7 @@ def test_hunyuanvideo_distributed():
 
     precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         use_cpu_offload=False,
+                         use_cpu_offload=True,
                          pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str))
     args.device = torch.device(f"cuda:{LOCAL_RANK}")
 

fastvideo/v1/tests/transformers/test_wanvideo.py

@@ -34,12 +34,12 @@ def test_wan_transformer():
     precision = torch.bfloat16
     precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         use_cpu_offload=False,
+                         use_cpu_offload=True,
                          pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
     args.device = device
 
     loader = TransformerLoader()
-    model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision)
+    model2 = loader.load(TRANSFORMER_PATH, "", args).to(dtype=precision)
 
     model1 = WanTransformer3DModel.from_pretrained(
         TRANSFORMER_PATH, device=device,

fastvideo/v1/tests/vaes/test_hunyuan_vae.py

@@ -30,7 +30,10 @@ CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
 
 # Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
 # torch 2.7.1
-REFERENCE_LATENT = -106.22467041015625
+A40_REFERENCE_LATENT = -106.22467041015625
+
+# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
+L40S_REFERENCE_LATENT = -158.32318115234375
 
 
 @pytest.mark.usefixtures("distributed_setup")
@@ -67,6 +70,14 @@ def test_hunyuan_vae():
         latent = model.encode(input_tensor).mean.double().sum().item()
 
     # Check if latents are similar
+    device_name = torch.cuda.get_device_name()
+    if "A40" in device_name:
+        REFERENCE_LATENT = A40_REFERENCE_LATENT
+    elif "L40S" in device_name:
+        REFERENCE_LATENT = L40S_REFERENCE_LATENT
+    else:
+        raise ValueError(f"Unknown device: {device_name}")
+
     diff_encoded_latents = abs(REFERENCE_LATENT - latent)
     logger.info(
         f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}"

fastvideo/v1/training/activation_checkpoint.py

@@ -0,0 +1,91 @@
+import collections
+from enum import Enum
+from typing import Optional
+
+import torch
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    checkpoint_wrapper)
+
+TRANSFORMER_BLOCK_NAMES = [
+    "blocks",
+    "double_blocks",
+    "single_blocks",
+    "transformer_blocks",
+    "temporal_transformer_blocks",
+    "transformer_double_blocks",
+    "transformer_single_blocks",
+]
+
+
+class CheckpointType(str, Enum):
+    FULL = "full"
+    OPS = "ops"
+    BLOCK_SKIP = "block_skip"
+
+
+_SELECTIVE_ACTIVATION_CHECKPOINTING_OPS = {
+    torch.ops.aten.mm.default,
+    torch.ops.aten._scaled_dot_product_efficient_attention.default,
+    torch.ops.aten._scaled_dot_product_flash_attention.default,
+    torch.ops._c10d_functional.reduce_scatter_tensor.default,
+}
+
+
+def apply_activation_checkpointing(
+        module: torch.nn.Module,
+        checkpointing_type: str = CheckpointType.FULL,
+        n_layer: int = 1) -> torch.nn.Module:
+    if checkpointing_type == CheckpointType.FULL:
+        module = _apply_activation_checkpointing_blocks(module)
+    elif checkpointing_type == CheckpointType.OPS:
+        module = _apply_activation_checkpointing_ops(
+            module, _SELECTIVE_ACTIVATION_CHECKPOINTING_OPS)
+    elif checkpointing_type == CheckpointType.BLOCK_SKIP:
+        module = _apply_activation_checkpointing_blocks(module, n_layer)
+    else:
+        raise ValueError(
+            f"Checkpointing type '{checkpointing_type}' not supported. Supported types are {CheckpointType.__members__.keys()}"
+        )
+    return module
+
+
+def _apply_activation_checkpointing_blocks(
+        module: torch.nn.Module,
+        n_layer: Optional[int] = None) -> torch.nn.Module:
+    for transformer_block_name in TRANSFORMER_BLOCK_NAMES:
+        blocks: torch.nn.Module = getattr(module, transformer_block_name, None)
+        if blocks is None:
+            continue
+        for index, (layer_id, block) in enumerate(blocks.named_children()):
+            if n_layer is None or index % n_layer == 0:
+                block = checkpoint_wrapper(block, preserve_rng_state=False)
+                blocks.register_module(layer_id, block)
+    return module
+
+
+def _apply_activation_checkpointing_ops(module: torch.nn.Module,
+                                        ops) -> torch.nn.Module:
+    from torch.utils.checkpoint import (CheckpointPolicy,
+                                        create_selective_checkpoint_contexts)
+
+    def _get_custom_policy(meta: dict[str, int]) -> CheckpointPolicy:
+
+        def _custom_policy(ctx, func, *args, **kwargs):
+            mode = "recompute" if ctx.is_recompute else "forward"
+            mm_count_key = f"{mode}_mm_count"
+            if func == torch.ops.aten.mm.default:
+                meta[mm_count_key] += 1
+            # Saves output of all compute ops, except every second mm
+            to_save = func in ops and not (func == torch.ops.aten.mm.default
+                                           and meta[mm_count_key] % 2 == 0)
+            return CheckpointPolicy.MUST_SAVE if to_save else CheckpointPolicy.PREFER_RECOMPUTE
+
+        return _custom_policy
+
+    def selective_checkpointing_context_fn():
+        meta: dict[str, int] = collections.defaultdict(int)
+        return create_selective_checkpoint_contexts(_get_custom_policy(meta))
+
+    return checkpoint_wrapper(module,
+                              context_fn=selective_checkpointing_context_fn,
+                              preserve_rng_state=False)

fastvideo/v1/training/training_pipeline.py

@@ -31,6 +31,8 @@ from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
                                     TrainingBatch)
+from fastvideo.v1.training.activation_checkpoint import (
+    apply_activation_checkpointing)
 from fastvideo.v1.training.training_utils import (
     clip_grad_norm_while_handling_failing_dtensor_cases,
     compute_density_for_timestep_sampling, get_sigmas, load_checkpoint,
@@ -82,6 +84,11 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
 
         self.transformer.requires_grad_(True)
         self.transformer.train()
+        if training_args.enable_gradient_checkpointing_type is not None:
+            self.transformer = apply_activation_checkpointing(
+                self.transformer,
+                checkpointing_type=training_args.
+                enable_gradient_checkpointing_type)
 
         noise_scheduler = self.modules["scheduler"]
         params_to_optimize = self.transformer.parameters()
@@ -245,9 +252,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         current_vsa_sparsity = training_batch.current_vsa_sparsity
 
         if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
-
             dit_seq_shape = [
-                latents.shape[2] // patch_size[0],
+                latents.shape[2] * self.sp_world_size // patch_size[0],
                 latents.shape[3] // patch_size[1],
                 latents.shape[4] // patch_size[2]
             ]
@@ -310,17 +316,18 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                 current_timestep=training_batch.current_timestep,
                 attn_metadata=training_batch.attn_metadata):
             model_pred = self.transformer(**input_kwargs)
-        if self.training_args.precondition_outputs:
-            model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
-        target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+            if self.training_args.precondition_outputs:
+                model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
+            target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
 
-        # make sure no implicit broadcasting happens
-        assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
-        loss = (torch.mean((model_pred.float() - target.float())**2) /
-                self.training_args.gradient_accumulation_steps)
+            # make sure no implicit broadcasting happens
+            assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
+            loss = (torch.mean((model_pred.float() - target.float())**2) /
+                    self.training_args.gradient_accumulation_steps)
+
+            loss.backward()
+            avg_loss = loss.detach().clone()
 
-        loss.backward()
-        avg_loss = loss.detach().clone()
         # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
         #             local_main_process_only=False)
         world_group = get_world_group()
@@ -547,6 +554,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             negative_prompt_attention_mask: torch.Tensor | None
     ) -> ForwardBatch:
 
+        assert len(validation_batch['info_list']
+                   ) == 1, "Only batch size 1 is supported for validation"
         prompt = validation_batch['info_list'][0]['prompt']
         prompt_embeds = validation_batch['text_embedding']
         prompt_attention_mask = validation_batch['text_attention_mask']
@@ -630,7 +639,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         # Process each validation prompt for each validation step
         for num_inference_steps in validation_steps:
             step_videos: List[np.ndarray] = []
-            step_captions: List[str | None] = []
+            step_captions: List[str] = []
 
             for validation_batch in validation_dataloader:
                 batch = self._prepare_validation_inputs(
@@ -638,7 +647,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                     num_inference_steps, negative_prompt_embeds,
                     negative_prompt_attention_mask)
 
-                step_captions.extend([None])  # TODO(peiyuan): add caption
+                assert batch.prompt is not None and isinstance(
+                    batch.prompt, str)
+                step_captions.append(batch.prompt)
 
                 # Run validation inference
                 with torch.no_grad(), torch.autocast("cuda",

fastvideo/v1/training/training_utils.py

@@ -162,8 +162,8 @@ def save_checkpoint(transformer,
                     weight_path,
                     local_main_process_only=False)
 
-        # Convert training format to diffusers format and save
-        diffusers_state_dict = convert_training_to_diffusers_format(
+        # Convert fastvideo custom format to diffusers format and save
+        diffusers_state_dict = convert_custom_format_to_diffusers_format(
             cpu_state, transformer)
         save_file(diffusers_state_dict, weight_path)
 
@@ -487,10 +487,10 @@ def _has_foreach_support(tensors: List[torch.Tensor],
         t is None or type(t) in [torch.Tensor] for t in tensors)
 
 
-def convert_training_to_diffusers_format(state_dict: Dict[str, Any],
-                                         transformer) -> Dict[str, Any]:
+def convert_custom_format_to_diffusers_format(state_dict: Dict[str, Any],
+                                              transformer) -> Dict[str, Any]:
     """
-    Convert training format state dict to diffusers format using reverse_param_names_mapping.
+    Convert fastvideo custom format state dict to diffusers format using reverse_param_names_mapping.
     
     Args:
         state_dict: State dict in training format

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.1.0"
+version = "0.1.1"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.8"

scripts/finetune/finetune_v1.sh

@@ -48,4 +48,5 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --weight_decay 0.01 \
     --not_apply_cfg_solver \
     --dit_precision "fp32" \
-    --max_grad_norm 1.0
+    --max_grad_norm 1.0 \
+    --enable_gradient_checkpointing_type "full"

scripts/finetune/finetune_v1_VSA.sh

@@ -58,5 +58,6 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --VSA_decay_sparsity 0.9 \
     --VSA_decay_rate 0.03 \
     --VSA_decay_interval_steps 30 \
-    --VSA_val_sparsity 0.9
+    --VSA_val_sparsity 0.9 \
+    --enable_gradient_checkpointing_type "full"
 # --resume_from_checkpoint "$CHECKPOINT_PATH"

scripts/huggingface/convert_to_hf.py

@@ -1,88 +0,0 @@
-from huggingface_hub import save_torch_state_dict, load_state_dict_from_file
-from safetensors.torch import save_file
-import torch
-import re
-from collections import OrderedDict
-
-_param_names_mapping: dict = {
-    r"^text_embedding\.0\.(.*)$":
-    r"condition_embedder.text_embedder.linear_1.\1",
-    r"^text_embedding\.2\.(.*)$":
-    r"condition_embedder.text_embedder.linear_2.\1",
-    r"^time_embedding\.0\.(.*)$":
-    r"condition_embedder.time_embedder.linear_1.\1",
-    r"^time_embedding\.2\.(.*)$":
-    r"condition_embedder.time_embedder.linear_2.\1",
-    r"^time_projection\.1\.(.*)$":
-    r"condition_embedder.time_proj.\1",
-    r"^img_emb\.proj\.0\.(.*)$":
-    r"condition_embedder.image_embedder.norm1.\1",
-    r"^img_emb\.proj\.1\.(.*)$":
-    r"condition_embedder.image_embedder.ff.net.0.proj.\1",
-    r"^img_emb\.proj\.3\.(.*)$":
-    r"condition_embedder.image_embedder.ff.net.2.\1",
-    r"^img_emb\.proj\.4\.(.*)$":
-    r"condition_embedder.image_embedder.norm2.\1",
-    r"^head\.modulation":
-    r"scale_shift_table",
-    r"^head\.head\.(.*)$":
-    r"proj_out.\1",
-    r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
-    r"blocks.\1.attn1.to_q.\2",
-    r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
-    r"blocks.\1.attn1.to_k.\2",
-    r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
-    r"blocks.\1.attn1.to_v.\2",
-    r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
-    r"blocks.\1.attn1.to_out.0.\2",
-    r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
-    r"blocks.\1.attn1.norm_q.\2",
-    r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
-    r"blocks.\1.attn1.norm_k.\2",
-    r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
-    r"blocks.\1.attn2.to_q.\2",
-    r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
-    r"blocks.\1.attn2.to_k.\2",
-    r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
-    r"blocks.\1.attn2.add_k_proj.\2",
-    r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
-    r"blocks.\1.attn2.to_v.\2",
-    r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
-    r"blocks.\1.attn2.add_v_proj.\2",
-    r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
-    r"blocks.\1.attn2.to_out.0.\2",
-    r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
-    r"blocks.\1.attn2.norm_q.\2",
-    r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
-    r"blocks.\1.attn2.norm_k.\2",
-    r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
-    r"blocks.\1.attn2.norm_added_k.\2",
-    r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
-    r"blocks.\1.ffn.net.0.proj.\2",
-    r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
-    r"blocks.\1.ffn.net.2.\2",
-    r"^blocks\.(\d+)\.modulation":
-    r"blocks.\1.scale_shift_table",
-    r"^blocks\.(\d+)\.norm3\.(.*)$":
-    r"blocks.\1.norm2.\2",
-}
-
-state_dict = load_state_dict_from_file("/mnt/weka/home/hao.zhang/wei/FastVideo/diffusion_pytorch_model.safetensors")
-new_state_dict = OrderedDict()
-for k, v in state_dict.items():
-    new_key = k
-    for pattern, replacement in _param_names_mapping.items():
-        if re.match(pattern, k):
-            new_key = re.sub(pattern, replacement, k)
-            break  # Stop at the first match
-    new_state_dict[new_key] = v
-    if "norm_added_k" in new_key:
-        dummy_key = new_key.replace("norm_added_k", "norm_added_q")
-        dummy_value = torch.zeros_like(v)
-        new_state_dict[dummy_key] = dummy_value
-del state_dict
-
-save_torch_state_dict(
-    new_state_dict,
-    "/mnt/weka/home/hao.zhang/wei/FastVideo/data/Wan2.1-Fun-1.3B-InP-Diffusers/transformer"
-)

scripts/huggingface/upload_hf.py

@@ -3,7 +3,7 @@ from huggingface_hub import HfApi
 api = HfApi()
 
 api.upload_folder(
-    folder_path="/mnt/weka/home/hao.zhang/wei/FastVideo/data/Wan2.1-Fun-1.3B-InP-Diffusers",
-    repo_id="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers",
-    repo_type="model",
+    folder_path="data/Black-Myth-Taylor-Src",
+    repo_id="FastVideo/Image-Vid-Finetune-Src",
+    repo_type="dataset",
 )

scripts/inference/v1_inference_wan_VSA.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+num_gpus=1
+export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
+# change model path to local dir if you want to inference using your checkpoint
+export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size $num_gpus \
+    --num-gpus $num_gpus \
+    --height 448 \
+    --width 832 \
+    --num-frames 77 \
+    --num-inference-steps 50 \
+    --fps 16 \
+    --guidance-scale 6.0 \
+    --flow-shift 8.0 \
+    --VSA-sparsity 0.9 \
+    --prompt "A beautiful woman in a red dress walking down a street" \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
+    --seed 1024 \
+    --output-path outputs_video_1.3B_VSA/sparsity_0.9/

scripts/inference/v1_inference_wan_i2v.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
-export HOME=/mnt/weka/home/hao.zhang/wei
+
 num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
-export MODEL_BASE=$HOME/FastVideo/data/Wan2.1-Fun-1.3B-InP-Diffusers
+export MODEL_BASE=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 # Note that the tp_size and sp_size should be the same and equal to the number
 # of GPUs. They are used for different parallel groups. sp_size is used for
@@ -14,13 +14,13 @@ fastvideo generate \
     --num-gpus $num_gpus \
     --height 480 \
     --width 832 \
-    --num-frames 29 \
+    --num-frames 77 \
     --num-inference-steps 40 \
     --fps 16 \
     --flow-shift 3.0 \
     --guidance-scale 5.0 \
-    --image-path "assets/dog.png" \
-    --prompt "A brown dog is shaking its head while sitting on a light-colored sofa in a cozy room. Behind the dog, there is a framed painting on the shelf, surrounded by pink flowers. The soft, warm lighting in the room creates a comfortable atmosphere." \
+    --image-path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" \
+    --prompt "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." \
     --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
     --seed 1024 \
-    --output-path outputs_videox_i2v/
+    --output-path outputs_i2v/