From 6b0ff1eb697db72541488a40dce824a75d3d1733 Mon Sep 17 00:00:00 2001 From: "Deng, Daisy" Date: Wed, 3 Jan 2024 19:36:07 -0800 Subject: [PATCH] fix t5 and mistral model load from config meta tensor bug --- deepspeed/module_inject/auto_tp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 20dd19fe32a4..10bc1bb94140 100644 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -122,7 +122,7 @@ class Loading(): def is_load_module(module): load_layers = [nn.Linear, nn.Embedding, nn.LayerNorm] load_layer_names = [ - "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "RMSNorm" + "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "RMSNorm", "MistralRMSNorm", "T5LayerNorm", ] return module.__class__ in load_layers or module._get_name() in load_layer_names