Update modeling file with complete recursive implementation

modeling_tiny_recursive.py

@@ -1,11 +1,11 @@
 
-from transformers import PreTrainedModel, PretrainedConfig
-from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
-from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
 import torch
 import torch.nn as nn
-
+from transformers import PreTrainedModel, PretrainedConfig, GPT2TokenizerFast, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
+from transformers.generation import GenerationMixin  # <--- FIXED: Import this explicitly
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from datasets import load_dataset
 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,12 +28,13 @@ 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
@@ -43,15 +44,17 @@ 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
 
-        # Required for transformers compatibility
+        # --- CRITICAL FIXES FOR COMPATIBILITY ---
+        # These map your custom names to what GPT2Attention expects
+        self.max_position_embeddings = n_positions 
         self.hidden_size = n_embd
-        self.num_attention_heads = n_head
+        self.num_attention_heads = n_head          # <--- FIXED: The missing attribute
         self.num_hidden_layers = n_physical_layers
-        self.n_inner = None
+        self.n_inner = None                        # Defaults to 4*hidden_size
 
 class TinyRecursiveModel(PreTrainedModel, GenerationMixin):
     config_class = TRMConfig
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = ["lm_head.weight"]  # <-- Add this line
 
     def __init__(self, config):
         super().__init__(config)
@@ -62,66 +65,90 @@ 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 - Physical transformer blocks
+        # 2. The Logic Core (The "7M" part)
         self.physical_blocks = nn.ModuleList([
-            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)
+            RecursiveBlock(config, layer_idx=i) 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)
 
-        # Initialize weights
+        # Weight tying
+        self.lm_head.weight = self.wte.weight
         self.post_init()
 
-    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
-        batch_size, seq_len = input_ids.shape
+    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
+
+        device = input_ids.device
+        b, t = input_ids.size()
 
-        # Get embeddings
-        token_embeds = self.wte(input_ids)
-        pos_ids = torch.arange(0, seq_len, dtype=torch.long, device=input_ids.device)
-        pos_embeds = self.wpe(pos_ids)
-        hidden_states = self.drop(token_embeds + pos_embeds)
+        # 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)
 
-        # 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 Mask Handling
+        if attention_mask is None:
+            attention_mask = torch.ones((b, t), device=device)
 
-            # 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
+        # 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
 
-            # MLP
-            ln_output = block["ln_2"](hidden_states)
-            mlp_output = block["mlp"](ln_output)
-            hidden_states = hidden_states + mlp_output
+        # =========================================================
+        # 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)
 
-        # 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,
-            attentions=None
+            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