LLMをゼロから作ってみる — スタンフォードCS336流の学習ロードマップ

+------------------+    +------------------+    +--------------------+
| 1. トークナイザー  | -> | 2. アーキテクチャ  | -> | 3. 学習インフラ      |
|  BPEをゼロから    |    |  Transformer     |    |  データ / 分散学習   |
+------------------+    +------------------+    +--------------------+
                                                        |
                                                        v
+------------------+    +------------------+    +--------------------+
| 6. 推論最適化     | <- | 5. アライメント    | <- | 4. スケーリング則    |
|  KV cache        |    |  SFT/RLHF/DPO    |    |  予算配分           |
+------------------+    +------------------+    +--------------------+

from collections import Counter

def get_pair_counts(corpus):
    """トークン列のリストから隣接ペアの頻度を数える。"""
    counts = Counter()
    for tokens in corpus:
        for a, b in zip(tokens, tokens[1:]):
            counts[(a, b)] += 1
    return counts

def merge_pair(corpus, pair, new_token):
    """corpus全体でpairをnew_tokenにマージする。"""
    merged = []
    for tokens in corpus:
        out, i = [], 0
        while i < len(tokens):
            if i < len(tokens) - 1 and (tokens[i], tokens[i + 1]) == pair:
                out.append(new_token)
                i += 2
            else:
                out.append(tokens[i])
                i += 1
        merged.append(out)
    return merged

def train_bpe(text, vocab_size):
    # バイトレベルから開始：基本語彙は256個
    corpus = [list(text.encode("utf-8"))]
    merges = {}
    next_id = 256
    while next_id < vocab_size:
        counts = get_pair_counts(corpus)
        if not counts:
            break
        best = counts.most_common(1)[0][0]
        merges[best] = next_id
        corpus = merge_pair(corpus, best, next_id)
        next_id += 1
    return merges

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head = n_head
        self.head_dim = d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, head_dim)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # アテンションスコア = Q @ K^T / sqrt(head_dim)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        # 因果マスク：未来のトークンを見ないように
        mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
        scores = scores.masked_fill(mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)

        out = attn @ v  # (B, n_head, T, head_dim)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)

生のHTML
   |  テキスト抽出（ボイラープレート除去）
   v
言語フィルタリング（fastText分類器）
   |
   v
品質フィルタリング（ヒューリスティック + 分類器）
   |
   v
重複除去（MinHash / 完全一致）
   |
   v
トークン化 + シャッフル + 固定長チャンク
   |
   v
学習用バイナリシャード (.bin)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
rank = dist.get_rank()
torch.cuda.set_device(rank)

model = GPT(config).cuda(rank)
model = DDP(model, device_ids=[rank])
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

for step, batch in enumerate(loader):
    x, y = batch
    x, y = x.cuda(rank), y.cuda(rank)
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

def chinchilla_optimal(compute_budget_flops):
    """6ND近似を使ったおおまかなcompute-optimal配分。
    C = 6 * N * D、D = 20 * N を仮定。"""
    # C = 6 * N * (20 * N) = 120 * N^2
    n_params = (compute_budget_flops / 120) ** 0.5
    n_tokens = 20 * n_params
    return n_params, n_tokens

# 例：A100を8枚で2週間（およそ1.2e21 FLOPsを仮定）
params, tokens = chinchilla_optimal(1.2e21)
print(f"params: {params/1e9:.2f}B, tokens: {tokens/1e9:.0f}B")
# およそ3Bパラメータ、63Bトークン規模

ベースモデル
   |
   v
SFT（教師あり微調整）
   |  高品質な指示-応答ペアで学習
   v
選好最適化（RLHF または DPO）
   |  人間が好む応答の方向へ調整
   v
最終的なアシスタントモデル

def dpo_loss(policy_chosen_logps, policy_rejected_logps,
             ref_chosen_logps, ref_rejected_logps, beta=0.1):
    """選ばれた応答と棄却された応答の対数確率差を、
    参照モデル比で広げる方向に学習する。"""
    chosen_ratio = policy_chosen_logps - ref_chosen_logps
    rejected_ratio = policy_rejected_logps - ref_rejected_logps
    logits = beta * (chosen_ratio - rejected_ratio)
    return -F.logsigmoid(logits).mean()

class KVCache:
    def __init__(self, batch, n_head, max_len, head_dim, device):
        shape = (batch, n_head, max_len, head_dim)
        self.k = torch.zeros(shape, device=device, dtype=torch.bfloat16)
        self.v = torch.zeros(shape, device=device, dtype=torch.bfloat16)
        self.pos = 0

    def update(self, k_new, v_new):
        t = k_new.size(2)
        self.k[:, :, self.pos:self.pos + t] = k_new
        self.v[:, :, self.pos:self.pos + t] = v_new
        self.pos += t
        return self.k[:, :, :self.pos], self.v[:, :, :self.pos]

# 仮想環境と主要な依存関係
python -m venv .venv
source .venv/bin/activate
pip install torch numpy tiktoken datasets wandb

# nanoGPTで最初の学習感覚をつかむ
git clone https://github.com/karpathy/nanoGPT.git
cd nanoGPT
python data/shakespeare_char/prepare.py
python train.py config/train_shakespeare_char.py

@torch.no_grad()
def estimate_perplexity(model, loader, max_batches=50):
    model.eval()
    total_loss, n = 0.0, 0
    for i, (x, y) in enumerate(loader):
        if i >= max_batches:
            break
        logits = model(x.cuda())
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)), y.cuda().view(-1)
        )
        total_loss += loss.item()
        n += 1
    model.train()
    return math.exp(total_loss / n)