Update modeling_neollm.py

modeling_neollm.py

@@ -4679,6 +4679,9 @@ class SpellingBeeEmbedding(nn.Module):
         d    = config.hidden_size
         base = getattr(config, "rope_theta", 10000.0)
 
+        # Guardado para poder recomputar los buffers RoPE en _reset_rope_buffers.
+        self._rope_base = base
+
         # 256 × d byte embedding lookup (one per UTF-8 byte value 0..255).
         self.byte_emb = nn.Embedding(256, d)
 
@@ -4723,6 +4726,28 @@ class SpellingBeeEmbedding(nn.Module):
         # float32 for stability, applied at vocab level before the batch gather.
         self.char_norm = nn.LayerNorm(d)
 
+        # ── Hook post-carga ────────────────────────────────────────────────
+        # Los buffers non-persistent (intra_cos, intra_sin, pos_idx) se
+        # calculan desde una fórmula fija y NUNCA deben venir del safetensors.
+        # Si el checkpoint fue guardado con persistent=True (versión anterior),
+        # from_pretrained los sobreescribiría con valores corruptos.
+        # Este hook los elimina del state_dict entrante antes de la carga.
+        self._register_load_state_dict_pre_hook(self._reset_rope_buffers_hook)
+
+    def _reset_rope_buffers_hook(self, state_dict, prefix, *args, **kwargs):
+        """
+        Pre-hook de carga: elimina intra_cos, intra_sin y pos_idx del
+        state_dict entrante para que no sobreescriban los valores correctos
+        calculados en __init__.
+
+        Necesario porque versiones anteriores los guardaban como persistent=True,
+        dejando valores corruptos en el safetensors que from_pretrained cargaba
+        silenciosamente, corrompiendo los buffers tras cada load.
+        """
+        for key in [f"{prefix}intra_cos", f"{prefix}intra_sin", f"{prefix}pos_idx"]:
+            if key in state_dict:
+                del state_dict[key]
+
     def set_byte_table(self, tokenizer) -> None:
         """
         Precompute the UTF-8 byte table and inv_sqrt_lens from a tokenizer.