[shardformer]: fix qwen2/qwen3 causal mask with pipeline + sequence parallelism and no flash attention

colossalai/shardformer/modeling/qwen2.py

@@ -138,6 +138,16 @@ def qwen2_model_forward(
                 q_padding_mask=attention_mask,
                 is_causal=True,
             )
+        elif shard_config.enable_sequence_parallelism:
+            # the attention forward is replaced for sequence parallelism (see the policy), and it takes
+            # a 4d causal mask, rather than the mask for the attention implementation of transformers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                hidden_states,
+                past_key_values_length,
+                sliding_window=self.config.sliding_window,
+            )
         else:
             if self.config._attn_implementation == "flash_attention_2":
                 # 2d mask is passed through the layers

colossalai/shardformer/modeling/qwen3.py

@@ -136,6 +136,16 @@ def qwen3_model_forward(
                 q_padding_mask=attention_mask,
                 is_causal=True,
             )
+        elif shard_config.enable_sequence_parallelism:
+            # the attention forward is replaced for sequence parallelism (see the policy), and it takes
+            # a 4d causal mask, rather than the mask for the attention implementation of transformers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                hidden_states,
+                past_key_values_length,
+                sliding_window=self.config.sliding_window,
+            )
         else:
             if self.config._attn_implementation == "flash_attention_2":
                 # 2d mask is passed through the layers

tests/test_checkpoint_io/test_gemini_checkpoint_io.py

@@ -23,7 +23,7 @@
 
 ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False]
 try:
-    import apex.amp
+    pass
 
     APEX_AVAILABLE = True
 except ImportError:

tests/test_shardformer/test_model/test_shard_qwen2.py

@@ -120,6 +120,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
             "precision": "fp16",
             "initial_scale": 1,
         },
+        {  # Ulysess without flash attention, the replaced attention forward needs a 4d causal mask
+            "tp_size": 1,
+            "pp_size": 2,
+            "sp_size": 2,
+            "num_microbatches": 2,
+            "enable_sequence_parallelism": True,
+            "sequence_parallelism_mode": "all_to_all",
+            "enable_flash_attention": False,
+            "use_lazy_init": True,
+            "zero_stage": 1,
+            "precision": "fp16",
+            "initial_scale": 1,
+        },
         {
             "tp_size": 2,
             "pp_size": 2,

tests/test_shardformer/test_model/test_shard_qwen3.py

@@ -116,6 +116,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
             "precision": "fp16",
             "initial_scale": 1,
         },
+        {  # Ulysess without flash attention, the replaced attention forward needs a 4d causal mask
+            "tp_size": 1,
+            "pp_size": 2,
+            "sp_size": 2,
+            "num_microbatches": 2,
+            "enable_sequence_parallelism": True,
+            "sequence_parallelism_mode": "all_to_all",
+            "enable_flash_attention": False,
+            "use_lazy_init": True,
+            "zero_stage": 1,
+            "precision": "fp16",
+            "initial_scale": 1,
+        },
         {
             "tp_size": 2,
             "pp_size": 2,

tests/test_zero/test_gemini/test_optim.py

@@ -15,7 +15,7 @@
 from tests.kit.model_zoo import model_zoo, run_fwd_bwd
 
 try:
-    import apex.amp
+    pass
 
     APEX_AMP_AVAILABLE = True
 except ImportError: