fix

2025-09-03 01:55:12 +00:00 · 2024-10-15 11:01:34 +08:00
parent 8ff7d0c780
commit 3dc08c8a5a
6 changed files with 18 additions and 15 deletions
--- a/colossalai/shardformer/modeling/gpt2.py
+++ b/colossalai/shardformer/modeling/gpt2.py
@@ -868,6 +868,7 @@ def get_gpt2_flash_attention_forward(shard_config: Optional[ShardConfig] = None)
                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)