Update modeling_chatglm.py

modeling_chatglm.py

@@ -489,9 +489,10 @@ class GLMBlock(torch.nn.Module):
         self.fp32_residual_connection = config.fp32_residual_connection
 
         LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
         # 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)
@@ -499,7 +500,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)
@@ -567,9 +568,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
 
@@ -690,10 +692,11 @@ class Embedding(torch.nn.Module):
 
         self.hidden_size = config.hidden_size
         # Word embeddings (parallel).
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
         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
@@ -728,12 +731,12 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         rotary_dim = (
             config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
         )
-
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
         self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio, original_impl=config.original_rope, 
-                                              device=device, dtype=config.torch_dtype)
+                                              device=device, 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)
 
     def get_input_embeddings(self):
         return self.embedding.word_embeddings
@@ -1153,8 +1156,9 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
 
         self.num_labels = config.num_labels
         self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
+        dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
 
-        self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=config.torch_dtype)
+        self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=dtype)
         if config.classifier_dropout is not None:
             self.dropout = nn.Dropout(config.classifier_dropout)
         else: