update

modeling_bailing_moe.py

@@ -1438,21 +1438,12 @@ class BailingMoeForCausalLM(BailingMoePreTrainedModel):
 
     def compute_logit(self, hidden_states):
         if self.norm_head:
-            if self.training:
-                norm_weight = (
-                    self.lm_head.weight / (torch.norm(self.lm_head.weight, p=2, dim=0, keepdim=True) + 1e-7).detach()
-                )
-                logits = F.linear(hidden_states, norm_weight, None)
-            else:
-                self.lm_head.weight.data = (
-                    self.lm_head.weight.data.float()
-                    / (torch.norm(self.lm_head.weight.data.float(), p=2, dim=0, keepdim=True) + 1e-7)
-                ).to(hidden_states.dtype)
-                logits = F.linear(hidden_states, self.lm_head.weight.data, None)
-                self.norm_head = False
+            weight_float = self.lm_head.weight.float()
+            norm = torch.norm(weight_float, p=2, dim=0, keepdim=True).clamp(min=1e-7)
+            norm_weight = (weight_float / norm).to(hidden_states.dtype)
+            logits = F.linear(hidden_states, norm_weight, None)
         else:
             logits = self.lm_head(hidden_states)
-        return logits
 
     @add_start_docstrings_to_model_forward(BAILINGMOE_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)