train.py

import argparse
import json
import sys
from dataclasses import dataclass

import torch

from dlhub.checkpoint import save_checkpoint
from dlhub.config import append_jsonl, dataclass_to_dict, write_json
from dlhub.device import resolve_device
from dlhub.logging import get_logger
from dlhub.paths import build_run_paths
from dlhub.seed import set_seed
from dlhub.training.loop import evaluate_token_classifier, fit_token_classifier

from .data import DataConfig, get_dataloaders
from .model import ModelConfig, ToyMambaLM


@dataclass(frozen=True)
class TrainConfig:
    epochs: int = 5
    learning_rate: float = 2e-3
    seed: int = 42
    device: str = "auto"
    max_train_batches: int | None = None
    max_eval_batches: int | None = None
    run_name: str = "dev"

    embed_dim: int = 128
    state_dim: int = 64
    num_layers: int = 2
    expansion_factor: int = 2
    dropout: float = 0.1


def parse_args() -> tuple[TrainConfig, DataConfig]:
    parser = argparse.ArgumentParser(
        description="Lesson 03 (LLM): toy language model with a simplified selective state-space block."
    )

    parser.add_argument("--num-samples", type=int, default=4096)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--seq-length", type=int, default=64)
    parser.add_argument("--base-vocab-size", type=int, default=64)
    parser.add_argument("--val-fraction", type=float, default=0.2)
    parser.add_argument("--data-seed", type=int, default=0)

    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--learning-rate", type=float, default=2e-3)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", type=str, default="auto")

    parser.add_argument("--embed-dim", type=int, default=128)
    parser.add_argument("--state-dim", type=int, default=64)
    parser.add_argument("--num-layers", type=int, default=2)
    parser.add_argument("--expansion-factor", type=int, default=2)
    parser.add_argument("--dropout", type=float, default=0.1)

    parser.add_argument("--max-train-batches", type=int, default=None)
    parser.add_argument("--max-eval-batches", type=int, default=None)
    parser.add_argument("--run-name", type=str, default="dev")

    args = parser.parse_args()

    train_cfg = TrainConfig(
        epochs=args.epochs,
        learning_rate=args.learning_rate,
        seed=args.seed,
        device=args.device,
        max_train_batches=args.max_train_batches,
        max_eval_batches=args.max_eval_batches,
        run_name=args.run_name,
        embed_dim=args.embed_dim,
        state_dim=args.state_dim,
        num_layers=args.num_layers,
        expansion_factor=args.expansion_factor,
        dropout=args.dropout,
    )
    data_cfg = DataConfig(
        num_samples=args.num_samples,
        batch_size=args.batch_size,
        seq_length=args.seq_length,
        base_vocab_size=args.base_vocab_size,
        val_fraction=args.val_fraction,
        seed=args.data_seed,
        num_workers=0,
    )
    return train_cfg, data_cfg


@torch.no_grad()
def _greedy_generate(
    *,
    model: ToyMambaLM,
    device: torch.device,
    prompt: list[int],
    max_new_tokens: int,
    pad_id: int,
) -> list[int]:
    model.eval()
    max_length = int(model.cfg.max_length)

    ids = prompt[:]
    for _ in range(int(max_new_tokens)):
        cur = ids[-max_length:]
        padded = cur + [int(pad_id)] * max(0, max_length - len(cur))
        mask = [1.0] * min(max_length, len(cur)) + [0.0] * max(0, max_length - len(cur))
        inputs = {
            "input_ids": torch.tensor([padded], dtype=torch.long, device=device),
            "attention_mask": torch.tensor([mask], dtype=torch.float32, device=device),
        }
        logits = model(inputs)
        pos = min(len(cur), max_length) - 1
        ids.append(int(logits[0, pos].argmax(dim=-1).item()))
    return ids


def _write_samples(*, model: ToyMambaLM, device: torch.device, vocab, out_path, epoch: int) -> None:
    prompt = [1, 2, 3, 4, 5]
    generated = _greedy_generate(
        model=model,
        device=device,
        prompt=prompt,
        max_new_tokens=20,
        pad_id=vocab.pad_id,
    )
    with open(out_path, "a", encoding="utf-8") as handle:
        handle.write(
            json.dumps({"epoch": int(epoch), "prompt_ids": prompt, "gen_ids": generated}) + "\n"
        )


def run_training(train_cfg: TrainConfig, data_cfg: DataConfig) -> int:
    set_seed(train_cfg.seed)
    device_info = resolve_device(train_cfg.device)

    paths = build_run_paths(
        track="llm", lesson="lesson_03_toy_mamba_language_model", run_name=train_cfg.run_name
    )
    logger = get_logger("llm.toy_mamba_lm", log_file=paths.logs_dir / "train.log")
    paths.run_dir.mkdir(parents=True, exist_ok=True)
    paths.checkpoints_dir.mkdir(parents=True, exist_ok=True)

    logger.info("Device: %s (%s)", device_info.name, device_info.torch_device)
    logger.info("Outputs: %s", paths.run_dir)

    train_loader, val_loader, vocab = get_dataloaders(data_cfg)
    model = ToyMambaLM(
        ModelConfig(
            vocab_size=vocab.size,
            pad_id=vocab.pad_id,
            max_length=int(data_cfg.seq_length),
            embed_dim=int(train_cfg.embed_dim),
            state_dim=int(train_cfg.state_dim),
            num_layers=int(train_cfg.num_layers),
            expansion_factor=int(train_cfg.expansion_factor),
            dropout=float(train_cfg.dropout),
        )
    ).to(device_info.torch_device)

    write_json(
        paths.run_dir / "config.json",
        {
            "train": dataclass_to_dict(train_cfg),
            "data": dataclass_to_dict(data_cfg),
            "versions": {"python": sys.version, "torch": torch.__version__},
        },
    )
    write_json(paths.run_dir / "vocab.json", vocab.to_dict())

    criterion = torch.nn.CrossEntropyLoss(ignore_index=int(vocab.pad_id))
    optimizer = torch.optim.Adam(model.parameters(), lr=float(train_cfg.learning_rate))

    metrics_path = paths.run_dir / "metrics.jsonl"
    samples_path = paths.run_dir / "samples.jsonl"
    for epoch in range(1, int(train_cfg.epochs) + 1):
        train_stats = fit_token_classifier(
            model=model,
            loader=train_loader,
            optimizer=optimizer,
            criterion=criterion,
            device=device_info.torch_device,
            max_batches=train_cfg.max_train_batches,
            ignore_index=int(vocab.pad_id),
        )
        eval_stats = evaluate_token_classifier(
            model=model,
            loader=val_loader,
            criterion=criterion,
            device=device_info.torch_device,
            max_batches=train_cfg.max_eval_batches,
            ignore_index=int(vocab.pad_id),
        )

        _write_samples(
            model=model,
            device=device_info.torch_device,
            vocab=vocab,
            out_path=samples_path,
            epoch=epoch,
        )
        logger.info(
            "Epoch %d/%d | train loss %.4f acc %.3f | eval loss %.4f acc %.3f",
            epoch,
            train_cfg.epochs,
            train_stats.loss,
            train_stats.accuracy,
            eval_stats.loss,
            eval_stats.accuracy,
        )
        append_jsonl(
            metrics_path,
            {
                "epoch": epoch,
                "train_loss": train_stats.loss,
                "train_acc": train_stats.accuracy,
                "eval_loss": eval_stats.loss,
                "eval_acc": eval_stats.accuracy,
                "lr": optimizer.param_groups[0]["lr"],
            },
        )

    ckpt_path = save_checkpoint(
        paths.checkpoints_dir / "checkpoint.pt",
        model=model,
        optimizer=optimizer,
        epoch=int(train_cfg.epochs),
        extra={
            "track": "llm",
            "lesson": "lesson_03_toy_mamba_language_model",
            "vocab_size": vocab.size,
        },
    )
    logger.info("Saved checkpoint to %s", ckpt_path)
    return 0


def main() -> int:
    if __package__ is None:
        raise RuntimeError(
            "Please run this lesson from the repo root as a module:\n"
            "  python -m tracks.llm.lesson_03_toy_mamba_language_model.train"
        )
    train_cfg, data_cfg = parse_args()
    return run_training(train_cfg, data_cfg)


if __name__ == "__main__":
    raise SystemExit(main())