Update configuration_trm_text_ism.py

configuration_trm_text_ism.py

@@ -1,17 +1,159 @@
+"""
+TRM-text-ISM: 単一ファイル構成。
+
+なぜ1ファイルにしたか:
+  config と modeling を別ファイル + relative import (`from .config import ...`) に
+  分けると、`save_pretrained()` → `push_to_hub()` → `from_pretrained(trust_remote_code=True)`
+  という往復で `ModuleNotFoundError` を起こす既知の不具合がある
+  (huggingface/transformers issue #40496, 2025-08)。
+  Falcon/ChatGLM2など実運用のHubモデルの多くも、複数ファイル構成を避けて
+  configとmodelingを1ファイルに収めることでこれを回避している。
+
+  このファイルだけを `modeling_trm_text_ism.py` としてHubに置けば、
+  Colabでの直importでも、Hubのtrust_remote_code経由でも、同一コードパスで動く。
+"""
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, PretrainedConfig
 from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
-try:
-    from .configuration_trm_text_ism import TRMTextISMConfig  # パッケージ context (trust_remote_code)
-except ImportError:
-    from configuration_trm_text_ism import TRMTextISMConfig    # 直import (Colab等)
 
 
+# ============================== Config ==============================
+
+class TRMTextISMConfig(PretrainedConfig):
+    """
+    TRM-text (ISM) config.
+
+    アーキテクチャ: RMSNorm + SwiGLU + RoPE + gated residual の `TRMBlock` を
+    `n_layers` 層積み、それを `recurrence_steps` 回ループ（recurrent-depth）。
+    n_layers=1 で「1ブロックをrecurrence_steps回」という最初の形と同一挙動。
+
+    制約: n_heads * head_dim == dim （qkvがdimにしか射影しないためMHAのみ）。
+    """
+
+    model_type = "trm_text_ism"
+    auto_map = {
+        "AutoConfig": "modeling_trm_text_ism.TRMTextISMConfig",
+        "AutoModelForCausalLM": "modeling_trm_text_ism.TRMTextISMForCausalLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 151936,
+        dim: int = 2048,
+        n_layers: int = 1,
+        n_heads: int = 16,
+        head_dim: int = 128,
+        mlp_ratio: float = 2.6875,
+        mlp_hidden_size: int | None = 5632,
+        recurrence_steps: int = 4,
+        max_seq_len: int = 2048,
+        residual_scale: float = 1.0,
+        tie_word_embeddings: bool = False,
+        pad_token_id: int | None = None,
+        bos_token_id: int | None = None,
+        eos_token_id: int | None = None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.mlp_ratio = mlp_ratio
+        self.mlp_hidden_size = mlp_hidden_size
+        self.recurrence_steps = recurrence_steps
+        self.max_seq_len = max_seq_len
+        self.residual_scale = residual_scale
+        kwargs["use_cache"] = False  # KVキャッシュ未実装。generate()のcache分岐を踏ませない
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+
+    @property
+    def hidden_size(self) -> int:
+        return self.dim
+
+    @property
+    def num_attention_heads(self) -> int:
+        return self.n_heads
+
+    @property
+    def num_hidden_layers(self) -> int:
+        return self.n_layers
+
+    @property
+    def num_key_value_heads(self) -> int:
+        return self.n_heads
+
+    @property
+    def _mlp_hidden(self) -> int:
+        return self.mlp_hidden_size or int(self.dim * self.mlp_ratio)
+
+    def param_breakdown(self) -> dict:
+        d, h, V = self.dim, self._mlp_hidden, self.vocab_size
+        attn = 3 * d * d + d * d
+        mlp = 3 * d * h
+        norms = 2 * d
+        gates = 2 * d
+        per_block = attn + mlp + norms + gates
+        blocks = self.n_layers * per_block
+        final_norm = d
+        token_emb = V * d
+        lm_head = 0 if self.tie_word_embeddings else V * d
+        non_emb = blocks + final_norm
+        total = non_emb + token_emb + lm_head
+        return {
+            "token_emb": token_emb, "lm_head": lm_head, "per_block": per_block,
+            "blocks_total": blocks, "final_norm": final_norm,
+            "non_embedding": non_emb, "total": total,
+            "embedding_share": (token_emb + lm_head) / total,
+        }
+
+    def num_parameters(self, include_embeddings: bool = True) -> int:
+        b = self.param_breakdown()
+        return b["total"] if include_embeddings else b["non_embedding"]
+
+    def __post_init_check__(self):
+        assert self.n_heads * self.head_dim == self.dim, (
+            f"n_heads*head_dim ({self.n_heads}*{self.head_dim}) != dim ({self.dim})."
+        )
+
+
+TRM_TEXT_PRESETS: dict[str, dict] = {
+    "debug": dict(dim=512, n_layers=4, n_heads=8, head_dim=64,
+                  mlp_hidden_size=1408, recurrence_steps=4, max_seq_len=1024),
+    "950m": dict(dim=2048, n_layers=7, n_heads=16, head_dim=128,
+                 mlp_hidden_size=5632, recurrence_steps=4, max_seq_len=2048),
+    "1b": dict(dim=2048, n_layers=8, n_heads=16, head_dim=128,
+               mlp_hidden_size=5632, recurrence_steps=4, max_seq_len=2048),
+    "1b-single": dict(dim=3072, n_layers=1, n_heads=24, head_dim=128,
+                      mlp_hidden_size=8192, recurrence_steps=4, max_seq_len=2048),
+    "1.3b": dict(dim=2304, n_layers=10, n_heads=18, head_dim=128,
+                 mlp_hidden_size=6144, recurrence_steps=4, max_seq_len=2048),
+}
+
+
+def trm_text_config(preset: str = "1b", **overrides) -> TRMTextISMConfig:
+    if preset not in TRM_TEXT_PRESETS:
+        raise KeyError(f"unknown preset {preset!r}. choices: {list(TRM_TEXT_PRESETS)}")
+    cfg_kwargs = {**TRM_TEXT_PRESETS[preset], **overrides}
+    cfg = TRMTextISMConfig(**cfg_kwargs)
+    cfg.__post_init_check__()
+    return cfg
+
+
+# ============================== Model ==============================
+
 def apply_rope(x, cos, sin):
     S = x.shape[2]
     c, s = cos[:, :, :S, :].to(x.dtype), sin[:, :, :S, :].to(x.dtype)
@@ -26,7 +168,7 @@ class SwiGLUMLP(nn.Module):
         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)
