File size: 3,050 Bytes
b615ad6 | 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 | """The assembler: layer 0 + layer 1 + 9x layer 2 + layer 11 = 12 layers.
The only file that knows the whole stack. forward(ids) -> logits (B, T, vocab).
"""
from dataclasses import dataclass, field
import torch.nn as nn
from .layer_0 import Layer0, Layer0Config
from .layer_1 import Layer1, Layer1Config
from .layer_2 import Layer2, Layer2Config
from .layer_11 import Layer11, Layer11Config
@dataclass
class NotioConfig:
layer0: Layer0Config = field(default_factory=Layer0Config)
layer1: Layer1Config = field(default_factory=Layer1Config)
block: Layer2Config = field(default_factory=Layer2Config)
head: Layer11Config = field(default_factory=Layer11Config)
n_blocks: int = 2 # layers 2..10; + head = layer 11 -> 12 layers total
class Notio(nn.Module):
def __init__(self, cfg: NotioConfig):
super().__init__()
self.cfg = cfg
# cross-layer contracts (the one place they can be checked)
assert cfg.layer1.block_size >= cfg.layer0.block_size, "layer 1 capacity must be >= layer 0 sequence length"
assert cfg.layer1.d_model == cfg.block.d_model == cfg.head.d_model, "d_model must match everywhere"
assert cfg.layer1.vocab_size == cfg.head.vocab_size, "vocab_size must match layer 1 and head"
self.layer0 = Layer0(cfg.layer0) # data (not an nn.Module)
self.layer1 = Layer1(cfg.layer1)
self.blocks = nn.ModuleList([Layer2(cfg.block) for _ in range(cfg.n_blocks)])
self.head = Layer11(cfg.head)
self.tie_head()
self.init_weights()
def init_weights(self):
for module in self.modules():
if isinstance(module, nn.Linear):
nn.init.normal_(module.weight, mean=0.0, std=self.cfg.layer1.init_std)
if module.bias is not None:
nn.init.zeros_(module.bias)
# residual-path projections get the GPT-2 depth scaling
for name, p in self.named_parameters():
if name.endswith("c_proj.weight"):
nn.init.normal_(p, mean=0.0,
std=self.cfg.layer1.init_std / (2 * self.cfg.n_blocks) ** 0.5)
def tie_head(self):
self.head.tie_to(self.layer1)
def forward(self, ids, states=None, pos_offset=0):
"""ids: (B, T) long -> (logits (B, T, vocab), states list).
states: per-block recurrent state from the previous window, or None
for zeros. The caller decides when to detach/reset (truncated BPTT).
pos_offset: position embedding offset for stateful generation.
"""
x = self.layer1(ids, pos_offset)
new_states = []
for i, block in enumerate(self.blocks):
st = None if states is None else states[i]
x, st = block(x, st)
new_states.append(st)
return self.head(x), new_states
@property
def n_params(self):
return sum(p.numel() for p in self.parameters())
def sample_batch(self):
return self.layer0.sample_batch()
|