| import torch |
| import torch.nn as nn |
|
|
| from .config import IvmeConfig |
| from .rmsnorm import RMSNorm |
| from .rope import precompute_rope_freqs |
| from .attention import CausalSelfAttention |
| from .feedforward import SwiGLU |
|
|
|
|
| class TransformerBlock(nn.Module): |
| """One dense transformer layer (Section 3): pre-norm attention + pre-norm SwiGLU, |
| with residual connections around each. Identical shape repeated n_layers times -- |
| no loops, no weight sharing (Section 3.1, distinguishing this from the shelved |
| Ivmetron design). |
| """ |
|
|
| def __init__(self, cfg: IvmeConfig): |
| super().__init__() |
| self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps) |
| self.attn = CausalSelfAttention(cfg.hidden_dim, cfg.n_heads, cfg.dropout) |
| self.ffn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps) |
| self.ffn = SwiGLU(cfg.hidden_dim, cfg.ffn_mult) |
|
|
| def forward(self, x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.attn_norm(x), rope_freqs) |
| x = x + self.ffn(self.ffn_norm(x)) |
| return x |
|
|
|
|
| class IvmeConversateV2(nn.Module): |
| """Ivme-Conversate-v2 (Dense) -- the full model described in Section 4. |
| |
| ~20M parameters, 10 layers, hidden_dim 384, 6 heads, RoPE, SwiGLU, RMSNorm, |
| tied embeddings, 16k vocab, 1024 context. See config.py for the exact spec. |
| """ |
|
|
| def __init__(self, cfg: IvmeConfig): |
| super().__init__() |
| self.cfg = cfg |
|
|
| self.tok_embed = nn.Embedding(cfg.vocab_size, cfg.hidden_dim) |
| self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)]) |
| self.final_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps) |
|
|
| |
| |
| self.lm_head = nn.Linear(cfg.hidden_dim, cfg.vocab_size, bias=False) |
| if cfg.tie_embeddings: |
| self.lm_head.weight = self.tok_embed.weight |
|
|
| rope_freqs = precompute_rope_freqs(cfg.head_dim, cfg.context_len, cfg.rope_theta) |
| self.register_buffer("rope_freqs", rope_freqs, persistent=False) |
|
|
| self.apply(self._init_weights) |
|
|
| 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.context_len, ( |
| f"sequence length {T} exceeds context_len {self.cfg.context_len}" |
| ) |
|
|
| x = self.tok_embed(idx) |
| for block in self.blocks: |
| x = block(x, self.rope_freqs) |
| x = self.final_norm(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
| if targets is not None: |
| loss = nn.functional.cross_entropy( |
| logits.view(-1, logits.size(-1)), |
| targets.view(-1), |
| ignore_index=-1, |
| ) |
| return logits, loss |
|
|
| def num_params(self, non_embedding: bool = False) -> int: |
| n = sum(p.numel() for p in self.parameters()) |
| if non_embedding: |
| n -= self.tok_embed.weight.numel() |
| return n |
|
|
