fix

colossalai/shardformer/policies/bloom.py

@@ -11,6 +11,7 @@ import colossalai.shardformer.layer as col_nn
 from ..modeling.bloom import (
     BloomPipelineForwards,
     build_bloom_alibi_tensor_fn,
+    get_bloom_sequence_parallel_attention_forward,
     get_bloom_sequence_parallel_forward_fn,
     get_jit_fused_bloom_attention_forward,
     get_jit_fused_bloom_gelu_forward,
@@ -61,6 +62,15 @@ class BloomPolicy(Policy):
 
         use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv
 
+        if self.shard_config.enable_sequence_parallelism:
+            self.append_or_create_method_replacement(
+                description={
+                    "forward": get_bloom_sequence_parallel_attention_forward(self.shard_config),
+                },
+                policy=policy,
+                target_key=BloomAttention,
+            )
+
         if self.shard_config.enable_tensor_parallelism:
             assert (
                 self.model.config.n_head % self.shard_config.tensor_parallel_size == 0