update test with multigpu (#207)

* update test

* update

Author: oahzxl

tests/pipelines/cogvideox/test_cogvideox.py

@@ -3,7 +3,7 @@
 from videosys import CogVideoXConfig, VideoSysEngine
 
 
-@pytest.mark.parametrize("num_gpus", [1])
+@pytest.mark.parametrize("num_gpus", [1, 2])
 def test_base(num_gpus):
     config = CogVideoXConfig(num_gpus=num_gpus)
     engine = VideoSysEngine(config)

tests/pipelines/latte/test_latte.py

@@ -13,9 +13,9 @@ def test_base(num_gpus):
     engine.save_video(video, f"./test_outputs/{prompt}_latte_{num_gpus}.mp4")
 
 
-@pytest.mark.parametrize("num_gpus", [1, 2])
+@pytest.mark.parametrize("num_gpus", [1])
 def test_pab(num_gpus):
-    config = LatteConfig(num_gpus=num_gpus)
+    config = LatteConfig(num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."

tests/pipelines/open_sora/test_open_sora.py

@@ -13,7 +13,7 @@ def test_base(num_gpus):
     engine.save_video(video, f"./test_outputs/{prompt}_open_sora_{num_gpus}.mp4")
 
 
-@pytest.mark.parametrize("num_gpus", [1, 2])
+@pytest.mark.parametrize("num_gpus", [1])
 def test_pab(num_gpus):
     config = OpenSoraConfig(num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)

tests/pipelines/open_sora_plan/test_open_sora_plan.py

@@ -10,12 +10,12 @@ def test_base(num_gpus):
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_pab_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_{num_gpus}.mp4")
 
 
-@pytest.mark.parametrize("num_gpus", [1, 2])
+@pytest.mark.parametrize("num_gpus", [1])
 def test_pab(num_gpus):
-    config = OpenSoraPlanConfig(num_gpus=num_gpus)
+    config = OpenSoraPlanConfig(num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."

videosys/models/transformers/open_sora_plan_transformer_3d.py

@@ -1258,23 +1258,9 @@ def __call__(
 
         # the output of sdp = (batch, num_heads, seq_len, head_dim)
         # TODO: add support for attn.scale when we move to Torch 2.1
-        if self.attention_mode == "flash":
-            assert attention_mask is None or torch.all(
-                attention_mask.bool()
-            ), "flash-attn do not support attention_mask"
-            with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
-                hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
-        elif self.attention_mode == "xformers":
-            with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
-                hidden_states = F.scaled_dot_product_attention(
-                    query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-                )
-        elif self.attention_mode == "math":
-            hidden_states = F.scaled_dot_product_attention(
-                query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-            )
-        else:
-            raise NotImplementedError(f"Found attention_mode: {self.attention_mode}")
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
         hidden_states = hidden_states.to(query.dtype)