v19: Fix tied weights and lm_head loading

config.json

@@ -1,8 +1,4 @@
 {
-  "auto_map": {
-    "AutoConfig": "configuration_trm_text_ism_v6.TRMTextISMConfig",
-    "AutoModelForCausalLM": "modeling_trm_text_ism_v6.TRMTextISMForCausalLM"
-  },
   "bos_token_id": null,
   "dim": 768,
   "dropout": 0.0,

configuration_trm_text_ism_v6.py

@@ -2,9 +2,8 @@
 from transformers import PretrainedConfig
 class TRMTextISMConfig(PretrainedConfig):
     model_type = "trm_text_ism"
-    def __init__(self, vocab_size=50257, max_seq_len=512, dim=768, n_heads=12, head_dim=64, recurrence_steps=4, mlp_ratio=2.6666666667, mlp_hidden_size=None, dropout=0.0, gate_style="stable", gate_init=-1.5, residual_scale=0.5, tie_word_embeddings=True, use_cache=False, **kwargs):
-        super().__init__(tie_word_embeddings=tie_word_embeddings, use_cache=use_cache, **kwargs)
+    def __init__(self, vocab_size=50257, max_seq_len=512, dim=768, n_heads=12, head_dim=64, recurrence_steps=4, mlp_ratio=2.6666666667, mlp_hidden_size=None, dropout=0.0, gate_style="stable", gate_init=-1.5, residual_scale=0.5, tie_word_embeddings=True, **kwargs):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
         self.vocab_size, self.max_seq_len, self.dim, self.n_heads, self.head_dim = vocab_size, max_seq_len, dim, n_heads, head_dim
         self.recurrence_steps, self.mlp_ratio, self.mlp_hidden_size, self.dropout = recurrence_steps, mlp_ratio, mlp_hidden_size, dropout
         self.gate_style, self.gate_init, self.residual_scale = gate_style, gate_init, residual_scale
-        self.hidden_size, self.num_attention_heads, self.num_hidden_layers = dim, n_heads, 1

modeling_trm_text_ism_v6.py

@@ -10,8 +10,11 @@ from .configuration_trm_text_ism_v6 import TRMTextISMConfig
 def apply_rope(x, cos, sin):
     B, H, S, D = x.shape
     half = D // 2
+    # Slice to current seq length and ensure dtypes match x
+    c = cos[:, :, :S, :].to(x.dtype)
+    s = sin[:, :, :S, :].to(x.dtype)
     x1, x2 = x[..., :half], x[..., half:]
-    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)
 
 class SwiGLUMLP(nn.Module):
     def __init__(self, config):
@@ -72,12 +75,11 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
     def forward(self, input_ids, attention_mask=None, response_starts=None, **kwargs):
         B, S = input_ids.shape
         x = self.token_emb(input_ids)
-        cos, sin = self.rope_cos[:, :, :S, :], self.rope_sin[:, :, :S, :]
         mask = torch.tril(torch.ones(S, S, device=x.device)).bool().unsqueeze(0).expand(B, -1, -1)
         if response_starts is not None:
             for b in range(B):
                 rs = response_starts[b]
                 mask[b, :rs, :rs] = True
         for _ in range(self.config.recurrence_steps):
-            x = self.block(x, mask, cos, sin)
+            x = self.block(x, mask, self.rope_cos, self.rope_sin)
         return CausalLMOutputWithPast(logits=self.lm_head(self.norm(x)))