Fix T5 and mistral model meta data error (#4958)

Fix 'NotImplementedError: Cannot copy out of meta tensor; no data!', when loading T5 and mistral from device meta. Co-authored-by: Logan Adams <[email protected]>
deepspeedai · Jan 19, 2024 · e62a47e · e62a47e
1 parent 96c5a87
commit e62a47e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
@@ -132,7 +132,8 @@ class Loading():
     def is_load_module(module):
         load_layers = [nn.Linear, nn.Embedding, nn.LayerNorm]
         load_layer_names = [
-            "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear"
+            "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear",
+            "MistralRMSNorm", "T5LayerNorm"
         ]
         return module.__class__ in load_layers or module._get_name() in load_layer_names