Fix precision error

modeling_chatglm.py

@@ -3,9 +3,7 @@
 import math
 import copy
 import warnings
-import re
 import sys
-
 import torch
 import torch.utils.checkpoint
 import torch.nn.functional as F
@@ -183,9 +181,14 @@ class RMSNorm(torch.nn.Module):
         self.eps = eps
 
     def forward(self, hidden_states: torch.Tensor):
-        input_dtype = hidden_states.dtype
-        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        if hidden_states == torch.bfloat16:
+            norm_x = torch.mean(hidden_states * hidden_states, dim=-1, keepdim=True)
+            x_normed = hidden_states * torch.rsqrt(norm_x + self.eps)
+            return self.weight * x_normed
+        else:
+            input_dtype = hidden_states.dtype
+            variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+            hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
 
         return (self.weight * hidden_states).to(input_dtype)
 
@@ -517,8 +520,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)
+        self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, dtype=config.torch_dtype)
 
         # Self attention.
         self.self_attention = SelfAttention(config, layer_number, device=device)