Update modeling file with complete recursive implementation

modeling_tiny_recursive.py

@@ -17,8 +17,15 @@ class TRMConfig(PretrainedConfig):
         n_physical_layers=3,
         n_loops=8,
         n_head=8,
+        activation_function="gelu_new",
+        resid_pdrop=0.1,
         embd_pdrop=0.1,
-        **kwargs
+        attn_pdrop=0.1,
+        layer_norm_epsilon=1e-5,
+        scale_attn_weights=True,
+        scale_attn_by_inverse_layer_idx=False,
+        reorder_and_upcast_attn=False,
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.vocab_size = vocab_size
@@ -27,7 +34,14 @@ class TRMConfig(PretrainedConfig):
         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
+        self.attn_pdrop = attn_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.scale_attn_weights = scale_attn_weights
+        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
         self.hidden_size = n_embd
@@ -48,17 +62,66 @@ 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 - Add your recursive layers here
-        # [Your recursive implementation from the notebook]
+        # 2. The Logic Core - Physical transformer blocks
+        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)
+        ])
 
-        # 3. Language modeling head
+        # 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
         self.post_init()
 
     def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
-        # Add your forward pass implementation
-        pass
+        batch_size, seq_len = input_ids.shape
+
+        # 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)
+
+        # 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:
+            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))
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None
+        )
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         return {"input_ids": input_ids}