fix

2025-09-25 19:55:03 +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/shard/shard_config.py
+++ b/colossalai/shardformer/shard/shard_config.py
@@ -51,6 +51,7 @@ class ShardConfig:
    extra_kwargs: Dict[str, Any] = field(default_factory=dict)

    # For ring attention
+    pg_mesh: Optional[int] = None
    inner_ring_size: Optional[int] = None
    # for moe related
    moe_dp_group: Optional[ProcessGroup] = None