File size: 3,501 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import torch
import torch.nn as nn
from .config import ARConfig
from .rmsnorm import RMSNorm
from .block import DecoderBlock
from .rope import build_rope_cache
class ARModel(nn.Module):
"""Decoder-only transformer: token IDs in, next-token scores out."""
def __init__(self, cfg: ARConfig):
super().__init__()
self.cfg = cfg
# token embedding: maps each token ID to a dim-wide vector
self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
# the stack of identical decoder blocks
self.blocks = nn.ModuleList([
DecoderBlock(cfg.dim, cfg.n_heads, cfg.n_kv_heads,
cfg.head_dim, cfg.hidden_dim, cfg.rms_norm_eps)
for _ in range(cfg.n_layers)
])
self.final_norm = RMSNorm(cfg.dim, cfg.rms_norm_eps)
# output head: dim-wide vector -> score for every token in the vocab
self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
# tie weights: embedding table and output head share one matrix
if cfg.tie_embeddings:
self.lm_head.weight = self.embed.weight
# precompute RoPE tables once; reused by every block (not a parameter)
cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
self.register_buffer("rope_cos", cos, persistent=False)
self.register_buffer("rope_sin", sin, persistent=False)
# initialize all weights small; without this, 24 stacked layers make
# the output logits explode (std ~32 instead of ~1)
self.apply(self._init_weights)
# extra: scale down the two residual-path projections per block by
# 1/sqrt(2*n_layers) so the signal doesn't compound across depth
scale = (2 * cfg.n_layers) ** -0.5
for block in self.blocks:
torch.nn.init.normal_(block.attn.o_proj.weight, mean=0.0, std=0.02 * scale)
torch.nn.init.normal_(block.mlp.down_proj.weight, mean=0.0, std=0.02 * scale)
def _init_weights(self, module):
# small normal init for all linear and embedding weights
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
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=0.02)
def forward(self, token_ids: torch.Tensor, block_mask) -> torch.Tensor:
B, T = token_ids.shape
# look up embeddings, then slice the RoPE tables to this sequence length
x = self.embed(token_ids)
cos, sin = self.rope_cos[:T], self.rope_sin[:T]
# pass through every decoder block in turn
for block in self.blocks:
x = block(x, cos, sin, block_mask)
# final clean-up, then project to vocab-sized scores (logits)
x = self.final_norm(x)
return self.lm_head(x)
# def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
# B, T = token_ids.shape
# # look up embeddings, then slice the RoPE tables to this sequence length
# x = self.embed(token_ids)
# cos, sin = self.rope_cos[:T], self.rope_sin[:T]
# # pass through every decoder block in turn
# for block in self.blocks:
# x = block(x, cos, sin)
# # final clean-up, then project to vocab-sized scores (logits)
# x = self.final_norm(x)
# return self.lm_head(x) |