一文要約: 訓練は「Batchを取る → Lossを計算 → Backward → Parameter更新」の反復です。しかし信頼できる train.py には、Data境界、Eval mode、学習率状態、Gradient norm、Checkpoint復元も必要です。

📦 2024年の原始Demo Repository: github.com/waylandzhang/Transformer-from-scratch。歴史的なDemoは訓練も一つの model.py に置いています。本章は第18章で監査した独立 model.py と組み合わせます。

Token dataとshiftしたBatchからForward、Backward、評価、再開可能Checkpointへ進む訓練Flow

19.1 Random Modelは空ではない

作ったばかりのModelにもRandom parameterがあり、logitsを出せます。ただし、Corpusから安定した規則をまだ学んでいません。

日本語版では、蒔絵道具の目録を例にします。

Token流： [蒔, 絵, 筆, を, 洗, い, 金, 粉, ...]

Input x： [蒔, 絵, 筆, を, 洗, い, 金]
Target y：[絵, 筆, を, 洗, い, 金, 粉]

各位置のTargetは次のTokenで、必ずしも次の文字や単語ではありません。図はShiftを見やすくするため単字で表しています。

V 個のTokenに本当に一様確率を与えるModelなら、Cross-Entropyは $\ln V$ です。これは参考線であり、Random initializationが必ず同じ値になるというTestではありません。第18章のModelはweight tyingとembedding scaleも使うため、初期logitsが一様とは限りません。

19.2 Model ConfigとTrain Configを分ける

ModelConfig はTensor shapeを記述します。訓練には別のConfigを使います。

from dataclasses import dataclass

@dataclass
class TrainConfig:
    batch_size: int = 8
    total_steps: int = 500
    eval_interval: int = 50
    eval_batches: int = 10
    peak_lr: float = 3e-4
    min_lr: float = 3e-5
    warmup_steps: int = 50
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    seed: int = 1337

復元時にこの分離が効きます。Batch sizeや評価間隔を変えたからといってModel shapeが変わってはいけません。またdevice文字列はArchitectureではありません。

数値は本書の小Modelを動かす候補であり、大Modelの万能Recipeではありません。第17章で見たように、学習率は実際のBatch、Data、計算予算で検証します。

19.3 Dataの元の構造を尊重する

19.3.1 CSVを普通のTextとして扱わない

完成Scriptは --data で整形済みのUTF-8 text fileを受け取ります。漆芸道具目録がplain textならそのまま使えます。CSVの一列だけを学習したいなら、Header、管理番号、別Columnまで混ぜず、先にCSV parserで対象Columnを取り出します。

import csv

def csv_column_to_text(path, column):
    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        if column not in (reader.fieldnames or []):
            raise ValueError(f"missing CSV column: {column}")
        rows = [
            row[column].strip()
            for row in reader
            if row[column].strip()
        ]
    if not rows:
        raise ValueError("the selected CSV column contains no text")
    return "\n".join(rows)

19.3.2 Tokenizer語彙は観測した最大IDではない

tokenizer_name = "cl100k_base"
tokenizer = tiktoken.get_encoding(tokenizer_name)
token_ids = tokenizer.encode(text, disallowed_special=())
tokens = torch.tensor(token_ids, dtype=torch.long)  # CorpusはCPUに置く

model_config = ModelConfig(
    vocab_size=tokenizer.n_vocab,
    context_length=128,
    d_model=80,
    n_layers=6,
    n_heads=4,
    dropout=0.1,
)

max(token_ids) + 1 は、このSampleで観測したIDの範囲にすぎません。正規のTokenizer語彙には tokenizer.n_vocab を使います。Corpus全体をAcceleratorへ置かず、現在のBatchだけを移せば、GPU memoryも無駄にしません。

19.3.3 先にSplitし、各領域の中でWindowを取る

def split_tokens(tokens, train_fraction, context_length):
    split = int(len(tokens) * train_fraction)
    train_data = tokens[:split]
    valid_data = tokens[split:]
    if min(len(train_data), len(valid_data)) <= context_length:
        raise ValueError("each split needs more than context_length tokens")
    return train_data, valid_data

これで一つのWindowが90%境界を跨ぎません。ただし同じ道具名が重複していれば、TrainとValidの両方に入ることがあります。実Dataでは必要に応じてEntity単位の重複除去やGroup splitを先に行います。

19.4 `get_batch()`：一つずれやすい場所

def get_batch(data, batch_size, context_length, device, generator):
    max_start = len(data) - context_length
    if max_start <= 0:
        raise ValueError("data is too short for context_length")
    starts = torch.randint(
        0, max_start, (batch_size,), generator=generator
    )
    offsets = torch.arange(context_length)
    x = data[starts[:, None] + offsets]
    y = data[starts[:, None] + offsets + 1]
    return x.to(device), y.to(device)

torch.randint は上限を含みません。最大Startは len(data) - context_length - 1 となり、Target用に右側の1 Tokenが残ります。

訓練Batch専用の torch.Generator を使えば、評価Samplingが次の訓練Batchを変えません。Generator stateもCheckpointへ保存できます。

19.5 LossにはProbabilityではなくLogitsを渡す

第18章のModelは次を計算します。

loss = F.cross_entropy(
    logits.reshape(-1, logits.size(-1)),
    targets.reshape(-1),
)

cross_entropy は数値的に安定したLogSoftmaxとNegative Log-Likelihoodを内部で組み合わせます。Inputはraw logits、Targetは整数Token IDsです。

Training lossの低下は、このObjectiveが改善した証拠です。事実性、安全性、読みやすい生成、Validation改善まで自動的に証明するものではありません。

19.6 評価で訓練状態を汚さない

@torch.inference_mode()
def estimate_loss(model, train_data, valid_data, config, device):
    was_training = model.training
    model.eval()
    try:
        result = {}
        eval_generator = torch.Generator().manual_seed(config.seed + 1)
        for name, data in (("train", train_data), ("valid", valid_data)):
            losses = []
            for _ in range(config.eval_batches):
                x, y = get_batch(
                    data,
                    config.batch_size,
                    model.config.context_length,
                    device,
                    eval_generator,
                )
                _, loss = model(x, y)
                losses.append(loss.detach().cpu())
            result[name] = torch.stack(losses).mean().item()
        return result
    finally:
        model.train(was_training)

inference_mode() はAutograd stateを作らず、eval() は本書ModelのDropoutを止め、finally は呼び出し前のModeを復元します。本書ModelにはBatchNormがないため、ここでModeを切り替える説明にBatchNormを持ち出す必要はありません。

固定Eval generatorは毎回同じTrain/Valid windowを取り、Train generatorを消費せずに比較Noiseを減らします。

19.7 AdamW Parameter Groupと第17章のSchedule

