一文要約: model.py は突然現れるコードの塊ではありません。これまでに導いた Embedding、位置Encoding、Causal Self-Attention、FFN、Norm、Residual path、出力射影を、同じData flowの順に接続したものです。

📦 2024年の原始Demo Repository: github.com/waylandzhang/Transformer-from-scratch。歴史的なDemoはModel、訓練、生成を一つの model.py に置いています。本章は今回の監査で修正した独立Teaching版であり、行ごとのMirrorではありません。

18.1 まずModelの境界を決める

実装するのは小さな decoder-only Transformer です。

Token IDs [B, T]
    ↓
Token Embedding + 固定Sinusoidal Position Encoding
    ↓
N × Pre-Norm Transformer Block
    ├── LayerNorm → Causal Multi-Head Attention → Residual Add
    └── LayerNorm → FFN → Residual Add
    ↓
Final LayerNorm
    ↓
Weightを共有するOutput Projection
    ↓
Logits [B, T, vocab_size]

仕組みを見えるままにするため、Attention は手で書きます。KV Cache、RoPE、Flash Attention、padding mask、分散訓練はまだ加えません。このコードが扱うのは同じ長さで PAD のない batch です。Paddingを使う場合は、第16章で扱った key padding mask も必要です。

また、これは GPT-2 の各部品をそのまま再現するModelではありません。GPT型の Pre-Norm decoder blockを使いながら、元の Transformer 型の固定Sinusoidal Positionを残します。これはTeaching上の選択であり、以下では各選択を隠さず明記します。

18.2 ConfigにShapeを守らせる

from dataclasses import dataclass

@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 256
    d_model: int = 512
    n_layers: int = 6
    n_heads: int = 8
    dropout: float = 0.1

    def __post_init__(self):
        if self.vocab_size <= 0 or self.context_length <= 0:
            raise ValueError("vocab_size and context_length must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0:
            raise ValueError("d_model, n_layers, and n_heads must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")

    @property
    def head_dim(self):
        return self.d_model // self.n_heads

head_dim = d_model // n_heads が成立するのは、割り切れる場合だけです。旧コードは割り切れると仮定しながら、出力射影には必ず d_model 個の特徴を要求していました。そのため、無効な設定がずっと後になってから壊れます。

vocab_size は、一つの訓練Fileで偶然観測した最大Token IDではなく、Tokenizerの語彙サイズから取ります。前者では、正規のTokenを扱えるかどうかがSampleに一度出現したかで変わってしまいます。

18.3 Feed-Forward Network

各Sequence位置は、独立して同じFFNを通ります。

[B, T, d_model]
       ↓ Linear
[B, T, 4 × d_model]
       ↓ GELU
[B, T, 4 × d_model]
       ↓ Linear + Dropout
[B, T, d_model]

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        self.net = nn.Sequential(
            nn.Linear(config.d_model, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)

完全なForward passの章で使ったGPT-2の基準と合わせるため、ここではGELUを使います。現代のModelはSwiGLUを採用し、中間幅も単純な 4 × d_model ではない場合があります。本章はまず古典的な2層FFNを実装します。

18.4 読みやすい単一Attention Headから始める

Causal maskを含む式は次のとおりです。

$M_{ij}=0$ は可視、$M_{ij}=-\infty$ はQuery位置 $i$ からKey位置 $j$ を見せない、という意味です。

class AttentionHead(nn.Module):
    def __init__(self, d_model, head_dim, context_length, dropout):
        super().__init__()
        self.head_dim = head_dim
        self.query = nn.Linear(d_model, head_dim, bias=False)
        self.key = nn.Linear(d_model, head_dim, bias=False)
        self.value = nn.Linear(d_model, head_dim, bias=False)
        self.attn_dropout = nn.Dropout(dropout)

        mask = torch.tril(
            torch.ones(context_length, context_length, dtype=torch.bool)
        )
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        _, T, _ = x.shape
        if T > self.causal_mask.size(0):
            raise ValueError("sequence is longer than context_length")

        q = self.query(x)                          # [B, T, head_dim]
        k = self.key(x)                            # [B, T, head_dim]
        v = self.value(x)                          # [B, T, head_dim]

        scores = q @ k.transpose(-2, -1)           # [B, T, T]
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))

        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)
        return weights @ v                         # [B, T, head_dim]

Maskはbool bufferとして登録します。訓練Parameterではありませんが、model.to(device) と一緒に移動します。persistent=False によって、再構成できる三角行列をcheckpointに保存しません。

Softmax後にDropoutを掛けると、訓練中のAttention weightの行和は必ずしも1ではありません。Inverted Dropoutが保つのは期待値です。model.eval() ではDropoutが無効になります。

18.5 分離したHeadから融合実装へ

18.5.1 Teaching版: HeadごとにModuleを作る

class TeachingMultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = nn.ModuleList([
            AttentionHead(
                config.d_model,
                config.head_dim,
                config.context_length,
                config.dropout,
            )
            for _ in range(config.n_heads)
        ])
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        y = torch.cat([head(x) for head in self.heads], dim=-1)
        return self.resid_dropout(self.out_proj(y))

数学上、各Headに順序依存はありません。しかし、このPythonのlist comprehensionは別々のModuleを呼び出します。読みやすい実装ですが、GPUに仕事を渡す最も効率的な形ではありません。

18.5.2 実用版: 一度Projectionしてreshapeする

元のTransformer論文は、Headごとの行列 $W_i^Q,W_i^K,W_i^V$ を定義しています。しかしPython objectをHeadごとに一つ作れとは規定していません。各行列を出力方向に連結すれば、一度の大きな行列積ですべてのQ、K、Vを作れます。これは論文式と代数的に等価な一般的融合実装であり、「論文のsource code」ではありません。

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.context_length = config.context_length

        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.context_length,
            config.context_length,
            dtype=torch.bool,
        )).view(1, 1, config.context_length, config.context_length)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        if T > self.context_length:
            raise ValueError("sequence is longer than context_length")

        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:, :, :T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        y = weights @ v                              # [B, H, T, D]
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.out_proj(y))

Shapeの変化:

x                 [B, T, C]
q, k, v           [B, T, C] each
split heads       [B, H, T, D]   where C = H × D
scores / weights  [B, H, T, T]
weighted values   [B, H, T, D]
merge heads       [B, T, C]

Production codeでは、手書き部分をPyTorchの scaled_dot_product_attention に置き換え、条件が合えばBackendにFlash Attentionなどのkernelを選ばせられます。本章ではMask、Softmax、Shapeを見えるままにするため、明示的な実装を残します。

18.6 Transformer Block

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

これはPre-Normです。NormをSub-layerの前に置き、Residual branchは正規化していない x を保持します。GPT-2はLayerNormによるPre-Normを使用します。LLaMAもPre-Norm topologyですが、NormはRMSNormです。Pre-Normは深いNetworkを最適化しやすいことが多いものの、あらゆる指標でPost-Normより無条件に優れるわけではありません。

18.7 TokenとPositionをModelへ入れる

18.7.1 Position tableは一度だけ作る

旧版は forward() のたびにSinusoidal table全体を作り直し、Configにdevice文字列を保存していました。計算が重複し、model.to(...) の後でModelと設定Deviceが食い違う可能性があります。

def sinusoidal_positions(max_length, d_model):
    position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1)
    frequency = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float32)
        * (-math.log(10000.0) / d_model)
    )
    angles = position * frequency

    table = torch.zeros(max_length, d_model)
    table[:, 0::2] = torch.sin(angles)
    table[:, 1::2] = torch.cos(angles[:, :table[:, 1::2].shape[1]])
    return table.unsqueeze(0)  # [1, context_length, d_model]

最後のsliceにより、奇数の d_model でもcos代入時のshape mismatchを防げます。一般的な設定は偶数でも、helperが暗黙に依存する必要はありません。

Model.__init__ で一度だけ登録します。

self.register_buffer(
    "position_encoding",
    sinusoidal_positions(config.context_length, config.d_model),
    persistent=False,
)

これでtableはModelと一緒にdevice・dtypeを移動し、古くなる self.device は不要です。

18.7.2 Weight tying

self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
self.apply(self._init_weights)
self.lm_head.weight = self.token_embedding.weight

最後の行はInput embeddingとOutput projectionを、同じ一つの Parameter にします。現在値を別MatrixへCopyしているのではありません。

PyTorch documentationでは、nn.Embedding のWeightはDefaultで $\mathcal{N}(0,1)$ から初期化されます。ここでは $\sqrt{d_{model}}$ のInput scalingとTied outputを併用するため、Defaultのままでは初期logitsが異常に大きくなります。そこでWeightを共有する前に、Model全体を初期化します。

@staticmethod
def _init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    if isinstance(module, nn.Linear) and module.bias is not None:
        nn.init.zeros_(module.bias)

0.02 はこのTeaching modelが明示的に採用するStart pointで、すべてのArchitectureに対する法則ではありません。より深いGPT recipeでは、一部のResidual projectionをLayer数に応じて追加Scaleすることもあります。

18.7.3 ForwardとLoss

def forward(self, idx, targets=None):
    if idx.ndim != 2:
        raise ValueError("idx must have shape [batch, time]")
    _, T = idx.shape
    if T == 0:
        raise ValueError("sequence must contain at least one token")
    if T > self.config.context_length:
        raise ValueError("sequence is longer than context_length")
    if targets is not None and targets.shape != idx.shape:
        raise ValueError("targets must have the same shape as idx")

    token = self.token_embedding(idx) * math.sqrt(self.config.d_model)
    position = self.position_encoding[:, :T, :]
    x = self.embedding_dropout(token + position)

    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.reshape(-1, logits.size(-1)),
            targets.reshape(-1),
        )
    return logits, loss

元のTransformerは、Sinusoidal positionと加算する前にToken embeddingを $\sqrt{d_{model}}$ 倍します。このModelもその約束を明示的に採用します。

F.cross_entropy は、Softmaxを手で掛けていない生のlogitsと整数Token targetsを受け取ります。第19章で訓練Dataを一つずらし、idx=[x0,...,xT-1]、targets=[x1,...,xT] を用意します。

18.8 生成関数にも境界がある

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    if idx.ndim != 2 or idx.size(1) == 0:
        raise ValueError("idx must have shape [batch, time] with time > 0")
    if max_new_tokens < 0:
        raise ValueError("max_new_tokens must be non-negative")
    if temperature < 0:
        raise ValueError("temperature must be non-negative")
    if top_k is not None and top_k <= 0:
        raise ValueError("top_k must be positive")

    for _ in range(max_new_tokens):
        idx_crop = idx[:, -self.config.context_length:]
        logits, _ = self(idx_crop)
        next_logits = logits[:, -1, :]

        if temperature == 0:
            idx_next = torch.argmax(next_logits, dim=-1, keepdim=True)
        else:
            next_logits = next_logits / temperature
            if top_k is None:
                probs = F.softmax(next_logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
            else:
                k = min(top_k, next_logits.size(-1))
                top_logits, top_indices = torch.topk(next_logits, k, dim=-1)
                top_probs = F.softmax(top_logits, dim=-1)
                choice = torch.multinomial(top_probs, num_samples=1)
                idx_next = top_indices.gather(-1, choice)

        idx = torch.cat((idx, idx_next), dim=1)
    return idx

三つの境界を曖昧にしません。

• temperature == 0 はGreedyへ分岐し、本当に0で割らない。
• torch.topk が返すK個のlogits内で直接Samplingするため、境界Scoreが同点でもK個より多いIndexを残さない。
• @torch.no_grad() が止めるのはGradient記録でありDropoutではない。生成前に呼び出し側が model.eval() を実行する。

このTeaching loopは毎Stepで現在Windowを再計算し、KV Cacheを持ちません。最古Tokenを切り捨てた後、固定Absolute positionはWindow内の0から再開します。これはこの実装で明示的に選んだSliding-window policyで、すべてのServing systemのDefaultではありません。

18.9 完成版 model.py

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
from torch.nn import functional as F


@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 256
    d_model: int = 512
    n_layers: int = 6
    n_heads: int = 8
    dropout: float = 0.1

    def __post_init__(self):
        if self.vocab_size <= 0 or self.context_length <= 0:
            raise ValueError("vocab_size and context_length must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0:
            raise ValueError("d_model, n_layers, and n_heads must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")

    @property
    def head_dim(self):
        return self.d_model // self.n_heads


def sinusoidal_positions(max_length, d_model):
    position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1)
    frequency = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float32)
        * (-math.log(10000.0) / d_model)
    )
    angles = position * frequency
    table = torch.zeros(max_length, d_model)
    table[:, 0::2] = torch.sin(angles)
    table[:, 1::2] = torch.cos(angles[:, :table[:, 1::2].shape[1]])
    return table.unsqueeze(0)


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        self.net = nn.Sequential(
            nn.Linear(config.d_model, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.context_length = config.context_length
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.context_length,
            config.context_length,
            dtype=torch.bool,
        )).view(1, 1, config.context_length, config.context_length)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        if T > self.context_length:
            raise ValueError("sequence is longer than context_length")

        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:, :, :T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        y = weights @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.out_proj(y))


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x


