[autoparallel] add numerical test for node strategies (#1760)

* [autoparallel] add numerical test for node strategies * polish code * polish code
2025-09-08 20:40:34 +00:00 · 2022-10-27 10:42:54 +08:00
parent 25952b67d7
commit b4cc59b61e
10 changed files with 283 additions and 60 deletions
--- a/colossalai/auto_parallel/tensor_shard/sharding_strategy.py
+++ b/colossalai/auto_parallel/tensor_shard/sharding_strategy.py
@@ -115,6 +115,7 @@ class CommAction:
    comm_spec: CommSpec = None
    comm_type: CommType = None
    arg_index: int = -1
+    key_for_kwarg: any = None


@dataclass