""" SimpleLLM - Pure Attention-based Language Model with DeepSeek MLA + RoPE. Architecture: - Token Embedding → Attention Blocks → Output Head - Attention Blocks: Multi-head Latent Attention with RoPE positional embeddings - Feed-forward: SwiGLU gates - No state-space models (SSM), pure transformer architecture - Full BF16 precision (no quantization) """ import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from taoTrain.core import BaseModel from taoTrain.config import ModelConfig from .registry import register_architecture from .mla_components import AttentionBlock from .embeddings import FactorizedEmbedding @register_architecture("taonet") class SimpleLLM(BaseModel): """ Pure attention-based language model with DeepSeek MLA + RoPE. Stateless architecture - no internal state management needed. Args: config: ModelConfig with: - vocab_size: Vocabulary size - hidden_dim: Model dimension (d_model) - hidden_dim_ff: Feed-forward dimension (default: 4 * hidden_dim) - num_layers: Number of attention blocks (n_layers) - num_heads: Number of attention heads (n_attn_heads) - d_latent_kv: KV compression dimension (default: 3/4 * hidden_dim) - d_rope: RoPE dimension per head (default: hidden_dim // num_heads) - max_seq_length: Maximum sequence length - dropout: Dropout rate - gqa_groups: Grouped Query Attention groups (default: 1) - use_factorized_embedding: Use low-rank embedding (default: False) """ def __init__(self, config: ModelConfig): super().__init__(config) # Parse config - use defaults if not specified self.vocab_size = config.vocab_size self.d_model = config.hidden_dim self.n_layers = config.num_layers self.n_heads = config.num_heads self.dropout = config.dropout # Optional parameters with smart defaults self.d_latent_kv = config.d_latent_kv if config.d_latent_kv is not None else int(self.d_model * 0.75) self.d_rope = config.d_rope if config.d_rope is not None else (self.d_model // self.n_heads) self.d_ff = config.hidden_dim_ff if config.hidden_dim_ff is not None else (self.d_model * 4) self.gqa_groups = getattr(config, 'gqa_groups', 1) self.use_factorized_embedding = getattr(config, 'use_factorized_embedding', False) self.d_embed_rank = getattr(config, 'd_embed_rank', 96) # YaRN parameters for context length extension self.rope_scale = getattr(config, 'rope_scale', 40.0) self.yarn_enabled = getattr(config, 'yarn_enabled', False) self.yarn_original_max_seq_length = getattr(config, 'yarn_original_max_seq_length', None) self.yarn_alpha = getattr(config, 'yarn_alpha', 1.0) self.max_seq_length = config.max_seq_length # Validate dimensions assert self.d_model % self.n_heads == 0, \ f"hidden_dim ({self.d_model}) must be divisible by num_heads ({self.n_heads})" assert self.d_latent_kv % self.n_heads == 0, \ f"d_latent_kv ({self.d_latent_kv}) must be divisible by num_heads ({self.n_heads})" # Token embedding if self.use_factorized_embedding: self.token_embedding = FactorizedEmbedding( self.vocab_size, self.d_model, self.d_embed_rank ) else: self.token_embedding = nn.Embedding(self.vocab_size, self.d_model) # Embedding dropout self.embedding_dropout = nn.Dropout(self.dropout) # Attention blocks with MLA + SwiGLU FFN self.blocks = nn.ModuleList() for _ in range(self.n_layers): self.blocks.append( AttentionBlock( d_model=self.d_model, d_latent_kv=self.d_latent_kv, n_heads=self.n_heads, d_rope=self.d_rope, d_ff=int(self.d_ff), dropout=self.dropout, gqa_groups=self.gqa_groups, rope_scale=self.rope_scale, max_seq_length=self.max_seq_length, yarn_enabled=self.yarn_enabled, yarn_original_max_seq_length=self.yarn_original_max_seq_length, yarn_alpha=self.yarn_alpha, ) ) # Final layer norm self.final_norm = nn.LayerNorm(self.d_model) # Output projection to vocabulary self.output_head = nn.Linear(self.d_model, self.vocab_size, bias=False) # Initialize weights self.apply(self._init_weights) # Cache for causal mask self.register_buffer("causal_mask_cache", None, persistent=False) self._print_architecture() def _init_weights(self, module): """Initialize weights for stable training.""" if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def _print_architecture(self): """Print model architecture summary.""" total_params = sum(p.numel() for p in self.parameters()) trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad) print(f"\n{'='*70}") print("MODEL ARCHITECTURE - TAONET (DeepSeek MLA + RoPE)") print(f"{'='*70}") print(f"Embedding:") if self.use_factorized_embedding: embed_rank_params = self.vocab_size * self.d_embed_rank embed_proj_params = self.d_embed_rank * self.d_model print(f" Type: Factorized (rank={self.d_embed_rank})") print(f" Rank layer: {embed_rank_params/1e6:>8.2f}M") print(f" Projection: {embed_proj_params/1e6:>8.2f}M") else: embed_params = self.vocab_size * self.d_model print(f" Type: Standard") print(f" Params: {embed_params/1e6:>8.2f}M") output_params = self.d_model * self.vocab_size print(f"Output Head: {output_params/1e6:>8.2f}M") print(f"Attention Blocks: {len(self.blocks):>10} layers x AttentionBlock") print(f"{'-'*70}") print(f"Total Parameters: {total_params/1e6:>8.2f}M (trainable: {trainable_params/1e6:.2f}M)") print(f"{'-'*70}") print(f"Configuration:") print(f" Model dimension (d_model): {self.d_model}") print(f" KV latent dimension (d_latent_kv): {self.d_latent_kv}") print(f" Attention heads: {self.n_heads}") print(f" Head dimension: {self.d_model // self.n_heads}") print(f" RoPE dimension: {self.d_rope}") print(f" Feed-forward dimension: {int(self.d_ff)}") print(f" Number of layers: {self.n_layers}") print(f" Max sequence length: {self.config.max_seq_length}") print(f" Dropout: {self.dropout}") print(f" GQA groups: {self.gqa_groups}") print(f"{'='*70}\n") def _get_causal_mask(self, seq_len, device): """Get or create causal mask for sequence.""" if self.causal_mask_cache is None or self.causal_mask_cache.size(-1) < seq_len: # [seq_len, seq_len] lower triangular matrix (1 = attend, 0 = mask) mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool)) self.register_buffer("causal_mask_cache", mask, persistent=False) return self.causal_mask_cache[:seq_len, :seq_len] def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, ) -> dict: """ Forward pass through the model. Args: input_ids: [batch_size, seq_len] tensor of token IDs attention_mask: [batch_size, seq_len] tensor where 1 = valid, 0 = padding labels: [batch_size, seq_len] target token IDs for loss computation Returns: Dictionary with: - 'logits': [batch_size, seq_len, vocab_size] output logits - 'loss': scalar loss (if labels provided, else None) """ if inputs_embeds is None: if input_ids is None: raise ValueError("Either input_ids or inputs_embeds must be provided") batch_size, seq_len = input_ids.shape device = input_ids.device x = self.token_embedding(input_ids) else: batch_size, seq_len, _ = inputs_embeds.shape device = inputs_embeds.device x = inputs_embeds # Get causal mask: [seq_len, seq_len] causal_mask = self._get_causal_mask(seq_len, device) # Combine causal mask with attention mask if provided if attention_mask is not None: # attention_mask: [batch, seq_len] where 1 = valid, 0 = padding # Expand to [batch, 1, 1, seq_len] padding_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool() # Combine with causal: [1, 1, seq_len, seq_len] * [batch, 1, 1, seq_len] combined_mask = causal_mask.unsqueeze(0).unsqueeze(0) & padding_mask # For MLA: convert to {0, 1} format combined_mask = combined_mask.float() else: # Just causal mask combined_mask = causal_mask.unsqueeze(0).unsqueeze(0).float() x = self.embedding_dropout(x) # Pass through attention blocks for block in self.blocks: x = block(x, attention_mask=combined_mask) # Final layer norm x = self.final_norm(x) # Output projection to vocabulary logits = self.output_head(x) # [batch_size, seq_len, vocab_size] # Compute loss if labels are provided loss = None if labels is not None: # Flatten for loss computation logits_flat = logits.view(-1, logits.size(-1)) # (batch * seq_len, vocab_size) labels_flat = labels.view(-1) valid_label_mask = labels_flat != -100 if not torch.any(valid_label_mask): raise ValueError( "All labels are masked out (-100), so loss cannot be computed. " "This usually indicates a dataset parsing or masking bug." ) # Only compute loss on valid targets (ignore -100 tokens for padding) loss = F.cross_entropy( logits_flat, labels_flat, reduction='mean', ignore_index=-100 ) return { 'logits': logits, 'loss': loss, }