| """Hugging Face-compatible wrapper for the from-scratch decoder transformer. |
| |
| This file is self-contained (no imports from the rest of this repo) because |
| Hugging Face's `trust_remote_code=True` loading only pulls in files that live |
| in the model's own HF repo, not the GitHub repo. The actual architecture is |
| identical to src/model/transformer.py, duplicated here so this file can be |
| uploaded on its own alongside config.json and the weights. |
| |
| Registered via config.json's "auto_map" so it loads with: |
| AutoModelForCausalLM.from_pretrained("<repo>", trust_remote_code=True) |
| """ |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PretrainedConfig, PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutput |
|
|
|
|
| class DecoderTransformerConfig(PretrainedConfig): |
| model_type = "decoder_transformer_scratch" |
|
|
| def __init__( |
| self, |
| vocab_size: int = 16384, |
| n_layer: int = 8, |
| n_head: int = 12, |
| n_embd: int = 768, |
| block_size: int = 1024, |
| dropout: float = 0.0, |
| bias: bool = False, |
| **kwargs, |
| ): |
| self.vocab_size = vocab_size |
| self.n_layer = n_layer |
| self.n_head = n_head |
| self.n_embd = n_embd |
| self.block_size = block_size |
| self.dropout = dropout |
| self.bias = bias |
| kwargs.setdefault("tie_word_embeddings", True) |
| super().__init__(**kwargs) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-5): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return norm * self.weight |
|
|
|
|
| def build_rope_cache(seq_len: int, head_dim: int, base: int = 10000, device=None): |
| inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) |
| t = torch.arange(seq_len, device=device).float() |
| freqs = torch.outer(t, inv_freq) |
| return torch.cos(freqs), torch.sin(freqs) |
|
|
|
|
| def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| x1, x2 = x[..., 0::2], x[..., 1::2] |
| cos = cos[: x.size(-2)].unsqueeze(0).unsqueeze(0) |
| sin = sin[: x.size(-2)].unsqueeze(0).unsqueeze(0) |
| rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) |
| return rotated.flatten(-2) |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, cfg: DecoderTransformerConfig): |
| super().__init__() |
| assert cfg.n_embd % cfg.n_head == 0 |
| self.n_head = cfg.n_head |
| self.head_dim = cfg.n_embd // cfg.n_head |
| self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias) |
| self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias) |
| self.dropout = cfg.dropout |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| B, T, C = x.shape |
| q, k, v = self.qkv(x).split(C, dim=2) |
| q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| q = apply_rope(q, cos, sin) |
| k = apply_rope(k, cos, sin) |
|
|
| y = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0 |
| ) |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.proj(y) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, cfg: DecoderTransformerConfig): |
| super().__init__() |
| hidden = int(8 * cfg.n_embd / 3) |
| hidden = ((hidden + 63) // 64) * 64 |
| self.w1 = nn.Linear(cfg.n_embd, hidden, bias=cfg.bias) |
| self.w2 = nn.Linear(cfg.n_embd, hidden, bias=cfg.bias) |
| self.w3 = nn.Linear(hidden, cfg.n_embd, bias=cfg.bias) |
| self.dropout = nn.Dropout(cfg.dropout) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.dropout(self.w3(F.silu(self.w1(x)) * self.w2(x))) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg: DecoderTransformerConfig): |
| super().__init__() |
| self.ln1 = RMSNorm(cfg.n_embd) |
| self.attn = CausalSelfAttention(cfg) |
| self.ln2 = RMSNorm(cfg.n_embd) |
| self.mlp = SwiGLU(cfg) |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.ln1(x), cos, sin) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
|
|
| class DecoderTransformerCore(nn.Module): |
| """Identical structure/state_dict keys to src/model/transformer.py's |
| DecoderTransformer, so a training checkpoint loads in directly.""" |
|
|
| def __init__(self, cfg: DecoderTransformerConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd) |
| self.drop = nn.Dropout(cfg.dropout) |
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]) |
| self.ln_f = RMSNorm(cfg.n_embd) |
| self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) |
| self.head.weight = self.tok_emb.weight |
|
|
| head_dim = cfg.n_embd // cfg.n_head |
| cos, sin = build_rope_cache(cfg.block_size, head_dim) |
| |
| |
| |
| |
| self.register_buffer("rope_cos", cos, persistent=True) |
| self.register_buffer("rope_sin", sin, persistent=True) |
|
|
| self.apply(self._init_weights) |
| for name, p in self.named_parameters(): |
| if name.endswith("proj.weight") or name.endswith("w3.weight"): |
| nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) |
|
|
| def _init_weights(self, module: nn.Module): |
| 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, idx: torch.Tensor, targets: torch.Tensor | None = None): |
| B, T = idx.shape |
| assert T <= self.cfg.block_size, f"sequence length {T} exceeds block_size {self.cfg.block_size}" |
| x = self.drop(self.tok_emb(idx)) |
| cos, sin = self.rope_cos.to(x.device), self.rope_sin.to(x.device) |
| for block in self.blocks: |
| x = block(x, cos, sin) |
| x = self.ln_f(x) |
| logits = self.head(x) |
|
|
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) |
| return logits, loss |
|
|
|
|
| class DecoderTransformerForCausalLM(PreTrainedModel): |
| """The HF-facing wrapper. Delegates all real computation to `self.model` |
| (DecoderTransformerCore), which keeps this file's forward pass identical |
| to what was actually trained rather than a reimplementation that could |
| drift from it.""" |
|
|
| config_class = DecoderTransformerConfig |
| |
| |
| |
| _tied_weights_keys = {"model.head.weight": "model.tok_emb.weight"} |
|
|
| def __init__(self, config: DecoderTransformerConfig): |
| super().__init__(config) |
| self.model = DecoderTransformerCore(config) |
| self.post_init() |
|
|
| def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None, **kwargs): |
| logits, loss = self.model(input_ids, labels) |
| return CausalLMOutput(loss=loss, logits=logits) |
|
|
| def get_input_embeddings(self): |
| return self.model.tok_emb |
|
|
| def set_input_embeddings(self, value): |
| self.model.tok_emb = value |
|
|
| def get_output_embeddings(self): |
| return self.model.head |
|
|
| def set_output_embeddings(self, new_embeddings): |
| self.model.head = new_embeddings |
|
|
| @torch.no_grad() |
| def generate_simple( |
| self, |
| input_ids: torch.Tensor, |
| max_new_tokens: int = 50, |
| temperature: float = 0.8, |
| top_k: int = 40, |
| repetition_penalty: float = 1.3, |
| ): |
| """Minimal sampling decode loop (temperature + top-k + repetition |
| penalty), recomputing the full sequence each step. This custom |
| architecture doesn't implement HF's KV-cache generation hooks, so |
| this is O(n^2) in sequence length instead of O(n) - fine for short |
| demo completions, not for serving at scale, which is exactly the |
| kind of tradeoff worth being explicit about rather than pretending |
| .generate() here is production-grade. |
| |
| repetition_penalty follows Keskar et al. 2019: logits of tokens |
| already in the sequence are softened (divided if positive, |
| multiplied if negative) before sampling. Without this, a ~57M-param |
| model trained on ~1B tokens readily falls into repetition loops |
| ("voltage-voltage-voltage...") at low temperature, since nothing |
| else in this minimal loop discourages repeating a high-probability |
| token indefinitely.""" |
| self.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = input_ids[:, -self.config.block_size :] |
| logits, _ = self.model(idx_cond) |
| logits = logits[:, -1, :] / max(temperature, 1e-5) |
| if repetition_penalty and repetition_penalty != 1.0: |
| for seen_id in set(input_ids[0].tolist()): |
| logit = logits[0, seen_id] |
| logits[0, seen_id] = logit / repetition_penalty if logit > 0 else logit * repetition_penalty |
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float("inf") |
| probs = F.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probs, num_samples=1) |
| input_ids = torch.cat([input_ids, next_id], dim=1) |
| return input_ids |
|
|
