v29: Strict weight key alignment (norm1, norm2, attn_gate, mlp_gate)

configuration_trm_text_ism.py

@@ -2,7 +2,7 @@
 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, **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=2048, 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

modeling_trm_text_ism.py

@@ -17,7 +17,8 @@ class SwiGLUMLP(nn.Module):
     def __init__(self, config):
         super().__init__()
         h = config.mlp_hidden_size or int(config.dim * config.mlp_ratio)
-        self.gate_proj, self.up_proj = nn.Linear(config.dim, h, bias=False), nn.Linear(config.dim, h, bias=False)
+        self.gate_proj = nn.Linear(config.dim, h, bias=False)
+        self.up_proj = nn.Linear(config.dim, h, bias=False)
         self.down_proj = nn.Linear(h, config.dim, bias=False)
     def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
 
@@ -25,7 +26,8 @@ class TRMAttention(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.n_heads, self.head_dim = config.n_heads, config.head_dim
-        self.qkv, self.out = nn.Linear(config.dim, 3*config.dim, bias=False), nn.Linear(config.dim, config.dim, bias=False)
+        self.qkv = nn.Linear(config.dim, 3*config.dim, bias=False)
+        self.out = nn.Linear(config.dim, config.dim, bias=False)
     def forward(self, x, mask, cos, sin):
         B, S, _ = x.shape
         q, k, v = self.qkv(x).chunk(3, dim=-1)
@@ -38,10 +40,10 @@ class TRMBlock(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.res = config.residual_scale
-        # 変数名をチェックポイントのキー (block.norm1.weight等) に合わせる
-        self.norm1, self.norm2 = nn.RMSNorm(config.dim), nn.RMSNorm(config.dim)
-        self.attn, self.mlp = TRMAttention(config), SwiGLUMLP(config)
-        # block.attn_gate, block.mlp_gate に合わせる
+        self.norm1 = nn.RMSNorm(config.dim)
+        self.attn = TRMAttention(config)
+        self.norm2 = nn.RMSNorm(config.dim)
+        self.mlp = SwiGLUMLP(config)
         self.attn_gate = nn.Parameter(torch.ones(config.dim))
         self.mlp_gate = nn.Parameter(torch.ones(config.dim))
     def forward(self, x, mask, c, s):
@@ -52,7 +54,9 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = TRMTextISMConfig
     def __init__(self, config):
         super().__init__(config)
-        self.token_emb, self.block, self.norm = nn.Embedding(config.vocab_size, config.dim), TRMBlock(config), nn.RMSNorm(config.dim)
+        self.token_emb = nn.Embedding(config.vocab_size, config.dim)
+        self.block = TRMBlock(config)
+        self.norm = nn.RMSNorm(config.dim)
         self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
         pos = torch.arange(config.max_seq_len).float()
         theta = 1.0 / (10000.0 ** (torch.arange(0, config.head_dim//2).float() / (config.head_dim//2)))
@@ -71,7 +75,8 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
         x = self.token_emb(input_ids)
         m = torch.tril(torch.ones(S, S, device=input_ids.device)).bool().unsqueeze(0).expand(B, -1, -1)
         if attention_mask is not None: m = m & attention_mask[:, None, :].bool()
+        c, s = self.rope_cos, self.rope_sin
         for _ in range(self.config.recurrence_steps):
-            x = self.block(x, m, self.rope_cos, self.rope_sin)
+            x = self.block(x, m, c, s)
         logits = self.lm_head(self.norm(x))
         return CausalLMOutputWithPast(logits=logits)