fix(model): fix norm input dtype

InternLM · Dec 3, 2024 · 8081e61 · 8081e61
1 parent 9d2d76a
commit 8081e61
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 8 deletions.
diff --git a/internlm/model/modeling_internlm.py b/internlm/model/modeling_internlm.py
@@ -195,7 +195,7 @@ def _forward(self, hidden_states, *args, **kwargs):
         def _dropout_and_norm_attn(_hidden_states):
             _dropped = self.dropout1(_hidden_states)
             _residual = _dropped
-            _hidden_states = self.norm1(_residual.float())
+            _hidden_states = self.norm1(_residual.to(self.norm1.weight.dtype))
             return _residual, _hidden_states
 
         if self.dropout_selective_checkpoint:
@@ -212,7 +212,7 @@ def _dropout_and_norm_attn(_hidden_states):
         def _dropout_and_norm_ffn(_residual, _hidden_states):
             _dropped = self.dropout2(_hidden_states)
             _residual = (_dropped + _residual) if _residual is not None else _dropped
-            _hidden_states = self.norm2(_residual.float())
+            _hidden_states = self.norm2(_residual.to(self.norm2.weight.dtype))
             return _residual, _hidden_states
 
         if self.dropout_selective_checkpoint:

diff --git a/internlm/model/modeling_internlm2.py b/internlm/model/modeling_internlm2.py
@@ -254,7 +254,7 @@ def _dropout_and_norm_attn(_residual, _hidden_states):
                     def _dropout_and_norm_ffn(_residual, _hidden_states):
                         _dropped = self.dropout2(_hidden_states)
                         _residual = (_dropped + _residual) if _residual is not None else _dropped
-                        _hidden_states = self.ffn_norm(_residual.to(torch.float32))
+                        _hidden_states = self.ffn_norm(_residual.to(self.ffn_norm.weight.dtype))
 
                         return _residual, _hidden_states
 

diff --git a/internlm/model/modeling_llama.py b/internlm/model/modeling_llama.py
@@ -246,7 +246,7 @@ def _dropout_and_norm_attn(_residual, _hidden_states):
                     def _dropout_and_norm_ffn(_residual, _hidden_states):
                         _dropped = self.dropout2(_hidden_states)
                         _residual = (_dropped + _residual) if _residual is not None else _dropped
-                        _hidden_states = self.ffn_norm(_residual.to(torch.float32))
+                        _hidden_states = self.ffn_norm(_residual.to(self.ffn_norm.weight.dtype))
 
                         return _residual, _hidden_states
 

diff --git a/internlm/model/modeling_mixtral.py b/internlm/model/modeling_mixtral.py
@@ -214,7 +214,7 @@ def _forward(self, hidden_states, *args, **kwargs):
         def _dropout_and_norm_attn(_hidden_states):
             _dropped = self.dropout1(_hidden_states)
             _residual = _dropped
-            _hidden_states = self.norm1(_residual.float())
+            _hidden_states = self.norm1(_residual.to(self.norm1.weight.dtype))
             return _residual, _hidden_states
 
         if self.dropout_selective_checkpoint:
@@ -231,7 +231,7 @@ def _dropout_and_norm_attn(_hidden_states):
         def _dropout_and_norm_ffn(_residual, _hidden_states):
             _dropped = self.dropout2(_hidden_states)
             _residual = (_dropped + _residual) if _residual is not None else _dropped
-            _hidden_states = self.norm2(_residual.float())
+            _hidden_states = self.norm2(_residual.to(self.norm2.weight.dtype))
             return _residual, _hidden_states
 
         if self.dropout_selective_checkpoint:

diff --git a/internlm/model/modeling_moe.py b/internlm/model/modeling_moe.py
@@ -205,7 +205,7 @@ def _forward(self, hidden_states, *args, **kwargs):
         def _dropout_and_norm_attn(_hidden_states):
             _dropped = self.dropout1(_hidden_states)
             _residual = _dropped
-            _hidden_states = self.norm1(_residual.float())
+            _hidden_states = self.norm1(_residual.to(self.norm1.weight.dtype))
             return _residual, _hidden_states
 
         if self.dropout_selective_checkpoint:
@@ -222,7 +222,7 @@ def _dropout_and_norm_attn(_hidden_states):
         def _dropout_and_norm_ffn(_residual, _hidden_states):
             _dropped = self.dropout2(_hidden_states)
             _residual = (_dropped + _residual) if _residual is not None else _dropped
-            _hidden_states = self.norm2(_residual.float())
+            _hidden_states = self.norm2(_residual.to(self.norm2.weight.dtype))
             return _residual, _hidden_states
 
         if self.dropout_selective_checkpoint: