| """Full dense decoder-only transformer model for SAGE.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Optional |
|
|
| import torch |
| from torch import nn |
|
|
| from model.block import TransformerBlock |
| from model.config import ModelConfig |
| from model.rope import build_rope_cache |
| from model.rmsnorm import RMSNorm |
|
|
|
|
| class SageTransformer(nn.Module): |
| """A dense Llama-style decoder-only transformer.""" |
|
|
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.config = config |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model) |
| self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)]) |
| self.norm = RMSNorm(config.d_model, eps=config.rms_norm_eps) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| if config.tie_word_embeddings: |
| self.lm_head.weight = self.embed_tokens.weight |
| cos, sin = build_rope_cache( |
| seq_len=config.context_length, |
| head_dim=config.head_dim, |
| base_frequency=config.rope_base_frequency, |
| scaling_factor=config.rope_scaling_factor, |
| ) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
| self._reset_parameters() |
|
|
| def _reset_parameters(self) -> None: |
| """Apply scaled initialization to the model.""" |
| embed_std = 1.0 / math.sqrt(self.config.d_model) |
| nn.init.normal_(self.embed_tokens.weight, mean=0.0, std=embed_std) |
| for module in self.modules(): |
| if not isinstance(module, nn.Linear): |
| continue |
| std = self.config.initializer_range |
| if module is self.lm_head and self.config.tie_word_embeddings: |
| continue |
| if module.out_features == self.config.d_model: |
| std = std / math.sqrt(2 * self.config.num_layers) |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| past_key_values: Optional[list[tuple[torch.Tensor, torch.Tensor]]] = None, |
| ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]]]: |
| """Return logits and the updated KV cache.""" |
| batch_size, seq_len = input_ids.shape |
| hidden_states = self.embed_tokens(input_ids) |
| past_key_values = past_key_values or [None] * self.config.num_layers |
| start = 0 |
| if past_key_values[0] is not None: |
| start = past_key_values[0][0].size(-2) |
| cos = self.rope_cos[start : start + seq_len].to(hidden_states.device) |
| sin = self.rope_sin[start : start + seq_len].to(hidden_states.device) |
| presents: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| for layer, past in zip(self.layers, past_key_values): |
| hidden_states, present = layer(hidden_states, cos, sin, past_key_value=past) |
| presents.append(present) |
| hidden_states = self.norm(hidden_states) |
| logits = self.lm_head(hidden_states) |
| return logits, presents |
|
|
