""" AresForCausalLM - Decoder-Only Transformer, from scratch. No external APIs for core intelligence. Includes: - Token embeddings (LOTS) - RoPE - MHA+GQA+KV Cache - RMSNorm - SwiGLU/GELU - Unembedding matrix multipliers - Token outputs / logits """ import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional, List, Tuple from ..config import AresConfig from .rmsnorm import RMSNorm from .transformer_block import AresDecoderLayer from .attention import KVCache class AresModel(nn.Module): def __init__(self, config: AresConfig): super().__init__() self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) # token embedding matrix self.layers = nn.ModuleList([AresDecoderLayer(config) for _ in range(config.num_hidden_layers)]) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = False def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, kv_caches=None): hidden_states = self.embed_tokens(input_ids) # [b,s,hidden] presents = [] if use_cache else None for i, layer in enumerate(self.layers): past_kv = past_key_values[i] if past_key_values is not None else None cache = kv_caches[i] if kv_caches is not None else None if self.gradient_checkpointing and self.training: # Flaw fix: memory blow -> gradient checkpointing def custom_forward(h): return layer(h, attention_mask=attention_mask, past_kv=past_kv, use_cache=use_cache, cache=cache)[0] hidden_states = torch.utils.checkpoint.checkpoint(custom_forward, hidden_states) # present not tracked in checkpoint mode for simplicity if use_cache: # dummy pass else: hidden_states, present = layer( hidden_states, attention_mask=attention_mask, past_kv=past_kv, use_cache=use_cache, cache=cache ) if use_cache: presents.append(present) hidden_states = self.norm(hidden_states) return hidden_states, presents class AresForCausalLM(nn.Module): def __init__(self, config: AresConfig): super().__init__() self.config = config self.model = AresModel(config) # Unembedding matrix multiplier self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # output projection if config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) 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=self.config.initializer_range) def forward( self, input_ids: torch.Tensor, labels: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, use_cache: bool = False, kv_caches: Optional[List[KVCache]] = None, ): hidden_states, presents = self.model( input_ids, attention_mask=attention_mask, past_key_values=past_key_values, use_cache=use_cache, kv_caches=kv_caches ) logits = self.lm_head(hidden_states) # token sections -> token outputs [b,s,vocab] loss = None if labels is not None: # Shift for causal LM shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = nn.CrossEntropyLoss() loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)) return {"logits": logits, "loss": loss, "past_key_values": presents, "hidden_states": hidden_states} @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 100, temperature: float = 0.8, top_p: float = 0.9, top_k: int = 50, do_sample: bool = True, eos_token_id: Optional[int] = None, ): self.eval() bsz, seq_len = input_ids.shape device = input_ids.device # Initialize KV caches for fast inference kv_caches = [ KVCache(bsz, self.config.num_key_value_heads, self.config.max_position_embeddings, self.config.head_dim, device, self.model.embed_tokens.weight.dtype) for _ in range(self.config.num_hidden_layers) ] # Prefill hidden, presents = self.model(input_ids[:, :-1], use_cache=False) if seq_len > 1 else (None, None) # Actually for cache prefill, we need to run first tokens through cache # Simplify: run all but last token to fill cache if seq_len > 1: # Fill cache iteratively - for simplicity use past_key_values not KVCache class here, but we will use cache # Prefill using kv_caches # First token cur_ids = input_ids[:, :1] for i in range(1, seq_len): _ = self.model(cur_ids[:, -1:], past_key_values=None, kv_caches=kv_caches if i>1 else None, use_cache=False) # Actually our model forward uses kv_caches list - let's call layer-wise? easier: just loop model forward for each token with kv_caches # We need to handle cache update inside attention, so we call model forward each step # Reset approach: use past_key_values accumulation instead of KVCache class for generate simplicity pass # For simplicity, abandon KVCache class in generate and use past_key_values past = None # Re-prefill correctly with past_key_values if seq_len > 1: out = self.forward(input_ids[:, :-1], use_cache=True) past = out["past_key_values"] logits = out["logits"][:, -1:, :] else: out = self.forward(input_ids, use_cache=True) past = out["past_key_values"] logits = out["logits"][:, -1:, :] else: out = self.forward(input_ids, use_cache=True) past = out["past_key_values"] logits = out["logits"][:, -1:, :] generated = input_ids for _ in range(max_new_tokens): # Last token logits last_logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0) if do_sample: if top_k > 0: top_k_vals, _ = torch.topk(last_logits, top_k) min_top_k = top_k_vals[:, -1].unsqueeze(-1) last_logits = torch.where(last_logits < min_top_k, torch.full_like(last_logits, float('-inf')), last_logits) if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(last_logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_mask = cumulative_probs > top_p sorted_mask[:, 1:] = sorted_mask[:, :-1].clone() sorted_mask[:, 0] = 0 indices_to_remove = sorted_mask.scatter(1, sorted_indices, sorted_mask) last_logits = last_logits.masked_fill(indices_to_remove, float('-inf')) probs = F.softmax(last_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(last_logits, dim=-1, keepdim=True) generated = torch.cat([generated, next_token], dim=1) if eos_token_id is not None and (next_token == eos_token_id).all(): break # next forward with cache out = self.forward(next_token, past_key_values=past, use_cache=True) past = out["past_key_values"] logits = out["logits"] if generated.shape[1] >= self.config.max_position_embeddings: break return generated def count_parameters(self): return sum(p.numel() for p in self.parameters())