Spaces:
Runtime error
Runtime error
File size: 2,605 Bytes
cc15101 8bc38db cc15101 | 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 | """
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
|