fix

README.md

@@ -8,7 +8,7 @@ It features a clean, consistent API that works across popular video models, maki
 With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>
 
 <div align="center">

docs/source/contributing/developer_env/runpod.md

@@ -8,14 +8,12 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](
 
 Choose a GPU that supports CUDA 12.4
 
-Pick 1 or 2 L40S GPU(s)
-
 ![RunPod CUDA selection](../../_static/images/runpod_cuda.png)
 
 When creating your pod template, use this image:
 
 ```
-ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
+ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
 ```
 
 Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):

docs/source/getting_started/installation.md

@@ -117,4 +117,4 @@ If you're planning to contribute to FastVideo please see the following page:
 
 If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
 
-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg) for additional support.

docs/source/inference/add_pipeline.md

@@ -12,7 +12,7 @@ This guide explains how to implement a custom diffusion pipeline in FastVideo, l
 4. **Register Your Pipeline** - Make it discoverable by the framework
 5. **Configure Your Pipeline** - (Coming soon)
 
-Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
 
 ## Step 1: Pipeline Modules
 

docs/source/inference/cli.md

@@ -27,7 +27,7 @@ fastvideo generate --help
 ### Hardware Configuration
 
 - `--num-gpus {NUM_GPUS}`: Number of GPUs to use
-- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
 - `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
 
 #### Video Configuration
@@ -68,7 +68,7 @@ Example configuration file (config.json):
     "output_path": "outputs/",
     "num_gpus": 2,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "num_frames": 45,
     "height": 720,
     "width": 1280,
@@ -102,7 +102,7 @@ prompt: "A beautiful woman in a red dress walking down a street"
 output_path: "outputs/"
 num_gpus: 2
 sp_size: 2
-tp_size: 1
+tp_size: 2
 num_frames: 45
 height: 720
 width: 1280

docs/source/inference/inference_quick_start.md

@@ -121,4 +121,4 @@ If the generated video doesn't match your prompt:
 - Learn about using [Optimizations](#inference-optimizations)
 - See [Examples](../examples/examples_inference_index.md) for more usage scenarios
 - Join our [Community Discord](https://discord.gg/JA7cksDz86).
-- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).

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

@@ -30,7 +30,7 @@ training_args=(
 parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 8
-  --tp_size 1
+  --tp_size 8
   --hsdp_replicate_dim 1
   --hsdp_shard_dim 8
 )

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm

@@ -66,7 +66,7 @@ training_args=(
 parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size $NUM_GPUS
-  --tp_size 1
+  --tp_size $NUM_GPUS
   --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
   --hsdp_shard_dim $NUM_GPUS
 )

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

@@ -30,7 +30,7 @@ training_args=(
 parallel_args=(
   --num_gpus $NUM_GPUS 
   --sp_size $NUM_GPUS 
-  --tp_size 1
+  --tp_size $NUM_GPUS
   --hsdp_replicate_dim 1
   --hsdp_shard_dim $NUM_GPUS
 )

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm

@@ -63,7 +63,7 @@ training_args=(
 parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 4
-  --tp_size 1
+  --tp_size 4
   --hsdp_replicate_dim 2
   --hsdp_shard_dim 4
 )

fastvideo/v1/configs/sample/registry.py

@@ -6,8 +6,7 @@ from typing import Any
 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                  HunyuanSamplingParam)
 from fastvideo.v1.configs.sample.stepvideo import StepVideoT2VSamplingParam
-from fastvideo.v1.configs.sample.wan import (Wan2_1_Fun_1_3B_InP_SamplingParam,
-                                             WanI2V_14B_480P_SamplingParam,
+from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
                                              WanI2V_14B_720P_SamplingParam,
                                              WanT2V_1_3B_SamplingParam,
                                              WanT2V_14B_SamplingParam)
@@ -24,8 +23,6 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
     # Add other specific weight variants
 }

fastvideo/v1/configs/sample/wan.py

@@ -94,20 +94,3 @@ class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
                 -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
                 -13.02252404
             ]))
-
-
-# =============================================
-# ============= Wan2.1 Fun Models =============
-# =============================================
-
-
-@dataclass
-class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
-    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
-    height: int = 480
-    width: int = 832
-    num_frames: int = 81
-    fps: int = 16
-    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
-    guidance_scale: float = 6.0
-    num_inference_steps: int = 50

fastvideo/v1/dataset/parquet_dataset_iterable_style.py

@@ -5,13 +5,14 @@ import random
 import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
+import torch
 import tqdm
 from torch.utils.data import IterableDataset, get_worker_info
 from torchdata.stateful_dataloader import StatefulDataLoader
 
 from fastvideo.v1.dataset.utils import collate_latents_embs_masks
-from fastvideo.v1.distributed import (get_sp_world_size, get_world_group,
-                                      get_world_rank, get_world_size)
+from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
+                                      get_world_size)
 from fastvideo.v1.logger import init_logger
 
 logger = init_logger(__name__)
@@ -156,123 +157,95 @@ def shard_parquet_files_across_sp_groups_and_workers(
     # Check if sharding plan already exists
     sharding_info_dir = os.path.join(
         path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
-
-    # Only rank 0 handles cache checking and file scanning
-    if get_world_rank() == 0:
-        cache_loaded = False
-        shard_parquet_files = None
-        shard_total_samples = None
-        shard_parquet_lengths = None
-
-        # First try to load existing sharding plan
-        if os.path.exists(sharding_info_dir):
-            logger.info("Loading sharding plan from %s", sharding_info_dir)
-            try:
-                with open(
-                        os.path.join(sharding_info_dir,
-                                     "shard_parquet_files.pkl"), "rb") as f:
-                    shard_parquet_files = pickle.load(f)
-                with open(
-                        os.path.join(sharding_info_dir,
-                                     "shard_total_samples.pkl"), "rb") as f:
-                    shard_total_samples = pickle.load(f)
-                with open(
-                        os.path.join(sharding_info_dir,
-                                     "shard_parquet_lengths.pkl"), "rb") as f:
-                    shard_parquet_lengths = pickle.load(f)
-                cache_loaded = True
-                logger.info("Successfully loaded sharding plan")
-            except Exception as e:
-                logger.error("Error loading sharding plan: %s", str(e))
-                logger.info("Falling back to creating new sharding plan")
-                cache_loaded = False
-
-        # If cache not loaded (either doesn't exist or failed to load), create sharding plan
-        if not cache_loaded:
-            logger.info("Creating new sharding plan")
-            logger.info("Scanning for parquet files in %s", path)
-
-            # Find all parquet files
-            parquet_files = []
-
-            for root, _, files in os.walk(path):
-                for file in files:
-                    if file.endswith('.parquet'):
-                        parquet_files.append(os.path.join(root, file))
-
-            if not parquet_files:
-                raise ValueError("No parquet files found in %s", path)
-
-            # Calculate file lengths efficiently using a single pass
-            logger.info("Calculating file lengths...")
-            lengths = []
-            for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
-                lengths.append(pq.ParquetFile(file).metadata.num_rows)
-
-            total_samples = sum(lengths)
-            logger.info("Found %d files with %d total samples",
-                        len(parquet_files), total_samples)
-
-            # Sort files by length for better balancing
-            sorted_indices = np.argsort(lengths)
-            sorted_files = [parquet_files[i] for i in sorted_indices]
-            sorted_lengths = [lengths[i] for i in sorted_indices]
-
-            # Create shards
-            num_shards = num_sp_groups * num_workers
-            shard_parquet_files = [[] for _ in range(num_shards)]
-            shard_total_samples = [0] * num_shards
-            shard_parquet_lengths = [{} for _ in range(num_shards)]
-
-            # Distribute files to shards using a greedy approach
-            logger.info("Distributing files to shards...")
-            for file, length in zip(reversed(sorted_files),
-                                    reversed(sorted_lengths),
-                                    strict=True):
-                # Find shard with minimum current length
-                target_shard = np.argmin(shard_total_samples)
-                shard_parquet_files[target_shard].append(file)
-                shard_total_samples[target_shard] += length
-                shard_parquet_lengths[target_shard][file] = length
-            #randomize each shard
-            for shard in shard_parquet_files:
-                random.seed(seed)
-                random.shuffle(shard)
-
-            # Save the sharding plan
-            os.makedirs(sharding_info_dir, exist_ok=True)
+    if os.path.exists(sharding_info_dir):
+        logger.info("Sharding plan already exists")
+        logger.info("Loading sharding plan from %s", sharding_info_dir)
+        try:
             with open(
                     os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
-                    "wb") as f:
-                pickle.dump(shard_parquet_files, f)
+                    "rb") as f:
+                shard_parquet_files = pickle.load(f)
             with open(
                     os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
-                    "wb") as f:
-                pickle.dump(shard_total_samples, f)
+                    "rb") as f:
+                shard_total_samples = pickle.load(f)
             with open(
                     os.path.join(sharding_info_dir,
-                                 "shard_parquet_lengths.pkl"), "wb") as f:
-                pickle.dump(shard_parquet_lengths, f)
-            logger.info("Saved sharding info to %s", sharding_info_dir)
+                                 "shard_parquet_lengths.pkl"), "rb") as f:
+                shard_parquet_lengths = pickle.load(f)
+            return shard_parquet_files, shard_total_samples, shard_parquet_lengths
+        except Exception as e:
+            logger.error("Error loading sharding plan: %s", str(e))
+            logger.info("Falling back to creating new sharding plan")
 
-    # Wait for rank 0 to finish creating/loading sharding plan
-    world_group = get_world_group()
-    world_group.barrier()
+    if get_world_rank() == 0:
+        logger.info("Scanning for parquet files in %s", path)
 
-    # Now all ranks load the sharding plan (it should exist and be valid now)
-    logger.info("Loading sharding plan from %s after barrier",
-                sharding_info_dir)
-    with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
-              "rb") as f:
-        shard_parquet_files = pickle.load(f)
-    with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
-              "rb") as f:
-        shard_total_samples = pickle.load(f)
-    with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"),
-              "rb") as f:
-        shard_parquet_lengths = pickle.load(f)
+        # Find all parquet files
+        parquet_files = []
 
-    return shard_parquet_files, shard_total_samples, shard_parquet_lengths
+        for root, _, files in os.walk(path):
+            for file in files:
+                if file.endswith('.parquet'):
+                    parquet_files.append(os.path.join(root, file))
+
+        if not parquet_files:
+            raise ValueError("No parquet files found in %s", path)
+
+        # Calculate file lengths efficiently using a single pass
+        logger.info("Calculating file lengths...")
+        lengths = []
+        for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
+            lengths.append(pq.ParquetFile(file).metadata.num_rows)
+
+        total_samples = sum(lengths)
+        logger.info("Found %d files with %d total samples", len(parquet_files),
+                    total_samples)
+
+        # Sort files by length for better balancing
+        sorted_indices = np.argsort(lengths)
+        sorted_files = [parquet_files[i] for i in sorted_indices]
+        sorted_lengths = [lengths[i] for i in sorted_indices]
+
+        # Create shards
+        num_shards = num_sp_groups * num_workers
+        shard_parquet_files = [[] for _ in range(num_shards)]
+        shard_total_samples = [0] * num_shards
+        shard_parquet_lengths = [{} for _ in range(num_shards)]
+
+        # Distribute files to shards using a greedy approach
+        logger.info("Distributing files to shards...")
+        for file, length in zip(reversed(sorted_files),
+                                reversed(sorted_lengths),
+                                strict=True):
+            # Find shard with minimum current length
+            target_shard = np.argmin(shard_total_samples)
+            shard_parquet_files[target_shard].append(file)
+            shard_total_samples[target_shard] += length
+            shard_parquet_lengths[target_shard][file] = length
+        #randomize each shard
+        for shard in shard_parquet_files:
+            random.seed(seed)
+            random.shuffle(shard)
+
+        save_dir = os.path.join(
+            path,
+            f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
+        os.makedirs(save_dir, exist_ok=True)
+        with open(os.path.join(save_dir, "shard_parquet_files.pkl"), "wb") as f:
+            pickle.dump(shard_parquet_files, f)
+        with open(os.path.join(save_dir, "shard_total_samples.pkl"), "wb") as f:
+            pickle.dump(shard_total_samples, f)
+        with open(os.path.join(save_dir, "shard_parquet_lengths.pkl"),
+                  "wb") as f:
+            pickle.dump(shard_parquet_lengths, f)
+        logger.info("Saved sharding info to %s", save_dir)
+
+    # wait for all ranks to finish
+    torch.distributed.barrier()
+    # recursive call
+    return shard_parquet_files_across_sp_groups_and_workers(
+        path, num_sp_groups, num_workers, seed)
 
 
 def build_parquet_iterable_style_dataloader(

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -13,8 +13,8 @@ from torch.utils.data import Dataset, Sampler
 from torchdata.stateful_dataloader import StatefulDataLoader
 
 from fastvideo.v1.dataset.utils import collate_rows_from_parquet_schema
-from fastvideo.v1.distributed import (get_sp_world_size, get_world_group,
-                                      get_world_rank, get_world_size)
+from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
+                                      get_world_size)
 from fastvideo.v1.logger import init_logger
 
 logger = init_logger(__name__)
@@ -97,64 +97,48 @@ def get_parquet_files_and_length(path: str):
     cache_dir = os.path.join(path, "map_style_cache")
     cache_file = os.path.join(cache_dir, "file_info.pkl")
 
-    # Only rank 0 checks for cache and scans files if needed
+    if os.path.exists(cache_file):
+        logger.info("Loading cached file info from %s", cache_file)
+        try:
+            with open(cache_file, "rb") as f:
+                file_names_sorted, lengths_sorted = pickle.load(f)
+            return file_names_sorted, lengths_sorted
+        except Exception as e:
+            logger.error("Error loading cached file info: %s", str(e))
+            logger.info("Falling back to scanning files")
+
+    # If no cache exists or loading failed, scan files
     if get_world_rank() == 0:
-        cache_loaded = False
-        file_names_sorted = None
-        lengths_sorted = None
+        lengths = []
+        file_names = []
+        for root, _, files in os.walk(path):
+            for file in sorted(files):
+                if file.endswith('.parquet'):
+                    file_path = os.path.join(root, file)
+                    file_names.append(file_path)
+        for file_path in tqdm.tqdm(file_names,
+                                   desc="Reading parquet files to get lengths"):
+            num_rows = pq.ParquetFile(file_path).metadata.num_rows
+            lengths.append(num_rows)
+        # sort according to file name to ensure all rank has the same order (in case os.walk is not sorted)
+        file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names,
+                                                            lengths,
+                                                            strict=True),
+                                                        key=lambda x: x[0]),
+                                                strict=True)
+        assert len(
+            file_names_sorted) != 0, "No parquet files found in the dataset"
 
-        # First try to load existing cache
-        if os.path.exists(cache_file):
-            logger.info("Loading cached file info from %s", cache_file)
-            try:
-                with open(cache_file, "rb") as f:
-                    file_names_sorted, lengths_sorted = pickle.load(f)
-                cache_loaded = True
-                logger.info("Successfully loaded cached file info")
-            except Exception as e:
-                logger.error("Error loading cached file info: %s", str(e))
-                logger.info("Falling back to scanning files")
-                cache_loaded = False
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(cache_file, "wb") as f:
+            pickle.dump((file_names_sorted, lengths_sorted), f)
+        logger.info("Saved file info to %s", cache_file)
 
-        # If cache not loaded (either doesn't exist or failed to load), scan files
-        if not cache_loaded:
-            logger.info("Scanning parquet files to get lengths")
-            lengths = []
-            file_names = []
-            for root, _, files in os.walk(path):
-                for file in sorted(files):
-                    if file.endswith('.parquet'):
-                        file_path = os.path.join(root, file)
-                        file_names.append(file_path)
-            for file_path in tqdm.tqdm(
-                    file_names, desc="Reading parquet files to get lengths"):
-                num_rows = pq.ParquetFile(file_path).metadata.num_rows
-                lengths.append(num_rows)
-            # sort according to file name to ensure all rank has the same order
-            file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names,
-                                                                lengths,
-                                                                strict=True),
-                                                            key=lambda x: x[0]),
-                                                    strict=True)
-            assert len(
-                file_names_sorted) != 0, "No parquet files found in the dataset"
+    # Wait for rank 0 to finish saving
+    if get_world_size() > 1:
+        torch.distributed.barrier()
 
-            # Save the cache
-            os.makedirs(cache_dir, exist_ok=True)
-            with open(cache_file, "wb") as f:
-                pickle.dump((file_names_sorted, lengths_sorted), f)
-            logger.info("Saved file info to %s", cache_file)
-
-    # Wait for rank 0 to finish creating/loading cache
-    world_group = get_world_group()
-    world_group.barrier()
-
-    # Now all ranks load the cache (it should exist and be valid now)
-    logger.info("Loading cached file info from %s after barrier", cache_file)
-    with open(cache_file, "rb") as f:
-        file_names_sorted, lengths_sorted = pickle.load(f)
-
-    return file_names_sorted, lengths_sorted
+    return get_parquet_files_and_length(path)
 
 
 def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
@@ -169,27 +153,17 @@ def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
     '''
     # find the parquet file and local row index
     cumulative = 0
-    file_index = 0
-    local_row_idx = 0
-
     for file_index in range(len(lengths)):
         if cumulative + lengths[file_index] > global_row_idx:
             local_row_idx = global_row_idx - cumulative
             break
         cumulative += lengths[file_index]
-    else:
-        # If we reach here, global_row_idx is out of bounds
-        raise IndexError(
-            f"global_row_idx {global_row_idx} is out of bounds for dataset")
 
     parquet_file = pq.ParquetFile(parquet_files[file_index])
 
     # Calculate the row group to read into memory and the local idx
     # This way we can avoid reading in the entire parquet file
     cumulative = 0
-    row_group_index = 0
-    local_index = 0
-
     for i in range(parquet_file.num_row_groups):
         num_rows = parquet_file.metadata.row_group(i).num_rows
         if cumulative + num_rows > local_row_idx:
@@ -197,11 +171,6 @@ def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
             local_index = local_row_idx - cumulative
             break
         cumulative += num_rows
-    else:
-        # If we reach here, local_row_idx is out of bounds for this parquet file
-        raise IndexError(
-            f"local_row_idx {local_row_idx} is out of bounds for parquet file {parquet_files[file_index]}"
-        )
 
     row_group = parquet_file.read_row_group(row_group_index).to_pydict()
     row_dict = {k: v[local_index] for k, v in row_group.items()}

fastvideo/v1/dataset/utils.py

@@ -140,7 +140,6 @@ def collate_rows_from_parquet_schema(rows,
             if shape_key in row and bytes_key in row:
                 shape = row[shape_key]
                 bytes_data = row[bytes_key]
-
                 if len(bytes_data) == 0:
                     tensor = torch.zeros(0, dtype=torch.bfloat16)
                 else:

fastvideo/v1/distributed/parallel_state.py

@@ -23,6 +23,7 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
+import gc
 import os
 import pickle
 import weakref
@@ -805,7 +806,6 @@ def get_dp_group() -> GroupCoordinator:
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     sequence_model_parallel_size: int = 1,
-    data_parallel_size: int = 1,
     backend: str | None = None,
 ) -> None:
     """
@@ -1015,6 +1015,15 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
+    gc.collect()
+    from fastvideo.v1.platforms import current_platform
+    if not current_platform.is_cpu():
+        torch.cuda.empty_cache()
+    try:
+        torch._C._host_emptyCache()
+    except AttributeError:
+        logger.warning(
+            "torch._C._host_emptyCache() only available in Pytorch >=2.5")
 
 
 def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,

fastvideo/v1/entrypoints/video_generator.py

@@ -6,6 +6,7 @@ This module provides a consolidated interface for generating videos using
 diffusion models.
 """
 
+import gc
 import math
 import os
 import time
@@ -276,3 +277,5 @@ class VideoGenerator:
         """
         self.executor.shutdown()
         del self.executor
+        gc.collect()
+        torch.cuda.empty_cache()

fastvideo/v1/fastvideo_args.py

@@ -292,7 +292,7 @@ class FastVideoArgs:
             assert self.sp_size != -1, "sp_size must be set for training"
 
         if self.tp_size == -1:
-            self.tp_size = 1
+            self.tp_size = self.num_gpus
         if self.sp_size == -1:
             self.sp_size = self.num_gpus
         if self.hsdp_shard_dim == -1:
@@ -305,6 +305,11 @@ class FastVideoArgs:
         if self.num_gpus < max(self.tp_size, self.sp_size):
             self.num_gpus = max(self.tp_size, self.sp_size)
 
+        if self.tp_size != self.sp_size:
+            raise ValueError(
+                f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"
+            )
+
         if self.enable_torch_compile and self.num_gpus > 1:
             logger.warning(
                 "Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"

fastvideo/v1/models/vaes/common.py

@@ -239,6 +239,7 @@ class ParallelTiledVAE(ABC):
 
         results = torch.cat(local_results, dim=0).contiguous()
         del local_results
+        torch.cuda.empty_cache()
         # first gather size to pad the results
         local_size = torch.tensor([results.size(0)],
                                   device=results.device,
@@ -252,7 +253,7 @@ class ParallelTiledVAE(ABC):
         padded_results = torch.zeros(max_size, device=results.device)
         padded_results[:results.size(0)] = results
         del results
-
+        torch.cuda.empty_cache()
         # Gather all results
         gathered_dim_metadata = [None] * world_size
         gathered_results = torch.zeros_like(padded_results).repeat(

fastvideo/v1/pipelines/stages/encoding.py

@@ -136,6 +136,7 @@ class EncodingStage(PipelineStage):
             self.maybe_free_model_hooks()
 
         self.vae.to("cpu")
+        torch.cuda.empty_cache()
 
         return batch
 

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -5,6 +5,8 @@ Image encoding stages for I2V diffusion pipelines.
 This module contains implementations of image encoding stages for diffusion pipelines.
 """
 
+import torch
+
 from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
@@ -66,6 +68,7 @@ class ImageEncodingStage(PipelineStage):
 
         if fastvideo_args.use_cpu_offload:
             self.image_encoder.to('cpu')
+            torch.cuda.empty_cache()
 
         return batch
 

fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -105,7 +105,7 @@ def run_training():
         "--num_latent_t", "8",
         "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
         "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
-        "--tp_size", 1,
+        "--tp_size", NUM_GPUS_PER_NODE_TRAINING,
         "--hsdp_replicate_dim", "1",
         "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
         "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,

fastvideo/v1/tests/ssim/README.md

@@ -24,7 +24,7 @@ FastHunyuan-diffusers: {
 "flow_shift": 17,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": true,
 "fps": 24
 }
@@ -41,7 +41,7 @@ Wan2.1-T2V-1.3B-Diffusers: {
 "flow_shift": 7.0,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": True,
 "fps": 24,
 "neg_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",
@@ -60,7 +60,7 @@ Wan2.1-I2V-14B-480P-Diffusers: {
 "flow_shift": 7.0,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": True,
 "fps": 24,
 "neg_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",

fastvideo/v1/tests/ssim/test_inference_similarity.py

@@ -33,7 +33,7 @@ HUNYUAN_PARAMS = {
     "flow_shift": 17,
     "seed": 1024,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "vae_sp": True,
     "fps": 24,
 }
@@ -50,7 +50,7 @@ WAN_T2V_PARAMS = {
     "flow_shift": 7.0,
     "seed": 1024,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "vae_sp": True,
     "fps": 24,
     "neg_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",
@@ -69,7 +69,7 @@ WAN_I2V_PARAMS = {
     "flow_shift": 7.0,
     "seed": 1024,
     "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
     "vae_sp": True,
     "fps": 24,
     "neg_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",
@@ -238,7 +238,7 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
         logger.error("Failed to write SSIM results to file")
 
     min_acceptable_ssim = 0.97
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
+    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"
 
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
 @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
@@ -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.93
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
+    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/Vanilla/test_training_loss.py

@@ -43,7 +43,7 @@ def run_worker():
         "--num_latent_t", "4",
         "--num_gpus", "4",
         "--sp_size", "4",
-        "--tp_size", "1",
+        "--tp_size", "4",
         "--hsdp_replicate_dim", "1",
         "--hsdp_shard_dim", "4",
         "--train_sp_batch_size", "1",

fastvideo/v1/training/training_pipeline.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import gc
 import math
 import os
 import time
@@ -429,7 +430,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             self.seed)
         logger.info("Initialized random seeds with seed: %s", self.seed)
 
-        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=3)
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
 
         if self.training_args.resume_from_checkpoint:
             self._resume_from_checkpoint()
@@ -439,7 +440,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         step_times: deque[float] = deque(maxlen=100)
 
         self._log_training_info()
-        # self._log_validation(self.transformer, self.training_args, 1)
+        self._log_validation(self.transformer, self.training_args, 1)
 
         # Train!
         progress_bar = tqdm(
@@ -662,6 +663,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                     frames.append((x * 255).numpy().astype(np.uint8))
                 step_videos.append(frames)
 
+            torch.distributed.barrier()
             # Only sp_group leaders (rank_in_sp_group == 0) need to send their
             # results to global rank 0
             if self.rank_in_sp_group == 0:
@@ -701,7 +703,10 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                     # Other sp_group leaders send their results to global rank 0
                     world_group.send_object(step_videos, dst=0)
                     world_group.send_object(step_captions, dst=0)
+            torch.distributed.barrier()
 
         # Re-enable gradients for training
         training_args.inference_mode = False
         transformer.train()
+        gc.collect()
+        torch.cuda.empty_cache()

fastvideo/v1/training/wan_i2v_training_pipeline.py

@@ -229,4 +229,4 @@ if __name__ == "__main__":
     parser = FastVideoArgs.add_cli_args(parser)
     args = parser.parse_args()
     args.use_cpu_offload = False
-    main(args)
+    main(args)

fastvideo/v1/worker/gpu_worker.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import contextlib
 import faulthandler
+import gc
 import multiprocessing as mp
 import os
 import signal
@@ -68,6 +69,8 @@ class Worker:
         torch.cuda.set_device(self.device)
 
         # _check_if_gpu_supports_dtype(self.model_config.dtype)
+        gc.collect()
+        torch.cuda.empty_cache()
         self.init_gpu_memory = torch.cuda.mem_get_info()[0]
 
         os.environ["MASTER_ADDR"] = "localhost"
@@ -99,6 +102,9 @@ class Worker:
         if hasattr(self, 'pipeline') and self.pipeline is not None:
             # Clean up pipeline resources if needed
             pass
+        # Release CUDA resources
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         # Destroy the distributed environment
         cleanup_dist_env_and_memory(shutdown_ray=False)
@@ -127,6 +133,8 @@ class Worker:
 
                 # Handle regular RPC calls
                 if method_name == 'execute_forward':
+                    gc.collect()
+                    torch.cuda.empty_cache()
                     forward_batch = recv_rpc['kwargs']['forward_batch']
                     fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
                     output_batch = self.execute_forward(forward_batch,

filter_parquet2.py

@@ -1,232 +0,0 @@
-#!/usr/bin/env python3
-"""
-scan_parquet_mt.py
-
-Recursively scans all Parquet files under the given root directory.
-If any row in a file contains a black frame (all pixels below threshold),
-writes a new parquet file with "filtered_" prefix and deletes the original.
-
-Features
---------
-• ThreadPoolExecutor for parallel I/O  
-• tqdm progress bar with per-file updates  
-• --dry-run flag for a safe preview  
-• --workers flag to control thread count
-• Writes filtered files in same location with "filtered_" prefix
-"""
-
-import argparse
-import hashlib
-import os
-import uuid
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from pathlib import Path
-
-import numpy as np
-import pyarrow.parquet as pq
-from PIL import Image
-from tqdm.auto import tqdm
-
-
-def process_file(path: Path,
-                 black_threshold: float = 5.0,
-                 dry_run: bool = False,
-                 images_output_dir: Path | None = None) -> int:
-    """Process a parquet file and write a filtered version with prefix. Returns number of rows removed."""
-
-    # Skip if already filtered
-    if path.stem.startswith("filtered_"):
-        # tqdm.write(f"[SKIP] Already filtered: {path}")
-        # return 0
-        truncate_prefix = True
-    else:
-        truncate_prefix = False
-
-    # Read the entire table
-    table = pq.read_table(path)
-    total_rows = len(table)
-
-    # Track which rows to keep
-    rows_to_keep = []
-    rows_removed = 0
-
-    # Check each row
-    for row_idx in range(total_rows):
-        row = table.slice(row_idx, 1).to_pylist()[0]
-
-        # Skip if any field is None
-        if row["pil_image_bytes"] is None or row[
-                "pil_image_shape"] is None or row["pil_image_dtype"] is None:
-            tqdm.write(
-                f"[WARN] Row {row_idx} in {path} has None values, keeping it")
-            rows_to_keep.append(row_idx)
-            continue
-
-        # Convert bytes to numpy array
-        image_bytes = row["pil_image_bytes"]
-        shape = row["pil_image_shape"]
-        dtype = row["pil_image_dtype"]
-
-        # Convert bytes to numpy array with proper shape and dtype
-        image_array = np.frombuffer(image_bytes,
-                                    dtype=np.float32).reshape(shape)
-        image_array = image_array.squeeze(
-            0)  # Remove single-dimensional entries if any
-
-        # Convert to uint8 for checking black frames
-        if image_array.dtype != np.uint8:
-            # Normalize to 0-255 range
-            img_min = image_array.min()
-            img_max = image_array.max()
-            if img_max > img_min:
-                image_uint8 = ((image_array - img_min) / (img_max - img_min) *
-                               255).astype(np.uint8)
-            else:
-                image_uint8 = np.zeros_like(image_array, dtype=np.uint8)
-        else:
-            image_uint8 = image_array
-
-        mean_value = np.mean(image_uint8)
-
-        # Check if the frame is black
-        if mean_value < black_threshold:
-            tqdm.write(
-                f"[INFO] Found black frame in {path} row {row_idx} (mean={mean_value:.2f})"
-            )
-            rows_removed += 1
-
-            # Save black frame for inspection with unique ID
-            if images_output_dir:
-                # Create unique filename using file path, row index, and UUID
-                # Hash the full path to handle duplicate filenames from different directories
-                path_hash = hashlib.md5(str(
-                    path.absolute()).encode()).hexdigest()[:8]
-                unique_id = f"{path.stem}_{path_hash}_row{row_idx}_{uuid.uuid4().hex[:8]}"
-
-                image_uint8 = np.transpose(image_uint8, (1, 2, 0))
-                img = Image.fromarray(image_uint8)
-                output_path = images_output_dir / f"{unique_id}.png"
-                img.save(output_path)
-                tqdm.write(f"[INFO] Saved black frame to {output_path}")
-        else:
-            # Keep this row
-            rows_to_keep.append(row_idx)
-
-    # Process based on what we found
-    if rows_removed > 0:
-        if dry_run:
-            tqdm.write(
-                f"[DRY-RUN] Would remove {rows_removed} rows from {path}")
-            tqdm.write(
-                f"[DRY-RUN] Would create filtered_{path.name} and delete original"
-            )
-        else:
-            # Create new table with only the rows to keep
-            if rows_to_keep:
-                # Filter the table to keep only non-black frames
-                new_table = table.take(rows_to_keep)
-
-                if truncate_prefix:
-                    name = path.name.replace("filtered_", "")
-                    print(f"[INFO] Truncating prefix for {path.name} to {name}")
-                else:
-                    name = path.name
-
-                # Create output path with "filtered_" prefix in same directory
-                output_path = path.parent / f"filtered_{name}"
-
-                # Write the filtered table
-                pq.write_table(new_table, output_path)
-                tqdm.write(
-                    f"[INFO] Wrote filtered parquet ({rows_removed} rows removed) to {output_path}"
-                )
-
-                # Delete original file
-                path.unlink()
-
-                tqdm.write(f"[INFO] Deleted original file: {path}")
-            else:
-                # All rows were black, just delete the file
-                path.unlink()
-                tqdm.write(
-                    f"[INFO] Deleted {path} (all {rows_removed} rows were black frames)"
-                )
-    else:
-        # No black frames found
-        if dry_run:
-            tqdm.write(f"[DRY-RUN] No black frames in {path}")
-        else:
-            # Create output path with "filtered_" prefix
-            output_path = path.parent / f"filtered_{path.name}"
-
-            # Just copy the table as-is
-            pq.write_table(table, output_path)
-            tqdm.write(
-                f"[INFO] No black frames in {path}, created {output_path}")
-
-            # Delete original file
-            path.unlink()
-            tqdm.write(f"[INFO] Deleted original file: {path}")
-
-    return rows_removed
-
-
-def handle_file(path: Path, dry_run: bool, black_threshold: float,
-                images_output_dir: Path | None) -> None:
-    """Process a single parquet file."""
-    try:
-        process_file(path,
-                     black_threshold=black_threshold,
-                     dry_run=dry_run,
-                     images_output_dir=images_output_dir)
-    except Exception as exc:
-        tqdm.write(f"[ERROR] Failed to process {path}: {exc}")
-
-
-def main(root: Path, dry_run: bool, workers: int,
-         black_threshold: float) -> None:
-    # Create output directory for black frame images
-    images_output_dir = Path.cwd() / f"filtered_{int(black_threshold)}"
-    images_output_dir.mkdir(exist_ok=True)
-    tqdm.write(f"[INFO] Saving black frames to {images_output_dir}")
-
-    parquet_files = list(root.rglob("*.parquet"))
-    if not parquet_files:
-        print(f"[INFO] No Parquet files found in {root}")
-        return
-
-    workers = max(1, workers)
-    with tqdm(total=len(parquet_files), desc="Scanning", unit="file") as bar:
-        with ThreadPoolExecutor(max_workers=workers) as pool:
-            futures = [
-                pool.submit(handle_file, fp, dry_run, black_threshold,
-                            images_output_dir) for fp in parquet_files
-            ]
-            for _ in as_completed(futures):
-                bar.update()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description=
-        "Filter black frames from Parquet files, write with 'filtered_' prefix, and delete originals."
-    )
-    parser.add_argument("--folder", type=Path, help="Root directory to scan")
-    parser.add_argument("--dry-run",
-                        action="store_true",
-                        help="Preview changes only")
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=os.cpu_count() or 32,
-        help="Number of worker threads (default: CPU count)",
-    )
-    parser.add_argument(
-        "--threshold",
-        type=float,
-        default=5.0,
-        help="Black frame threshold (default: 5.0)",
-    )
-    args = parser.parse_args()
-
-    main(args.folder, args.dry_run, args.workers, args.threshold)

finetune_1.3B_VSA_1e5_bs64_3k.sh

@@ -1,98 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=v-i-1
-#SBATCH --partition=main
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=vsa-i2v/1.3B-1e5.out
-#SBATCH --error=vsa-i2v/1.3B-1e5.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export WANDB_API_KEY='73190d8c0de18a14eb3444e222f9432d247d1e30'
-# will key
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents_i2v/train/
-DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents_i2v/test_filter/
-VALIDATION_DIR=/mnt/weka/home/hao.zhang/wl/FastVideo/data/mixkit/validation.json
-NUM_GPUS=8
-# export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# OUTPUT_PATH="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn_77x768x1280/VSA_I2V_1.3B_1e5_bs64"
-OUTPUT_PATH="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn_77x768x1280/VSA_I2V_1.3B_1e5_bs32"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-srun torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
-    --node_rank $SLURM_PROCID \
-    --rdzv_backend=c10d \
-    --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
-    --model_path $MODEL_PATH \
-    --inference_mode False\
-    --pretrained_model_name_or_path $MODEL_PATH \
-    --data_path "$DATA_DIR" \
-    --validation_dataset_file "$VALIDATION_DIR" \
-    --train_batch_size 1 \
-    --num_latent_t 16 \
-    --sp_size 1 \
-    --tp_size 1 \
-    --num_gpus 8 \
-    --hsdp_replicate_dim 8 \
-    --hsdp-shard-dim 1  \
-    --train_sp_batch_size 1 \
-    --dataloader_num_workers 4 \
-    --gradient_accumulation_steps 1 \
-    --max_train_steps 4500 \
-    --learning_rate 2e-5 \
-    --mixed_precision "bf16" \
-    --checkpointing_steps 1000 \
-    --validation_steps 300 \
-    --validation_sampling_steps "50" \
-    --log_validation \
-    --checkpoints_total_limit 3 \
-    --allow_tf32 \
-    --ema_start_step 0 \
-    --training_cfg_rate 0.1 \
-    --seed 1024 \
-    --output_dir $OUTPUT_PATH \
-    --tracker_project_name VSA_finetune \
-    --num_height 448 \
-    --num_width 832 \
-    --num_frames  61 \
-    --flow_shift 3 \
-    --validation_guidance_scale "6.0" \
-    --num_euler_timesteps 50 \
-    --master_weight_type "fp32" \
-    --dit_precision "fp32" \
-    --weight_decay 1e-4 \
-    --max_grad_norm 1.0 \
-    --VSA_decay_rate 0.03 \
-    --VSA_decay_interval_steps 50 \
-    --VSA_sparsity 0.9 

scripts/finetune/finetune_v1.sh

@@ -20,7 +20,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --train_batch_size=4 \
     --num_latent_t 20 \
     --sp_size 4 \
-    --tp_size 1 \
+    --tp_size 4 \
     --hsdp_replicate_dim 1 \
     --hsdp_shard_dim 4 \
     --num_gpus $NUM_GPUS \

scripts/inference/v1_inference_fasthunyuan.sh

@@ -3,10 +3,13 @@
 num_gpus=4
 export MODEL_BASE=FastVideo/FastHunyuan-Diffusers
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# 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 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 720 \
     --width 1280 \

scripts/inference/v1_inference_hunyuan.sh

@@ -4,10 +4,13 @@ num_gpus=4
 export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 # 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 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 720 \
     --width 1280 \

scripts/inference/v1_inference_hunyuan_STA.sh

@@ -5,10 +5,13 @@ export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_hunyuan.json
 export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 # 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 1 \
+    --tp-size ${num_gpus} \
     --height 768 \
     --width 1280 \
     --num-frames 117 \

scripts/inference/v1_inference_wan.sh

@@ -4,10 +4,13 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
 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 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 480 \
     --width 832 \

scripts/inference/v1_inference_wan_STA.sh

@@ -4,10 +4,13 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
 export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
+# 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 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 768 \
     --width 1280 \

scripts/inference/v1_inference_wan_VSA.sh

@@ -4,11 +4,14 @@ 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 
+# 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 1 \
+    --tp-size $num_gpus \
     --num-gpus $num_gpus \
     --height 448 \
     --width 832 \

scripts/inference/v1_inference_wan_i2v.sh

@@ -4,6 +4,9 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
 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
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
     --model-path $MODEL_BASE \
     --sp-size $num_gpus \