diff --git a/colossalai/shardformer/modeling/gpt2.py b/colossalai/shardformer/modeling/gpt2.py
index d7d1827625d0..30b4a3019a09 100644
--- a/colossalai/shardformer/modeling/gpt2.py
+++ b/colossalai/shardformer/modeling/gpt2.py
@@ -43,7 +43,7 @@ def _get_attention_mask(
     # Received input is already split for non-first pipeline stages,
     # but attn mask isn't
     batch_size = hidden_states.size(0)
-    seq_len = attention_mask.size(-1)
+    seq_len = attention_mask.size(-1) if attention_mask is not None else hidden_states.size(1)
 
     sp_mode = shard_config.sequence_parallelism_mode
     # If a 2D or 3D attention mask is provided for the cross-attention
diff --git a/colossalai/shardformer/policies/gpt2.py b/colossalai/shardformer/policies/gpt2.py
index ba7a5c5bc4f6..e662531990e0 100644
--- a/colossalai/shardformer/policies/gpt2.py
+++ b/colossalai/shardformer/policies/gpt2.py
@@ -277,7 +277,9 @@ def module_policy(self):
                 target_key=GPT2Attention,
             )
 
-        if not self.shard_config.pipeline_stage_manager and self.shard_config.enable_sequence_parallelism:
+        if not self.shard_config.pipeline_stage_manager and (
+            self.shard_config.enable_sequence_parallelism or use_flash_attention
+        ):
             policy[GPT2Model].method_replacement = {
                 "forward": partial(GPT2PipelineForwards.gpt2_model_forward, shard_config=self.shard_config)
             }
diff --git a/tests/test_shardformer/test_model/test_shard_gpt2.py b/tests/test_shardformer/test_model/test_shard_gpt2.py
index b67c494a6492..1eaddf463221 100644
--- a/tests/test_shardformer/test_model/test_shard_gpt2.py
+++ b/tests/test_shardformer/test_model/test_shard_gpt2.py
@@ -138,6 +138,14 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
 @parameterize(
     "test_config",
     [
+        {
+            "tp_size": 1,
+            "pp_size": 1,
+            "enable_all_optimization": True,
+            "use_lazy_init": False,
+            "precision": "fp32",
+            "initial_scale": 1,
+        },
         {
             "sp_size": 2,
             "tp_size": 1,