[INF] Add Qwen2RMSNorm to loaded layers in auto_tp (microsoft#5786)

Co-authored-by: Logan Adams <[email protected]>
samadejacobs · Jul 23, 2024 · 830d0c0 · 830d0c0
1 parent 85c66fd
commit 830d0c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
@@ -134,7 +134,7 @@ def is_load_module(module):
         load_layer_names = [
             "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear",
             "MistralRMSNorm", "T5LayerNorm", "MixtralRMSNorm", "Phi3RotaryEmbedding", "Phi3SuScaledRotaryEmbedding",
-            "Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding"
+            "Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding", "Qwen2RMSNorm"
         ]
         return module.__class__ in load_layers or module._get_name() in load_layer_names