Fix model architecture and generation compatibility

modeling_tiny_recursive.py

@@ -1,11 +1,11 @@
 
-import torch
-import torch.nn as nn
-from transformers import PreTrainedModel, PretrainedConfig, GPT2TokenizerFast, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers import PreTrainedModel, PretrainedConfig
 from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
-from transformers.generation import GenerationMixin  # <--- FIXED: Import this explicitly
+from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
-from datasets import load_dataset
+import torch
+import torch.nn as nn
+
 class TRMConfig(PretrainedConfig):
     model_type = "recursive_gpt"
 
@@ -14,9 +14,9 @@ class TRMConfig(PretrainedConfig):
         vocab_size=50257,
         n_positions=1024,
         n_embd=512,
+        n_physical_layers=3,
+        n_loops=8,
         n_head=8,
-        n_physical_layers=2,
-        n_loops=6,
         activation_function="gelu_new",
         resid_pdrop=0.1,
         embd_pdrop=0.1,
@@ -28,13 +28,12 @@ class TRMConfig(PretrainedConfig):
         **kwargs,
     ):
         super().__init__(**kwargs)
-        # Standard config
         self.vocab_size = vocab_size
         self.n_positions = n_positions
         self.n_embd = n_embd
-        self.n_head = n_head
         self.n_physical_layers = n_physical_layers
         self.n_loops = n_loops
+        self.n_head = n_head
         self.activation_function = activation_function
         self.resid_pdrop = resid_pdrop
         self.embd_pdrop = embd_pdrop
@@ -44,17 +43,16 @@ class TRMConfig(PretrainedConfig):
         self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
         self.reorder_and_upcast_attn = reorder_and_upcast_attn
 
-        # --- CRITICAL FIXES FOR COMPATIBILITY ---
-        # These map your custom names to what GPT2Attention expects
-        self.max_position_embeddings = n_positions 
+        # Required for transformers compatibility
         self.hidden_size = n_embd
-        self.num_attention_heads = n_head          # <--- FIXED: The missing attribute
+        self.num_attention_heads = n_head
         self.num_hidden_layers = n_physical_layers
-        self.n_inner = None                        # Defaults to 4*hidden_size
+        self.n_inner = None
+        self.is_encoder_decoder = False
 
 class TinyRecursiveModel(PreTrainedModel, GenerationMixin):
     config_class = TRMConfig
-    _tied_weights_keys = ["lm_head.weight"]  # <-- Add this line
+    _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
         super().__init__(config)
@@ -65,90 +63,74 @@ class TinyRecursiveModel(PreTrainedModel, GenerationMixin):
         self.wpe = nn.Embedding(config.n_positions, config.n_embd)
         self.drop = nn.Dropout(config.embd_pdrop)
 
-        # 2. The Logic Core (The "7M" part)
+        # 2. Physical blocks - matching your saved model structure
         self.physical_blocks = nn.ModuleList([
-            RecursiveBlock(config, layer_idx=i) for i in range(config.n_physical_layers)
+            nn.ModuleDict({
+                "ln_1": nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon),
+                "attn": GPT2Attention(config, layer_idx=i),
+                "ln_2": nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon),
+                "mlp": GPT2MLP(4 * config.n_embd, config)
+            }) for i in range(config.n_physical_layers)
         ])
 
+        # 3. Final layer norm
         self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+
+        # 4. Language modeling head
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 
-        # Weight tying
-        self.lm_head.weight = self.wte.weight
+        # Initialize weights
         self.post_init()
 
-    def forward( self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
-        # Default to True if not specified, required for generation
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        if input_ids is None:
+            return None
 
+        batch_size, seq_len = input_ids.shape
         device = input_ids.device
-        b, t = input_ids.size()
-
-        # Positions & Embeddings
-        pos = torch.arange(0, t, dtype=torch.long, device=device)
-        tok_emb = self.wte(input_ids)
-        pos_emb = self.wpe(pos)
-        hidden_states = self.drop(tok_emb + pos_emb)
-
-        # Attention Mask Handling
-        if attention_mask is None:
-            attention_mask = torch.ones((b, t), device=device)
-
-        # Broadcast mask to (batch, head, seq, seq)
-        # We preserve the original mask for the loss calculation later if needed, 
-        # but for the blocks we need the 4D version.
-        extended_attention_mask = attention_mask.view(b, 1, 1, t)
-        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
-
-        # =========================================================
-        # THE RECURSIVE LOOP 
-        # =========================================================
-        for loop_i in range(self.config.n_loops):
-            for block in self.physical_blocks:
-                hidden_states = block(hidden_states, attention_mask=extended_attention_mask)
 
+        # Get embeddings
+        token_embeds = self.wte(input_ids)
+        pos_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
+        pos_embeds = self.wpe(pos_ids)
+        hidden_states = self.drop(token_embeds + pos_embeds)
+
+        # Apply recursive loops through physical blocks
+        for loop in range(self.config.n_loops):
+            block_idx = loop % self.config.n_physical_layers
+            block = self.physical_blocks[block_idx]
+
+            # Attention
+            ln_output = block["ln_1"](hidden_states)
+            attn_output = block["attn"](ln_output, attention_mask=attention_mask)[0]
+            hidden_states = hidden_states + attn_output
+
+            # MLP
+            ln_output = block["ln_2"](hidden_states)
+            mlp_output = block["mlp"](ln_output)
+            hidden_states = hidden_states + mlp_output
+
+        # Final layer norm and projection
         hidden_states = self.ln_f(hidden_states)
         logits = self.lm_head(hidden_states)
 
         loss = None
         if labels is not None:
-            loss_fct = nn.CrossEntropyLoss()
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
 
-        # <--- CRITICAL FIX: Return CausalLMOutputWithCrossAttentions
-        if not return_dict:
-            output = (logits,)
-            return ((loss,) + output) if loss is not None else output
-
         return CausalLMOutputWithCrossAttentions(
             loss=loss,
             logits=logits,
-            past_key_values=None, # We are not using KV-cache for simplicity in this recursive setup
-            hidden_states=None,
+            hidden_states=hidden_states,
             attentions=None,
+            cross_attentions=None
         )
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         return {"input_ids": input_ids}
 
-class RecursiveBlock(nn.Module):
-    def __init__(self, config, layer_idx):
-        super().__init__()
-        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.attn = GPT2Attention(config, layer_idx=layer_idx)
-        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.mlp = GPT2MLP(config.n_embd, config)
-
-    def forward(self, x, layer_past=None, attention_mask=None):
-        residual = x
-        x = self.ln_1(x)
-        # We disable caching (use_cache=False) to simplify the recursion loop
-        attn_outputs = self.attn(x, layer_past=layer_past, attention_mask=attention_mask, use_cache=False)
-        x = residual + attn_outputs[0]
-
-        residual = x
-        x = self.ln_2(x)
-        x = residual + self.mlp(x)
-        return x
+    def _reorder_cache(self, past, beam_idx):
+        return past