fix

2025-09-25 19:55:03 +00:00 · 2024-10-15 11:56:49 +08:00
parent 3dc08c8a5a
commit 6be9862aaf
6 changed files with 23 additions and 6 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
+    sp_axis: Optional[int] = None
    pg_mesh: Optional[int] = None
    inner_ring_size: Optional[int] = None
    # for moe related