-        self.down_proj._scale_init = True  # 残差直前の射影は深さでスケールダウン初期化
+        self.down_proj._scale_init = True
 
     def forward(self, x):
         return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
@@ -48,7 +190,6 @@ class TRMAttention(nn.Module):
         q, k, v = [t.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) for t in (q, k, v)]
         q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
         if attn_mask is None:
-            # packed full系列（paddingなし）: flash/efficient SDPAが効く経路
             y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
         else:
             y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask[:, None, :, :])
@@ -78,7 +219,6 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
     def __init__(self, config):
         super().__init__(config)
         self.token_emb = nn.Embedding(config.vocab_size, config.dim)
-        # リカレントコア: n_layers個のユニークブロック。n_layers=1 で元コードと同一挙動。
         self.blocks = nn.ModuleList([TRMBlock(config) for _ in range(config.n_layers)])
         self.norm = nn.RMSNorm(config.dim)
         self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
@@ -87,16 +227,20 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
         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)))
         f = torch.outer(pos, theta)
-        # persistent=False: configから再計算できるのでcheckpointに焼かない
-        self.register_buffer("rope_cos", f.cos().view(1, 1, config.max_seq_len, -1), persistent=False)
-        self.register_buffer("rope_sin", f.sin().view(1, 1, config.max_seq_len, -1), persistent=False)
+        # persistent=True (デフォルト): from_pretrained の low_cpu_mem_usage 経路では
+        # モデルが meta device 上に一旦構築され、その後 state_dict から重みがロードされる。
+        # persistent=False のバッファは state_dict に乗らないため、このロード経路では
+        # meta device 上の未初期化値のまま残ってしまい、cos()/sin() の出力が
+        # 1e+34 のような異常値になってNaNが全体に伝播する事故が起きた。
+        # config から再計算可能な値であっても、ロード安全性のため persistent のままにする。
+        self.register_buffer("rope_cos", f.cos().view(1, 1, config.max_seq_len, -1))
+        self.register_buffer("rope_sin", f.sin().view(1, 1, config.max_seq_len, -1))
         self.post_init()
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
             std = 0.02
             if getattr(module, "_scale_init", False):
-                # 実効深度 = n_layers * recurrence_steps ぶん残差が積み上がるのでスケールダウン
                 eff_depth = self.config.n_layers * self.config.recurrence_steps
                 std = 0.02 / math.sqrt(2 * max(1, eff_depth))
             nn.init.normal_(module.weight, mean=0.0, std=std)
@@ -138,7 +282,6 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
         B, S = input_ids.shape
         x = self.token_emb(input_ids)
 
-        # paddingが無ければ causal flash 経路。あれば明示マスク（生成・可変長用）。
         if attention_mask is None:
             m, is_causal = None, True
         else:
@@ -148,7 +291,6 @@ class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
 
         c, s = self.rope_cos, self.rope_sin
 
-        # recurrent depth: コア(n_layers層)を recurrence_steps 回まわす
         for _ in range(self.config.recurrence_steps):
             for blk in self.blocks:
                 if self.gradient_checkpointing and self.training: