diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py old mode 100644 new mode 100755 index 52d7c95ec9d8..221d490a37d2 --- a/deepspeed/module_inject/auto_tp.py +++ b/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): diff --git a/docs/_tutorials/automatic-tensor-parallelism.md b/docs/_tutorials/automatic-tensor-parallelism.md old mode 100644 new mode 100755 index e1903ed05892..d5a08b27bf4d --- a/docs/_tutorials/automatic-tensor-parallelism.md +++ b/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