fix: RoPE inv_freq fp32 construction + clean-copy self-heal (NPU state-dict load corruption)

modeling_moss_vl.py

@@ -209,11 +209,14 @@ class MossVLVisionRotaryEmbedding(nn.Module):
 
     def __init__(self, dim: int, theta: float = 10000.0) -> None:
         super().__init__()
-        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float, device="cpu") / dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._inv_freq_clean = inv_freq.clone()
 
     def forward(self, seqlen: int) -> torch.Tensor:
-        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        if torch.isnan(self.inv_freq).any() or torch.isinf(self.inv_freq).any():
+            self.inv_freq = self._inv_freq_clean.to(self.inv_freq.device)
+        seq = torch.arange(seqlen, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
         freqs = torch.outer(seq, self.inv_freq)
         return freqs
 
@@ -445,12 +448,23 @@ class MossVLTextRotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        if self.rope_type == "default" or self.rope_type not in ROPE_INIT_FUNCTIONS:
+            import torch as _torch
+            base = config.rope_theta if hasattr(config, "rope_theta") else 10000.0
+            head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+            inv_freq = 1.0 / (base ** (_torch.arange(0, head_dim, 2, device="cpu").float() / head_dim))
+            self.attention_scaling = 1.0
+        else:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.original_inv_freq = self.inv_freq
 
+        # Guard: from_pretrained may corrupt inv_freq with garbage from the
+        # state dict (persistent=False buffers can still be overwritten during
+        # weight loading on some transformers versions). Re-validate on forward.
+        self._inv_freq_computed = inv_freq.clone()
+
 
         if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
              self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
@@ -477,6 +491,10 @@ class MossVLTextRotaryEmbedding(nn.Module):
     @torch.no_grad()
     @dynamic_rope_update
     def forward(self, x, position_ids):
+        # Guard: from_pretrained can corrupt inv_freq — restore from clean copy
+        if torch.isnan(self.inv_freq).any() or torch.isinf(self.inv_freq).any():
+            self.inv_freq = self._inv_freq_computed.to(self.inv_freq.device)
+            self.original_inv_freq = self._inv_freq_computed.to(self.inv_freq.device)
 
         if position_ids.ndim == 2:
             position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)