diff --git a/colossalai/shardformer/modeling/qwen2.py b/colossalai/shardformer/modeling/qwen2.py
index 80dfb003dccf..c1a4ae161407 100644
--- a/colossalai/shardformer/modeling/qwen2.py
+++ b/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
diff --git a/colossalai/shardformer/modeling/qwen3.py b/colossalai/shardformer/modeling/qwen3.py
index 9b6670ba051b..01982ef1e0d0 100644
--- a/colossalai/shardformer/modeling/qwen3.py
+++ b/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
diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
index ba60de0bd1c1..3649b18ec8bb 100644
--- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
+++ b/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:
diff --git a/tests/test_shardformer/test_model/test_shard_qwen2.py b/tests/test_shardformer/test_model/test_shard_qwen2.py
index 865563adc625..193be8bce467 100644
--- a/tests/test_shardformer/test_model/test_shard_qwen2.py
+++ b/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,
diff --git a/tests/test_shardformer/test_model/test_shard_qwen3.py b/tests/test_shardformer/test_model/test_shard_qwen3.py
index 9670a5999e8c..7c76cd834f1b 100644
--- a/tests/test_shardformer/test_model/test_shard_qwen3.py
+++ b/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,
diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py
index b89c591882d6..a0970e43db9e 100644
--- a/tests/test_zero/test_gemini/test_optim.py
+++ b/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: