Merge pull request #6071 from wangbluo/ring_attention

[Ring Attention] fix the 2d ring attn when using multiple machine

colossalai/shardformer/modeling/gpt2.py

@@ -857,17 +857,17 @@ def get_gpt2_flash_attention_forward(shard_config: Optional[ShardConfig] = None)
         dropout_p = self.attn_dropout.p if self.training else 0.0
 
         sp_mode = shard_config.sequence_parallelism_mode
-        sp_group = shard_config.sequence_parallel_process_group
         if sp_mode == "ring_attn":
             attn_output = RingAttention.attention(
                 query,
                 key,
                 value,
-                sp_group,
+                sp_axis=shard_config.sp_axis,
                 **attention_mask,
                 dropout_p=dropout_p,
                 scale=scale,
                 inner_ring_size=shard_config.inner_ring_size,
+                pg_mesh=shard_config.pg_mesh,
             )
         else:
             attn_output = ColoAttention.attention(query, key, value, **attention_mask, dropout_p=dropout_p, scale=scale)

colossalai/shardformer/modeling/llama.py

@@ -569,9 +569,10 @@ def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, s
                 query_states,
                 key_states,
                 value_states,
-                sp_group,
+                sp_axis=shard_config.sp_axis,
                 **attention_mask,
                 inner_ring_size=shard_config.inner_ring_size,
+                pg_mesh=shard_config.pg_mesh,
             )
 
         elif shard_config.enable_flash_attention: