torch_core.py

  1"""Differentiable counterparts and a small encoder-decoder transformer.
  2
  3Install numpy and torch. Run: python torch_core.py
  4No nn.Transformer wrapper: the attention and block operations are explicit.
  5"""
  6from __future__ import annotations
  7
  8import math
  9import numpy as np
 10import torch
 11from torch import Tensor, nn
 12from torch.nn import functional as F
 13
 14# region embeddings
 15def embedding_demo() -> tuple[Tensor, nn.Embedding, Tensor]:
 16    vocab = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "the": 3,
 17             "small": 4, "robot": 5, "moves": 6, "<unk>": 7}
 18    tokens = "the small robot moves".split()
 19    ids = torch.tensor([vocab.get(token, vocab["<unk>"]) for token in tokens])
 20    # Match the NumPy weights exactly, rather than assuming RNGs match.
 21    weights = torch.from_numpy(np.random.default_rng(0).normal(size=(len(vocab), 6)))
 22    table = nn.Embedding.from_pretrained(weights, freeze=False)
 23    x = table(ids)
 24    torch.testing.assert_close(x[2], table.weight[vocab["robot"]])
 25    return ids, table, x
 26# endregion embeddings
 27
 28# region positions
 29def positions(length: int, width: int, start: int = 0, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32) -> Tensor:
 30    if width <= 0 or width % 2:
 31        raise ValueError("This example uses a positive, even model width")
 32    pos = torch.arange(start, start + length, device=device, dtype=dtype)[:, None]
 33    frequency = 10000.0 ** (-torch.arange(0, width, 2, device=device, dtype=dtype) / width)
 34    angles = pos * frequency[None, :]
 35    pe = torch.empty(length, width, device=device, dtype=dtype)
 36    pe[:, 0::2] = torch.sin(angles)
 37    pe[:, 1::2] = torch.cos(angles)
 38    return pe
 39# endregion positions
 40
 41# region attention
 42def attention(q: Tensor, k: Tensor, v: Tensor, allowed: Tensor | None = None) -> tuple[Tensor, Tensor]:
 43    """(..., T, dk), (..., S, dk), (..., S, dv). True means visible."""
 44    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
 45        raise ValueError("Q/K feature widths and K/V source lengths must match")
 46    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
 47    if allowed is not None:
 48        visible = torch.broadcast_to(allowed.to(device=q.device, dtype=torch.bool), scores.shape)
 49        if not visible.any(dim=-1).all():
 50            raise ValueError("Every query needs at least one visible key")
 51        scores = scores.masked_fill(~visible, float("-inf"))
 52    weights = scores.softmax(dim=-1)
 53    return weights @ v, weights
 54# endregion attention
 55
 56# region masks
 57def attention_masks(source_ids: Tensor, target_ids: Tensor, pad: int = 0) -> tuple[Tensor, Tensor]:
 58    source_visible = (source_ids != pad)[:, None, None, :]
 59    target_visible = (target_ids != pad)[:, None, None, :]
 60    length = target_ids.shape[1]
 61    index = torch.arange(length, device=target_ids.device)
 62    causal = index[None, :] <= index[:, None]
 63    target_visible = target_visible & causal[None, None, :, :]
 64    return source_visible, target_visible
 65# endregion masks
 66
 67# region heads
 68class MultiHead(nn.Module):
 69    def __init__(self, width: int, heads: int, seed: int = 0) -> None:
 70        super().__init__()
 71        if heads <= 0 or width % heads:
 72            raise ValueError("Model width must be divisible by head count")
 73        self.heads, self.head_width = heads, width // heads
 74        self.wq, self.wk, self.wv, self.wo = [nn.Linear(width, width, bias=False) for _ in range(4)]
 75        # Shared initialization makes the two libraries directly comparable.
 76        rng = np.random.default_rng(seed)
 77        with torch.no_grad():
 78            for layer in [self.wq, self.wk, self.wv, self.wo]:
 79                w = rng.normal(size=(width, width)) / np.sqrt(width)
 80                layer.weight.copy_(torch.tensor(w.T, dtype=layer.weight.dtype))
 81
 82    def split(self, x: Tensor) -> Tensor:
 83        batch, length, _ = x.shape
 84        return x.reshape(batch, length, self.heads, self.head_width).transpose(1, 2)
 85
 86    def forward(self, query: Tensor, context: Tensor, allowed: Tensor | None = None) -> tuple[Tensor, Tensor]:
 87        q, k, v = self.split(self.wq(query)), self.split(self.wk(context)), self.split(self.wv(context))
 88        heads, weights = attention(q, k, v, allowed)
 89        batch, _, length, _ = heads.shape
 90        joined = heads.transpose(1, 2).reshape(batch, length, -1)
 91        return self.wo(joined), weights
 92# endregion heads
 93
 94# region norm_ffn
 95def layer_norm(x: Tensor, gamma: Tensor | None = None, beta: Tensor | None = None, eps: float = 1e-5) -> Tensor:
 96    # normalized_shape=(D,) means statistics over features of EACH token.
 97    return F.layer_norm(x, (x.shape[-1],), weight=gamma, bias=beta, eps=eps)
 98
 99
