[Sharderformer] Support zbv in Sharderformer Policy

colossalai/pipeline/weight_grad_store.py

@@ -8,7 +8,6 @@ class WeightGradStore:
 
     @classmethod
     def put(cls, total_input, grad_output, weight, func):
-        # func(total_input, grad_output, weight.main_grad)
         cls.cache.append((total_input, grad_output, weight, func))
 
     @classmethod
@@ -18,7 +17,6 @@ def flush(cls, chunk=0):
 
     @classmethod
     def pop(cls, chunk=0):
-        # print(f"chunk id {chunk} queue size {cls.weight_grad_queue[chunk].qsize()}")
         if cls.weight_grad_queue[chunk].qsize() > 0:
             stored_grads = cls.weight_grad_queue[chunk].get()
             for total_input, grad_output, weight, func in stored_grads:

colossalai/shardformer/layer/__init__.py

@@ -6,7 +6,14 @@
 from .loss import cross_entropy_1d, dist_cross_entropy
 from .normalization import FusedLayerNorm, FusedRMSNorm, LayerNorm, RMSNorm
 from .parallel_module import ParallelModule
-from .qkv_fused_linear import FusedLinear1D_Col, FusedLinear1D_Row, GPT2FusedLinearConv1D_Col, GPT2FusedLinearConv1D_Row
+from .qkv_fused_linear import (
+    FusedLinear1D,
+    FusedLinear1D_Col,
+    FusedLinear1D_Row,
+    GPT2FusedLinearConv1D,
+    GPT2FusedLinearConv1D_Col,
+    GPT2FusedLinearConv1D_Row,
+)
 
 __all__ = [
     "Embedding1D",
@@ -16,6 +23,7 @@
     "Linear1D_Row",
     "GPT2FusedLinearConv1D_Col",
     "GPT2FusedLinearConv1D_Row",
+    "GPT2FusedLinearConv1D_Col",
     "DropoutForParallelInput",
     "DropoutForReplicatedInput",
     "cross_entropy_1d",
@@ -26,6 +34,7 @@
     "FusedLayerNorm",
     "FusedRMSNorm",
     "FusedLinear1D_Col",
+    "FusedLinear1D",
     "ParallelModule",
     "PaddingEmbedding",
     "PaddingLMHead",