| from typing import Optional |
| """ |
| Decoder-only Transformer language model. |
| |
| Training objective: |
| Next-token prediction. |
| |
| Input: |
| input_ids: [batch, seq] |
| |
| Output: |
| logits: [batch, seq, vocab] |
| loss: optional cross-entropy loss if targets are given. |
| """ |
|
|
| import math |
| from dataclasses import asdict |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from src.model.config import ModelConfig |
| from src.model.layers import RMSNorm, TransformerBlock |
|
|
|
|
| class DecoderOnlyTransformer(nn.Module): |
| """ |
| GPT/LLaMA-style decoder-only language model. |
| """ |
|
|
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.config = config |
|
|
| self.tok_embeddings = nn.Embedding(config.vocab_size, config.n_embd) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| self.blocks = nn.ModuleList( |
| [TransformerBlock(config) for _ in range(config.n_layer)] |
| ) |
|
|
| self.norm = RMSNorm(config.n_embd, config.norm_eps) |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
|
|
| self.apply(self._init_weights) |
|
|
| |
| self.lm_head.weight = self.tok_embeddings.weight |
|
|
| |
| for name, param in self.named_parameters(): |
| if name.endswith("out_proj.weight") or name.endswith("down_proj.weight"): |
| nn.init.normal_( |
| param, |
| mean=0.0, |
| std=0.02 / math.sqrt(2 * config.n_layer), |
| ) |
|
|
| def _init_weights(self, module: nn.Module) -> None: |
| """ |
| Initializes model weights. |
| """ |
| 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 forward( |
| self, |
| input_ids: torch.Tensor, |
| targets: Optional[torch.Tensor] = None, |
| ): |
| """ |
| Input: |
| input_ids: [batch, seq] |
| targets: [batch, seq], optional |
| |
| Output: |
| logits: [batch, seq, vocab] |
| loss: scalar tensor if targets are provided. |
| """ |
| batch, seq_len = input_ids.shape |
|
|
| if seq_len > self.config.block_size: |
| raise ValueError( |
| f"Sequence length {seq_len} exceeds block_size {self.config.block_size}" |
| ) |
|
|
| x = self.tok_embeddings(input_ids) |
| x = self.dropout(x) |
|
|
| for block in self.blocks: |
| x = block(x) |
|
|
| x = self.norm(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
|
|
| if targets is not None: |
| loss = F.cross_entropy( |
| logits.reshape(-1, logits.size(-1)), |
| targets.reshape(-1), |
| ) |
|
|
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| input_ids: torch.Tensor, |
| max_new_tokens: int, |
| temperature: float = 1.0, |
| top_k: Optional[int] = None, |
| eos_token_id: Optional[int] = None, |
| ) -> torch.Tensor: |
| """ |
| Input: |
| input_ids: [batch, seq] |
| max_new_tokens: number of tokens to generate. |
| temperature: sampling temperature. |
| top_k: keep only top-k logits if provided. |
| eos_token_id: stop generation if this token is produced. |
| |
| Output: |
| Generated token IDs: [batch, seq + max_new_tokens] |
| """ |
| self.eval() |
|
|
| |
| |
| |
| batch = input_ids.size(0) |
| finished = torch.zeros(batch, dtype=torch.bool, device=input_ids.device) |
|
|
| for _ in range(max_new_tokens): |
| context = input_ids[:, -self.config.block_size:] |
|
|
| logits, _ = self(context) |
| logits = logits[:, -1, :] / temperature |
|
|
| if top_k is not None: |
| values, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < values[:, [-1]]] = -float("inf") |
|
|
| probs = F.softmax(logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
|
|
| input_ids = torch.cat([input_ids, next_token], dim=1) |
|
|
| if eos_token_id is not None: |
| finished = finished | (next_token.squeeze(-1) == eos_token_id) |
| if bool(finished.all()): |
| break |
|
|
| return input_ids |
|
|
| def get_num_params(self) -> int: |
| """ |
| Output: |
| Number of trainable model parameters. |
| """ |
| return sum(p.numel() for p in self.parameters() if p.requires_grad) |
|
|
| def config_dict(self) -> dict: |
| return asdict(self.config) |
|
|