def configure_optimizer(model, config):
    decay, no_decay = [], []
    for _, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        (decay if parameter.ndim >= 2 else no_decay).append(parameter)

    groups = [
        {"params": decay, "weight_decay": config.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(
        groups,
        lr=config.peak_lr,
        betas=(0.9, 0.95),
    )

一般的な分け方として、2次元以上のEmbedding/Linear weightをdecayし、1次元のbias・Normはdecayしません。Tied embedding/output Parameterは named_parameters() に一度だけ現れます。

AdamWは「Adam + L2」ではありません。Weight decayをAdamのMoment正規化Loss-gradient updateから分離します。また各Parameterに完璧な個別学習率を自動発見するわけでもありません。

Scheduleは第17章と同じ境界を使います。

def lr_at_step(step, total_steps, warmup_steps, peak_lr, min_lr):
    if not 1 <= step <= total_steps:
        raise ValueError("step must be in [1, total_steps]")
    if not 0 < warmup_steps < total_steps:
        raise ValueError("warmup_steps must be in (0, total_steps)")
    if step <= warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (peak_lr - min_lr) * cosine

19.8 Updateした後で、そのStepを報告する

各Iterationは現在LRを設定し、Batchを取り、Gradientを消し、Lossを計算し、Backward、Global norm clipping、optimizer.step() の順に進みます。step=500 の評価は500回目のUpdate後です。

旧版はindex 499で評価し、その後もう一度Updateしてから、Printed metricsとは別のModel stateを保存していました。

clip_grad_norm_ が返すのはClipping前のtotal normです。Spikeを観測するための値であり、表示値が必ず 1.0 以下になるという意味ではありません。

19.9 Weight保存だけでは訓練を再開できない

正しいRecoveryには、Model parameters、Optimizer moments、完了Step数、両Config、Tokenizer identity、Corpus fingerprint、Batch generator、PyTorch RNG stateが必要です。

完成Scriptは一時Fileへ書いてからDestinationを置換します。Load時は weights_only=True を使い、ConfigやCorpus hashの不一致を拒否します。それでも、自分で作成したか出所を信頼できるCheckpointだけをLoadしてください。

RNG stateを保存しても、GPU、PyTorch version、非決定的Kernelが異なる環境でbit-for-bit一致するとは約束できません。訓練状態は保存できますが、異なる環境を同じ数学装置にはできません。

19.10 完成版 `train.py`

import argparse
import hashlib
import math
from dataclasses import asdict, dataclass
from pathlib import Path

import tiktoken
import torch

from model import Model, ModelConfig


@dataclass
class TrainConfig:
    batch_size: int = 8
    total_steps: int = 500
    eval_interval: int = 50
    eval_batches: int = 10
    peak_lr: float = 3e-4
    min_lr: float = 3e-5
    warmup_steps: int = 50
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    seed: int = 1337

    def __post_init__(self):
        if min(self.batch_size, self.total_steps, self.eval_interval,
               self.eval_batches, self.warmup_steps) <= 0:
            raise ValueError("step and batch settings must be positive")
        if self.warmup_steps >= self.total_steps:
            raise ValueError("warmup_steps must be smaller than total_steps")
        if not 0.0 <= self.min_lr <= self.peak_lr:
            raise ValueError("learning rates must satisfy 0 <= min_lr <= peak_lr")
        if self.weight_decay < 0 or self.grad_clip <= 0:
            raise ValueError("weight_decay and grad_clip are invalid")


def select_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def split_tokens(tokens, train_fraction, context_length):
    if not 0.0 < train_fraction < 1.0:
        raise ValueError("train_fraction must be between 0 and 1")
    split = int(len(tokens) * train_fraction)
    train_data = tokens[:split]
    valid_data = tokens[split:]
    if min(len(train_data), len(valid_data)) <= context_length:
        raise ValueError("each split needs more than context_length tokens")
    return train_data, valid_data


def get_batch(data, batch_size, context_length, device, generator):
    max_start = len(data) - context_length
    if max_start <= 0:
        raise ValueError("data is too short for context_length")
    starts = torch.randint(
        0, max_start, (batch_size,), generator=generator
    )
    offsets = torch.arange(context_length)
    x = data[starts[:, None] + offsets]
    y = data[starts[:, None] + offsets + 1]
    return x.to(device), y.to(device)


@torch.inference_mode()
def estimate_loss(model, train_data, valid_data, config, device):
    was_training = model.training
    model.eval()
    try:
        result = {}
        eval_generator = torch.Generator().manual_seed(config.seed + 1)
        for name, data in (("train", train_data), ("valid", valid_data)):
            losses = []
            for _ in range(config.eval_batches):
                x, y = get_batch(
                    data,
                    config.batch_size,
                    model.config.context_length,
                    device,
                    eval_generator,
                )
                _, loss = model(x, y)
                losses.append(loss.detach().cpu())
            result[name] = torch.stack(losses).mean().item()
        return result
    finally:
        model.train(was_training)


def configure_optimizer(model, config):
    decay, no_decay = [], []
    for _, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        (decay if parameter.ndim >= 2 else no_decay).append(parameter)
    groups = [
        {"params": decay, "weight_decay": config.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(
        groups,
        lr=config.peak_lr,
        betas=(0.9, 0.95),
    )


def lr_at_step(step, total_steps, warmup_steps, peak_lr, min_lr):
    if not 1 <= step <= total_steps:
        raise ValueError("step must be inside the training range")
    if not 0 < warmup_steps < total_steps:
        raise ValueError("warmup_steps must be inside the training range")
    if step <= warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (peak_lr - min_lr) * cosine


def save_checkpoint(path, model, optimizer, step, model_config,
                    train_config, tokenizer_name, corpus_sha256,
                    train_generator):
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    cuda_rng = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
    mps_rng = (
        torch.mps.get_rng_state()
        if hasattr(torch, "mps") and torch.backends.mps.is_available()
        else None
    )
    torch.save({
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "step": step,
        "model_config": asdict(model_config),
        "train_config": asdict(train_config),
        "tokenizer_name": tokenizer_name,
        "corpus_sha256": corpus_sha256,
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state": cuda_rng,
        "mps_rng_state": mps_rng,
        "train_generator_state": train_generator.get_state(),
    }, temporary)
    temporary.replace(path)


def load_checkpoint(path, model, optimizer, model_config,
                    train_config, tokenizer_name, corpus_sha256,
                    train_generator, device):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    if checkpoint["model_config"] != asdict(model_config):
        raise ValueError("checkpoint model_config does not match")
    if checkpoint["train_config"] != asdict(train_config):
        raise ValueError("checkpoint train_config does not match")
    if checkpoint["tokenizer_name"] != tokenizer_name:
        raise ValueError("checkpoint tokenizer does not match")
    if checkpoint["corpus_sha256"] != corpus_sha256:
        raise ValueError("checkpoint corpus does not match")
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    torch.set_rng_state(checkpoint["torch_rng_state"])
    train_generator.set_state(checkpoint["train_generator_state"])
    if device.type == "cuda" and checkpoint["cuda_rng_state"] is not None:
        torch.cuda.set_rng_state_all(checkpoint["cuda_rng_state"])
    if device.type == "mps" and checkpoint["mps_rng_state"] is not None:
        torch.mps.set_rng_state(checkpoint["mps_rng_state"])
    return int(checkpoint["step"])


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--output", type=Path,
                        default=Path("model/checkpoint.pt"))
    parser.add_argument("--resume", type=Path)
    return parser.parse_args()


def main():
    args = parse_args()
    train_config = TrainConfig()
    tokenizer_name = "cl100k_base"
    tokenizer = tiktoken.get_encoding(tokenizer_name)

    text = args.data.read_text(encoding="utf-8")
    corpus_sha256 = hashlib.sha256(text.encode("utf-8")).hexdigest()
    token_ids = tokenizer.encode(text, disallowed_special=())
    tokens = torch.tensor(token_ids, dtype=torch.long)  # keep corpus on CPU

    model_config = ModelConfig(
        vocab_size=tokenizer.n_vocab,
        context_length=128,
        d_model=80,
        n_layers=6,
        n_heads=4,
        dropout=0.1,
    )
    train_data, valid_data = split_tokens(
        tokens, train_fraction=0.9,
        context_length=model_config.context_length,
    )

    device = select_device()
    torch.manual_seed(train_config.seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(train_config.seed)

    model = Model(model_config).to(device)
    optimizer = configure_optimizer(model, train_config)
    train_generator = torch.Generator().manual_seed(train_config.seed)
    start_step = 0

    if args.resume is not None:
        start_step = load_checkpoint(
            args.resume, model, optimizer, model_config,
            train_config, tokenizer_name, corpus_sha256,
            train_generator, device,
        )

    model.train()
    if start_step == 0:
        metrics = estimate_loss(
            model, train_data, valid_data, train_config, device
        )
        print(f"step=0 train={metrics['train']:.4f} "
              f"valid={metrics['valid']:.4f}")

    for step in range(start_step + 1, train_config.total_steps + 1):
        lr = lr_at_step(
            step,
            train_config.total_steps,
            train_config.warmup_steps,
            train_config.peak_lr,
            train_config.min_lr,
        )
        for group in optimizer.param_groups:
            group["lr"] = lr

        x, y = get_batch(
            train_data,
            train_config.batch_size,
            model_config.context_length,
            device,
            train_generator,
        )
        optimizer.zero_grad(set_to_none=True)
        _, loss = model(x, y)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(
            model.parameters(), train_config.grad_clip
        )
        optimizer.step()

        should_evaluate = (
            step % train_config.eval_interval == 0
            or step == train_config.total_steps
        )
        if should_evaluate:
            metrics = estimate_loss(
                model, train_data, valid_data, train_config, device
            )
            print(
                f"step={step} lr={lr:.2e} grad={grad_norm.item():.3f} "
                f"train={metrics['train']:.4f} valid={metrics['valid']:.4f}"
            )
            save_checkpoint(
                args.output, model, optimizer, step,
                model_config, train_config, tokenizer_name, corpus_sha256,
                train_generator,
            )


if __name__ == "__main__":
    main()

19.11 訓練が正しく進んでいるかを見る

「Lossは必ず10.8から2.8へ下がる」という綺麗な曲線を先に作ってはいけません。次を確認します。

Step 0のLossが有限か;
optimizer.step() 後にParameterが本当に変わったか;
Training lossがWindow全体で下向きか。毎Step単調である必要はない;
Validation lossも改善しているか;
Gradient normが頻繁にSpikeしていないか;
Resume後にStep、LR、Optimizer state、Batch RNGが接続しているか。

Training lossが下がりValidation lossが上がるのが典型的なGeneralization gapの拡大です。Training loss自体が下がってから上がる場合は、LR、Data順、Checkpoint復元、数値問題も疑います。

章末チェックリスト

TargetがInputを1 Tokenずらしたものだと説明できる
本物のCSV parserで対象Text columnを選べる
CorpusはCPUに置き、現在Batchだけを移動できる
Train modeとRNGを汚さず評価できる
AdamW group、Warmup + Cosine、Gradient clippingを使える
訓練再開に十分な状態を保存・検証できる

次章予告

今あるのは単なるWeight fileではありません。完了Step、Optimizer state、Config、Corpus identityを持つCheckpointです。

第20章では inference.py を作り、信頼できるCheckpointからModelを再構築し、Eval modeへ切り替え、Promptをencodeし、明示したSampling境界の中で生成し、結果をdecodeします。