Spaces:
Runtime error
Runtime error
| """ | |
| Abyss-1 the full de coder-only transformer. | |
| tokens -> embedding -> [block] x n_layer -> final RMSNorm -> LM head | |
| forward() returns logits, and (when targets are given) the cross-entropy | |
| loss for next-token prediction. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .norm import AbyssRMSNorm | |
| from .block import AbyssBlock | |
| 8934 | |
| class TheAbyssModel(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| self.token_emb = nn.Embedding(config.vocab_size, config.n_embd) | |
| self.drop = nn.Dropout(config.dropout) | |
| self.blocks = nn.ModuleList([AbyssBlock(config) for _ in range(config.n_layer)]) | |
| self.final_norm = AbyssRMSNorm(config.n_embd, config.norm_eps) | |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True) | |
| if config.tie_embeddings: | |
| self.lm_head.weight = self.token_emb.weight | |
| self.apply(self._init_weights) | |
| # GPT-2 trick: shrink the projections that feed residual adds, so the | |
| # residual stream doesn't blow up as depth grows | |
| for name, p in self.named_parameters(): | |
| if name.endswith("proj.weight") or name.endswith("down.weight"): | |
| nn.init.normal_(p, mean=0.0, std=0.02 / (2 * config.n_layer) ** 0.5) | |
| def _init_weights(self, 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, targets=None): | |
| B, T = idx.shape | |
| assert T <= self.config.block_size, f"sequence {T} > block_size {self.config.block_size}" | |
| x = self.drop(self.token_emb(idx)) | |
| for block in self.blocks: | |
| x = block(x) | |
| x = self.final_norm(x) | |
| logits = self.lm_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 | |
| def num_params(self, non_embedding=False): | |
| n = sum(p.numel() for p in self.parameters()) | |
| if non_embedding: | |
| n -= self.token_emb.weight.numel() | |
| if not self.config.tie_embeddings: | |
| n -= self.lm_head.weight .numel() | |
| return n | |