Fix WanVideo (#461)

Author: JerryZhou54

fastvideo/v1/dataset/parquet_datasets.py

@@ -43,7 +43,8 @@ def __init__(self,
         self.cfg_rate = cfg_rate
         self.num_latent_t = num_latent_t
         self.local_indices = None
-        self.plan_output_dir = os.path.join(self.path, "data_plan.json")
+        self.plan_output_dir = os.path.join(
+            self.path, f"data_plan_{self.world_size}_{self.sp_world_size}.json")
 
         ranks = get_sp_group().ranks
         group_ranks: List[List] = [[] for _ in range(self.world_size)]
@@ -54,6 +55,7 @@ def __init__(self,
             # This will be useful when resume training
             if os.path.exists(self.plan_output_dir):
                 print(f"Using existing plan from {self.plan_output_dir}")
+                dist.barrier()
                 return
 
             # Find all parquet files recursively, and record num_rows for each file
@@ -87,6 +89,7 @@ def __init__(self,
 
             with open(self.plan_output_dir, "w") as f:
                 json.dump(plan, f)
+        dist.barrier()
 
     def __len__(self):
         if self.local_indices is None:

fastvideo/v1/models/dits/wanvideo.py

@@ -319,6 +319,7 @@ def forward(
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+
         # Apply rotary embeddings
         cos, sin = freqs_cis
         query, key = _apply_rotary_emb(query, cos, sin,