Merge branch 'main' into sync/npu

colossalai/legacy/zero/gemini/__init__.py

+from .colo_init_context import ColoInitContext, post_process_colo_init_ctx
 from .ophooks import BaseOpHook, register_ophooks_recursively
 from .stateful_tensor import StatefulTensor
 from .stateful_tensor_mgr import StatefulTensorMgr
@@ -11,4 +12,6 @@ __all__ = [
     "AutoTensorPlacementPolicy",
     "register_ophooks_recursively",
     "BaseOpHook",
+    "ColoInitContext",
+    "post_process_colo_init_ctx",
 ]

colossalai/shardformer/_utils.py

@@ -32,7 +32,7 @@ def set_obj_list_element(obj, attr: str, value):
     r"""
     Set the element to value of a list object
 
-    It used like set_obj_list_element(obj, 'lyaers[0]', new_layer), it will set obj.layers[0] to value
+    It used like set_obj_list_element(obj, 'layers[0]', new_layer), it will set obj.layers[0] to value
 
     Args:
         obj (object): The object to set

colossalai/shardformer/layer/linear.py

@@ -408,7 +408,7 @@ class Linear1D_Row(ParallelModule):
                     handle.wait()
                 output = torch.cat(output_parallel_list, dim=-1)
         else:
-            output_parallel = F.linear(input_, self.weight)
+            output_parallel = linear_with_async_comm(input_, self.weight, None, None, False)
             if self.seq_parallel:
                 output = linear_reducescatter_forward_gather_backward(
                     output_parallel, self.process_group, self.seq_parallel_dim

colossalai/shardformer/layer/normalization.py

@@ -275,8 +275,8 @@ class FusedRMSNorm(BaseLayerNorm):
             )
 
         LazyInitContext.materialize(module)
-        # to check if it is huggingface LlamaRMSNorm
-        if module.__class__.__name__ == "LlamaRMSNorm":
+        # to check if it is huggingface LlamaRMSNorm or MistralRMSNorm
+        if module.__class__.__name__ in ["LlamaRMSNorm", "MistralRMSNorm"]:
             normalized_shape = module.weight.shape[0]
             eps = module.variance_epsilon
             elementwise_affine = True

colossalai/shardformer/policies/bloom.py

@@ -21,6 +21,15 @@ from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDe
 
 
 class BloomPolicy(Policy):
+    def __init__(self) -> None:
+        super().__init__()
+        import transformers
+        from packaging.version import Version
+
+        assert Version(transformers.__version__) <= Version(
+            "4.33.0"
+        ), "The Bloom model should run on a transformers version not greater than 4.33.0."
+
     def config_sanity_check(self):
         pass