Update modeling_neollm.py

modeling_neollm.py

@@ -853,15 +853,24 @@ class LeviathanGenerator(nn.Module):
         # head_proj_weight [M*d_seed, d_seed] — los pesos del cabezal m
         # son las filas [m*d_seed : (m+1)*d_seed].
         proj_w = self.head_proj_weight[m * d : (m + 1) * d]       # [d_seed, d_seed]
-        zh = F.linear(z.float(), proj_w)                           # [N, d_seed]
+        # Keep the matmul in the parameter dtype so eager inference matches
+        # mixed-precision training, then promote to float32 for the reduction-
+        # heavy normalization and KHRONOS path below.
+        zh = F.linear(
+            z.to(dtype=proj_w.dtype, device=proj_w.device),
+            proj_w,
+        )                                                          # [N, d_seed]
+        zh = zh.float()
 
         # ── LayerNorm manual por cabezal ──────────────────────────────────
         # Equivalente a nn.LayerNorm(d_seed) con parámetros independientes
         # head_norm_weight[m] y head_norm_bias[m].
+        norm_w = self.head_norm_weight[m].float()
+        norm_b = self.head_norm_bias[m].float()
         mean = zh.mean(dim=-1, keepdim=True)
         var  = zh.var(dim=-1, keepdim=True, unbiased=False)
         zh   = (zh - mean) / (var + self.head_norm_eps).sqrt()
-        zh   = zh * self.head_norm_weight[m] + self.head_norm_bias[m]
+        zh   = zh * norm_w + norm_b
 
         # ── Sigmoid(x/2) → coordenada latente en [0,1]^d_seed ────────────
         zh = torch.sigmoid(zh / 2.0).clamp(0.0, 1.0)              # [N, d_seed]
@@ -4230,7 +4239,12 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             output_attentions if output_attentions is not None
             else self.config.output_attentions
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if return_dict is None:
+            cfg_dict = vars(self.config)
+            return_dict = cfg_dict.get(
+                "return_dict",
+                cfg_dict.get("use_return_dict", True),
+            )
 
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("Specify exactly one of input_ids or inputs_embeds")
@@ -4723,4 +4737,4 @@ __all__ = [
 
 AutoConfig.register("neollm", NeoLLMConfig)
 AutoModel.register(NeoLLMConfig, NeoLLMModel)
-AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)
+AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)