Update modeling_chatglm.py

modeling_chatglm.py

@@ -516,6 +516,7 @@ class GLMBlock(torch.nn.Module):
     def __init__(self, config: ChatGLMConfig, layer_number, device=None):
         super(GLMBlock, self).__init__()
         self.layer_number = layer_number
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
 
         self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
 
@@ -524,7 +525,7 @@ class GLMBlock(torch.nn.Module):
         LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
         # Layernorm on the input data.
         self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
-                                             dtype=config.torch_dtype)
+                                             dtype=dtype)
 
         # Self attention.
         self.self_attention = SelfAttention(config, layer_number, device=device)
@@ -532,7 +533,7 @@ class GLMBlock(torch.nn.Module):
 
         # Layernorm on the attention output
         self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
-                                                      dtype=config.torch_dtype)
+                                                      dtype=dtype)
 
         # MLP
         self.mlp = MLP(config, device=device)
@@ -600,9 +601,10 @@ class GLMTransformer(torch.nn.Module):
 
         if self.post_layer_norm:
             LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
+            dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
             # Final layer norm before output.
             self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
-                                                 dtype=config.torch_dtype)
+                                                 dtype=dtype)
 
         self.gradient_checkpointing = False
 
@@ -711,13 +713,14 @@ class Embedding(torch.nn.Module):
 
     def __init__(self, config: ChatGLMConfig, device=None):
         super(Embedding, self).__init__()
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
 
         self.hidden_size = config.hidden_size
         # Word embeddings (parallel).
         self.word_embeddings = nn.Embedding(
             config.padded_vocab_size,
             self.hidden_size,
-            dtype=config.torch_dtype,
+            dtype=dtype,
             device=device
         )
         self.fp32_residual_connection = config.fp32_residual_connection
@@ -748,6 +751,7 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         self.num_layers = config.num_layers
         self.multi_query_group_num = config.multi_query_group_num
         self.kv_channels = config.kv_channels
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
 
         # Rotary positional embeddings
         self.seq_length = config.seq_length
@@ -756,10 +760,10 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         )
 
         self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
-                                              dtype=config.torch_dtype)
+                                              dtype=dtype)
         self.encoder = init_method(GLMTransformer, config, **init_kwargs)
         self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
-                                        dtype=config.torch_dtype, **init_kwargs)
+                                        dtype=dtype, **init_kwargs)
         self.pre_seq_len = config.pre_seq_len
         self.prefix_projection = config.prefix_projection
         if self.pre_seq_len is not None: