Merge branch 'master' into jomayeri/aio-locked-tensor

deepspeed/module_inject/auto_tp.py

@@ -333,7 +333,8 @@ def _replace(self, child, name, conv_linear_layer):
         weight_shape = child.weight.shape
         mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
         # For mixtral-7x8b, need to skip MoE gate linear replace.
-        if name == "block_sparse_moe.gate":
+        if name == "block_sparse_moe.gate" or (('mlp.shared_expert_gate' == name or 'mlp.gate' == name)
+                                               and 'qwen2_moe' in str(type(self.module))):
             return child
         # For Yuan model
         if 'Yuan' in str(self.module):

docs/_tutorials/automatic-tensor-parallelism.md

@@ -158,6 +158,7 @@ The following model families have been successfully tested with automatic tensor
 - plbart
 - qwen
 - qwen2
+- qwen2-moe
 - reformer
 - roberta
 - roformer