Resolving inference compatibility issues in the Kormo model’s Transformer 5.2

_modeling_kormo.py

@@ -95,6 +95,12 @@ def rotate_half(x):
     x2 = x[..., x.shape[-1] // 2 :]
     return torch.cat((-x2, x1), dim=-1)
 
+def copy_(tensor: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        with torch.no_grad():
+            return tensor.copy_(other)
+    return tensor
+
 class Attention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
@@ -237,12 +243,25 @@ class RotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        rope_init_fn = self.compute_default_rope_parameters
+        rope_init_fn = ROPE_INIT_FUNCTIONS.get(self.rope_type, rope_init_fn)
+        self.rope_init_fn = rope_init_fn
 
         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
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(config: KORMoConfig, device=None, seq_len =None):
+        base = config.rope_theta
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
 
+        attention_factor = 1.0
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+    
     @torch.no_grad()
     @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
     def forward(self, x, position_ids):
@@ -258,10 +277,6 @@ class RotaryEmbedding(nn.Module):
             return cos, sin
     
 
-
-
-
-
 class KORMoPreTrainedModel(PreTrainedModel):
     config_class = KORMoConfig
     base_model_prefix = "model"
@@ -289,7 +304,15 @@ class KORMoPreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
         elif isinstance(module, RMSNorm):
             module.weight.data.fill_(1.0)
-
+        elif "RotaryEmbedding" in module.__class__.__name__ and hasattr(module, "original_inv_freq"):
+            rope_fn = (
+                ROPE_INIT_FUNCTIONS[module.rope_type]
+                if module.rope_type != "default"
+                else module.compute_default_rope_parameters
+            )
+            buffer_value, _ = rope_fn(module.config)
+            copy_(module.inv_freq, buffer_value)
+            copy_(module.original_inv_freq, buffer_value)
 
 class KORMoModel(KORMoPreTrainedModel):
     def __init__(self, config: KORMoConfig):