Create model.py

model.py

@@ -0,0 +1,273 @@
+"""
+Nano Reasoning Model (NRM) - Main Architecture
+
+ARCHITECTURE DESIGN PHILOSOPHY:
+================================
+This model maximizes reasoning ability per parameter through several key innovations:
+
+1. SHARED LAYERS: The middle layers are shared (looped through multiple times).
+   This creates a form of "iterative refinement" - the model processes information
+   multiple passes, similar to how recurrent networks process sequences but applied
+   to depth instead. This is inspired by Universal Transformers and ALBERT.
+   
+   WHY IT HELPS REASONING: Reasoning often requires iterative refinement of
+   intermediate representations. Shared layers let the model "think more" without
+   more parameters.
+
+2. THINKING TOKENS: Special <THINK> and </THINK> tokens create a "scratchpad"
+   where the model can show intermediate reasoning steps. The model is trained to
+   use <STEP> tokens for each logical step.
+   
+   WHY IT HELPS: Decomposing complex problems into steps is THE key capability
+   for reasoning. Even large models benefit from chain-of-thought prompting.
+
+3. WEIGHT TYING: Input and output embeddings share the same weight matrix.
+   This halves the embedding parameter count and creates a natural link between
+   token understanding and token generation.
+   
+   WHY IT HELPS CPU: Fewer parameters = faster forward/backward passes.
+
+4. LOW-RANK PROJECTIONS: All attention and MLP projections use LoRA-style
+   factored matrices, cutting parameter count by ~8x in linear layers.
+
+5. GROUPED QUERY ATTENTION: KV heads are shared across query heads,
+   reducing KV projection parameters and memory.
+
+PARAMETER BUDGET (~10M):
+  Embedding: 2048 * 256 = 524K (shared with output head)
+  Per unique layer: ~200K
+  4 unique + 2 shared (run 2x) = 6 effective layers
+  Total: ~2.1M (layers) + 524K (embed) ≈ 2.6M unique params
+  Effective computation: ~3.1M param equivalent
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Dict
+from .components import TransformerBlock, RMSNorm
+
+
+class NanoReasoningModel(nn.Module):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        
+        d_model = config['d_model']
+        n_heads = config['n_heads']
+        n_layers = config['n_layers']
+        n_shared = config.get('n_shared_layers', 2)
+        d_ff = config['d_ff']
+        vocab_size = config['vocab_size']
+        max_seq_len = config['max_seq_len']
+        dropout = config.get('dropout', 0.05)
+        rank = config.get('lora_rank', 16)
+        self.use_thinking = config.get('use_thinking_tokens', True)
+        self.n_thinking_steps = config.get('n_thinking_steps', 2)
+        n_kv_heads = config.get('n_kv_heads', n_heads // 2)
+        
+        # Token embeddings (will be tied with output head)
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.embedding_dropout = nn.Dropout(dropout)
+        
+        # Entry layers (unique)
+        n_unique = n_layers - n_shared
+        self.entry_layers = nn.ModuleList([
+            TransformerBlock(d_model, n_heads, d_ff, rank, dropout, max_seq_len, n_kv_heads)
+            for _ in range(n_unique // 2)
+        ])
+        
+        # Shared layers (looped)
+        self.shared_layers = nn.ModuleList([
+            TransformerBlock(d_model, n_heads, d_ff, rank, dropout, max_seq_len, n_kv_heads)
+            for _ in range(n_shared)
+        ])
+        
+        # Exit layers (unique)
+        self.exit_layers = nn.ModuleList([
+            TransformerBlock(d_model, n_heads, d_ff, rank, dropout, max_seq_len, n_kv_heads)
+            for _ in range(n_unique - n_unique // 2)
+        ])
+        
+        # Final norm
+        self.final_norm = RMSNorm(d_model)
+        
+        # Output head (tied with embeddings)
+        self.output_head = nn.Linear(d_model, vocab_size, bias=False)
+        
+        if config.get('weight_tying', True):
+            self.output_head.weight = self.token_embedding.weight
+        
+        # Thinking step gate: learned scalar for blending thinking iterations
+        if self.use_thinking:
+            self.think_gate = nn.Parameter(torch.tensor(0.5))
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+        
+        # Count parameters
+        self._count_parameters()
+
+    def _init_weights(self, module: nn.Module):
+        """Initialize weights with scaled initialization for stability."""
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _count_parameters(self):
+        """Count and report parameters."""
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        
+        # Count unique parameters (shared layers counted once)
+        unique = sum(p.numel() for p in self.parameters())
+        
+        self.total_params = total
+        self.trainable_params = trainable
+        print(f"\n{'='*50}")
+        print(f"NRM Model Configuration:")
+        print(f"  d_model: {self.config['d_model']}")
+        print(f"  n_heads: {self.config['n_heads']}")
+        print(f"  n_layers: {self.config['n_layers']} "
+              f"({len(self.entry_layers)} entry + {len(self.shared_layers)} shared + {len(self.exit_layers)} exit)")
+        print(f"  d_ff: {self.config['d_ff']}")
+        print(f"  vocab_size: {self.config['vocab_size']}")
+        print(f"  LoRA rank: {self.config.get('lora_rank', 16)}")
+        print(f"  Thinking: {'enabled' if self.use_thinking else 'disabled'}")
+        print(f"  Total parameters: {total:,}")
+        print(f"  Trainable parameters: {trainable:,}")
+        print(f"{'='*50}\n")
+
+    def forward(self, input_ids: torch.Tensor, 
+                attention_mask: Optional[torch.Tensor] = None,
+                labels: Optional[torch.Tensor] = None,
+                n_think_loops: int = 1) -> Dict[str, torch.Tensor]:
+        """
+        Forward pass with optional thinking loops.
+        
+        n_think_loops: How many times to loop through shared layers.
+        During reasoning, we increase this to give the model more "thinking time".
+        """
+        B, T = input_ids.shape
+        
+        # Embeddings
+        x = self.token_embedding(input_ids)
+        x = self.embedding_dropout(x)
+        
+        # Padding mask
+        pad_mask = None
+        if attention_mask is not None:
+            pad_mask = (attention_mask == 0)  # True where padded
+        
+        # Entry layers
+        for layer in self.entry_layers:
+            x = layer(x, pad_mask)
+        
+        # Shared layers with thinking loops
+        actual_loops = max(1, n_think_loops)
+        if self.use_thinking and actual_loops > 1:
+            # Store the "pre-think" state
+            x_original = x
+            for loop in range(actual_loops):
+                for layer in self.shared_layers:
+                    x = layer(x, pad_mask)
+                if loop < actual_loops - 1:
+                    # Blend with original (residual thinking)
+                    gate = torch.sigmoid(self.think_gate)
+                    x = gate * x + (1 - gate) * x_original
+        else:
+            for layer in self.shared_layers:
+                x = layer(x, pad_mask)
+        
+        # Exit layers
+        for layer in self.exit_layers:
+            x = layer(x, pad_mask)
+        
+        # Output
+        x = self.final_norm(x)
+        logits = self.output_head(x)
+        
+        result = {"logits": logits}
+        
+        if labels is not None:
+            # Shift for autoregressive loss
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=0,  # PAD token
+                label_smoothing=0.05  # Slight smoothing for better generalization
+            )
+            result["loss"] = loss
+        
+        return result
+
+    @torch.no_grad()
+    def generate(self, input_ids: torch.Tensor, max_new_tokens: int = 100,
+                 temperature: float = 0.7, top_k: int = 50, top_p: float = 0.9,
+                 n_think_loops: int = 1, eos_token_id: int = 2) -> torch.Tensor:
+        """
+        Autoregressive generation with temperature, top-k, and top-p sampling.
+        
+        Uses nucleus (top-p) sampling for diverse but coherent generation.
+        """
+        self.eval()
+        generated = input_ids.clone()
+        
+        for _ in range(max_new_tokens):
+            # Truncate to max_seq_len
+            context = generated[:, -self.config['max_seq_len']:]
+            
+            outputs = self.forward(context, n_think_loops=n_think_loops)
+            logits = outputs["logits"][:, -1, :] / max(temperature, 1e-5)
+            
+            # Top-k filtering
+            if top_k > 0:
+                top_k_val = min(top_k, logits.size(-1))
+                indices_to_remove = logits < torch.topk(logits, top_k_val)[0][..., -1, None]
+                logits[indices_to_remove] = float('-inf')
+            
+            # Top-p (nucleus) filtering
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                indices_to_remove = sorted_indices_to_remove.scatter(
+                    1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = float('-inf')
+            
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            generated = torch.cat([generated, next_token], dim=1)
+            
+            if next_token.item() == eos_token_id:
+                break
+        
+        return generated
+
+    def save(self, path: str):
+        """Save model state dict and config."""
+        import os, json
+        os.makedirs(path, exist_ok=True)
+        torch.save(self.state_dict(), os.path.join(path, "model.pt"))
+        with open(os.path.join(path, "config.json"), 'w') as f:
+            json.dump(self.config, f, indent=2)
+        print(f"Model saved to {path}")
+
+    @classmethod
+    def load(cls, path: str, device: str = 'cpu') -> 'NanoReasoningModel':
+        """Load model from saved state."""
+        import os, json
+        with open(os.path.join(path, "config.json"), 'r') as f:
+            config = json.load(f)
+        model = cls(config)
+        model.load_state_dict(torch.load(os.path.join(path, "model.pt"), 
+                                          map_location=device))
+        return model