[shardformer] integrated linear 1D with dtensor (#3996)

* [shardformer] integrated linear 1D with dtensor * polish code
2025-09-16 14:41:53 +00:00 · 2023-06-15 18:03:38 +08:00
parent d3bc530849
commit 015af592f8
9 changed files with 707 additions and 408 deletions
--- a/colossalai/nn/layer/base_layer.py
+++ b/colossalai/nn/layer/base_layer.py
@@ -10,6 +10,7 @@ from colossalai.core import global_context as gpc


 class ParallelLayer(nn.Module):
+
    global_state_dict: bool = True

    def __init__(self):