class Model(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.embedding_dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)
        ])
        self.final_norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.apply(self._init_weights)
        self.lm_head.weight = self.token_embedding.weight

        self.register_buffer(
            "position_encoding",
            sinusoidal_positions(config.context_length, config.d_model),
            persistent=False,
        )

    @staticmethod
    def _init_weights(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    def forward(self, idx, targets=None):
        if idx.ndim != 2:
            raise ValueError("idx must have shape [batch, time]")
        _, T = idx.shape
        if T == 0:
            raise ValueError("sequence must contain at least one token")
        if T > self.config.context_length:
            raise ValueError("sequence is longer than context_length")
        if targets is not None and targets.shape != idx.shape:
            raise ValueError("targets must have the same shape as idx")

        token = self.token_embedding(idx) * math.sqrt(self.config.d_model)
        position = self.position_encoding[:, :T, :]
        x = self.embedding_dropout(token + position)

        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.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        if idx.ndim != 2 or idx.size(1) == 0:
            raise ValueError("idx must have shape [batch, time] with time > 0")
        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if temperature < 0:
            raise ValueError("temperature must be non-negative")
        if top_k is not None and top_k <= 0:
            raise ValueError("top_k must be positive")

        for _ in range(max_new_tokens):
            idx_crop = idx[:, -self.config.context_length:]
            logits, _ = self(idx_crop)
            next_logits = logits[:, -1, :]

            if temperature == 0:
                idx_next = torch.argmax(next_logits, dim=-1, keepdim=True)
            else:
                next_logits = next_logits / temperature
                if top_k is None:
                    probs = F.softmax(next_logits, dim=-1)
                    idx_next = torch.multinomial(probs, num_samples=1)
                else:
                    k = min(top_k, next_logits.size(-1))
                    top_logits, top_indices = torch.topk(next_logits, k, dim=-1)
                    top_probs = F.softmax(top_logits, dim=-1)
                    choice = torch.multinomial(top_probs, num_samples=1)
                    idx_next = top_indices.gather(-1, choice)

            idx = torch.cat((idx, idx_next), dim=1)
        return idx

18.10 「動いた」だけでなくSemanticsを検証する

torch.manual_seed(7)
config = ModelConfig(
    vocab_size=32,
    context_length=8,
    d_model=12,
    n_layers=2,
    n_heads=3,
    dropout=0.0,
)
model = Model(config).eval()

idx = torch.tensor([[1, 2, 3, 4]])
targets = torch.tensor([[2, 3, 4, 5]])

with torch.inference_mode():
    logits, loss = model(idx, targets)
assert logits.shape == (1, 4, 32)
assert loss.ndim == 0

# 未来Tokenを変えても、最初の2位置のlogitsは変わらない。
future_changed = torch.tensor([[1, 2, 9, 10]])
with torch.inference_mode():
    a, _ = model(idx)
    b, _ = model(future_changed)
torch.testing.assert_close(a[:, :2], b[:, :2])

# Greedyで3 Token生成すると、長さは3増える。
generated = model.generate(idx[:, :2], 3, temperature=0)
assert generated.shape == (1, 5)

三つのassertionは、出力Contract、Causal Semantics、生成長をそれぞれ検証します。Exceptionが出なかっただけでは十分ではありません。

18.11 Parameter数: TiedとUntiedでは大きく違う

d_model=512、n_heads=8、n_layers=6、vocab_size=50,000 の場合:

Weightを共有しなければ、Output matrixとしてさらに 512 × 50,000 = 25,600,000 が加わり、合計70,103,040です。旧版はOutputを独立Parameterとして数えながら、biasとNormの仮定を「約7000万」に隠していました。ここでは仮定と正確な合計を明示します。

章末チェックリスト

• Matrix multiplicationが失敗する前にShape constraintを検証できる
• Bool causal maskを持つ単一Attention headを書ける
• 分離HeadとFused QKVが代数的に等価でも実行方法は違うと説明できる
• Pre-Norm block、固定Position buffer、本当のweight tyingを実装できる
• temperature=0、Top-Kの同点、model.eval()を正しく扱える
• Causal invarianceをテストし、未来を見ていないことを確認できる

次章予告

ModelはForwardを実行し、Random parameterからToken IDsを生成できるようになりました。しかしDataからはまだ何も学んでいません。

第19章では train.py を書きます。1 TokenずらしたInputとTargetsを作り、Train/Validation Dataを分け、Backwardを行い、Optimizerで本当にParameterを更新します。