100def feed_forward(x: Tensor, w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> Tensor:
101    hidden = torch.relu(x @ w1 + b1)
102    return hidden @ w2 + b2
103# endregion norm_ffn
104
105# region encoder
106class EncoderBlock(nn.Module):
107    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 10) -> None:
108        super().__init__()
109        self.attn = MultiHead(width, heads, seed)
110        self.norm1, self.norm2 = nn.LayerNorm(width), nn.LayerNorm(width)
111        self.ffn = nn.Sequential(nn.Linear(width, inner), nn.ReLU(), nn.Linear(inner, width))
112
113    def forward(self, x: Tensor, source_visible: Tensor) -> Tensor:
114        normalized = self.norm1(x)
115        x = x + self.attn(normalized, normalized, source_visible)[0]
116        return x + self.ffn(self.norm2(x))
117# endregion encoder
118
119# region decoder
120class DecoderBlock(nn.Module):
121    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 20) -> None:
122        super().__init__()
123        self.self_attn = MultiHead(width, heads, seed)
124        self.cross_attn = MultiHead(width, heads, seed + 1)
125        self.norm1, self.norm2, self.norm3 = [nn.LayerNorm(width) for _ in range(3)]
126        self.ffn = nn.Sequential(nn.Linear(width, inner), nn.ReLU(), nn.Linear(inner, width))
127
128    def forward(self, y: Tensor, memory: Tensor, target_visible: Tensor, source_visible: Tensor) -> Tensor:
129        normalized = self.norm1(y)
130        y = y + self.self_attn(normalized, normalized, target_visible)[0]
131        y = y + self.cross_attn(self.norm2(y), memory, source_visible)[0]
132        return y + self.ffn(self.norm3(y))
133# endregion decoder
134
135# region architectures
136def architecture_demo() -> tuple[Tensor, Tensor]:
137    source = torch.tensor([[4, 7, 9], [6, 8, 0]])
138    prefix = torch.tensor([[1, 4, 7], [1, 6, 8]])
139    width, vocabulary = 32, 12
140    embedding = nn.Embedding(vocabulary, width)
141    source_mask, target_mask = attention_masks(source, prefix)
142    x = embedding(source) + positions(source.shape[1], width)
143    y = embedding(prefix) + positions(prefix.shape[1], width)
144
145    # Encoder-only: pool only real source positions for classification.
146    encoded = layer_norm(EncoderBlock()(x, source_mask))
147    real = (source != 0)[..., None]
148    pooled = (encoded * real).sum(1) / real.sum(1)
149    class_logits = nn.Linear(width, 3)(pooled)
150
151    # Decoder-only: same two-sublayer structure, with CAUSAL visibility.
152    causal_hidden = layer_norm(EncoderBlock(seed=30)(y, target_mask))
153    next_token_logits = nn.Linear(width, vocabulary)(causal_hidden)
154    return class_logits, next_token_logits  # (2, 3), (2, 3, 12)
155# endregion architectures
156
157# region model
158class TinyTransformer(nn.Module):
159    def __init__(self, vocabulary: int = 12, width: int = 32, heads: int = 4, inner: int = 64) -> None:
160        super().__init__()
161        self.embedding = nn.Embedding(vocabulary, width, padding_idx=0)
162        self.encoder = EncoderBlock(width, heads, inner)
163        self.decoder = DecoderBlock(width, heads, inner)
164        self.encoder_norm, self.decoder_norm = nn.LayerNorm(width), nn.LayerNorm(width)
165        self.readout = nn.Linear(width, vocabulary)
166
167    def embed(self, ids: Tensor) -> Tensor:
168        x = self.embedding(ids)
169        return x + positions(ids.shape[1], x.shape[-1], device=x.device, dtype=x.dtype)
170
171    def encode(self, source: Tensor) -> tuple[Tensor, Tensor]:
172        visible = (source != 0)[:, None, None, :]
173        memory = self.encoder_norm(self.encoder(self.embed(source), visible))
174        return memory, visible
175
176    def decode(self, target_prefix: Tensor, memory: Tensor, source_visible: Tensor) -> Tensor:
177        _, target_visible = attention_masks(target_prefix, target_prefix)
178        y = self.decoder(self.embed(target_prefix), memory, target_visible, source_visible)
179        return self.readout(self.decoder_norm(y))
180
181    def forward(self, source: Tensor, target_prefix: Tensor) -> Tensor:
182        memory, visible = self.encode(source)
183        return self.decode(target_prefix, memory, visible)
184# endregion model
185
186# region generate
187@torch.no_grad()
188def generate(model: TinyTransformer, source: Tensor, max_new_tokens: int = 4, bos: int = 1, eos: int = 2, pad: int = 0) -> Tensor:
189    was_training = model.training
190    model.eval()
191    try:
192        memory, visible = model.encode(source)
193        prefix = torch.full((source.shape[0], 1), bos, dtype=torch.long, device=source.device)
194        finished = torch.zeros(source.shape[0], dtype=torch.bool, device=source.device)
195        for _ in range(max_new_tokens):
196            logits = model.decode(prefix, memory, visible)[:, -1, :]
197            token = logits.argmax(dim=-1)  # Greedy decoding for this experiment.
198            token = torch.where(finished, pad, token)
199            prefix = torch.cat([prefix, token[:, None]], dim=1)
200            finished |= token == eos
201            if finished.all():
202                break
203        return prefix[:, 1:]
204    finally:
205        model.train(was_training)
206# endregion generate
207
208if __name__ == "__main__":
209    torch.manual_seed(7)
210    model = TinyTransformer()
211    source = torch.tensor([[4, 7, 9], [6, 8, 0]])
212    prefix = torch.tensor([[1, 4, 7], [1, 6, 8]])
213    print("Vocabulary logits:", tuple(model(source, prefix).shape))