test_scripts.py

import argparse
from dataclasses import asdict
import importlib.util
import json
import math
import subprocess
import sys
from pathlib import Path

import torch

from nano_transformer import Tokenizer, TransformerLM
from nano_transformer.training import (
    TrainerConfig,
    create_optimizer,
    default_document_ids_path,
    prepare_pretraining_data,
    save_checkpoint,
)


ROOT = Path(__file__).resolve().parent.parent
FIXTURES = Path(__file__).resolve().parent / "fixtures"
PRETRAIN_SCRIPT = ROOT / "scripts" / "pretrain.py"
SAMPLE_SCRIPT = ROOT / "scripts" / "sample.py"
EVAL_SCRIPT = ROOT / "scripts" / "eval.py"


def load_script_module(script_path: Path, module_name: str):
    spec = importlib.util.spec_from_file_location(module_name, script_path)
    if spec is None or spec.loader is None:
        raise AssertionError(f"Unable to load script module from {script_path}.")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def build_sample_run_dir(tmp_path: Path) -> Path:
    tokenizer = Tokenizer.from_files(
        vocab_filepath=FIXTURES / "gpt2_vocab.json",
        merges_filepath=FIXTURES / "gpt2_merges.txt",
        special_tokens=["<|endoftext|>"],
    )
    model = TransformerLM(
        vocab_size=len(tokenizer.vocab),
        context_length=8,
        d_model=16,
        num_layers=2,
        num_heads=4,
        d_ff=32,
        rope_theta=10000.0,
    )
    optimizer = create_optimizer(
        model,
        optimizer_name="adamw",
        lr=3e-4,
        betas=(0.9, 0.95),
        eps=1e-8,
        weight_decay=0.1,
        device_type="cpu",
        use_fused=False,
    )
    run_dir = tmp_path / "run"
    run_dir.mkdir()
    save_checkpoint(model, optimizer, iteration=0, out=run_dir / "latest.pt")

    run_config = {
        "vocab_path": str(FIXTURES / "gpt2_vocab.json"),
        "merges_path": str(FIXTURES / "gpt2_merges.txt"),
        "special_tokens": ["<|endoftext|>"],
        "model_config": model.get_config(),
    }
    (run_dir / "run_config.json").write_text(json.dumps(run_config), encoding="utf-8")
    return run_dir


def run_sample_script(run_dir: Path, *args: str) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [sys.executable, str(SAMPLE_SCRIPT), "--run-dir", str(run_dir), *args],
        cwd=ROOT,
        capture_output=True,
        text=True,
        check=True,
    )


def build_eval_run_dir(tmp_path: Path) -> tuple[Path, Tokenizer, Path]:
    train_text = tmp_path / "train.txt"
    valid_text = tmp_path / "valid.txt"
    train_text.write_text("alpha\nbeta\n", encoding="utf-8")
    valid_text.write_text("alpha\nbeta\n", encoding="utf-8")

    train_cache = tmp_path / "train-tokens.npy"
    valid_cache = tmp_path / "valid-tokens.npy"
    train_doc_ids_path = default_document_ids_path(train_cache)
    valid_doc_ids_path = default_document_ids_path(valid_cache)
    _, _, tokenizer, _, _ = prepare_pretraining_data(
        train_text_path=train_text,
        valid_text_path=valid_text,
        vocab_path=FIXTURES / "gpt2_vocab.json",
        merges_path=FIXTURES / "gpt2_merges.txt",
        train_cache_path=train_cache,
        valid_cache_path=valid_cache,
        train_document_ids_path=train_doc_ids_path,
        valid_document_ids_path=valid_doc_ids_path,
        return_document_ids=True,
    )

    model = TransformerLM(
        vocab_size=len(tokenizer.vocab),
        context_length=8,
        d_model=16,
        num_layers=2,
        num_heads=4,
        d_ff=32,
        rope_theta=10000.0,
        num_experts=4,
        moe_aux_loss_coef=0.1,
        moe_z_loss_coef=0.1,
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    run_dir = tmp_path / "eval-run"
    run_dir.mkdir()
    save_checkpoint(model, optimizer, iteration=0, out=run_dir / "latest.pt")

    trainer_config = TrainerConfig(
        batch_size=2,
        context_length=4,
        max_iters=1,
        eval_iters=1,
        device="cpu",
        output_dir=run_dir,
    )
    serialized_config = asdict(trainer_config)
    serialized_config["output_dir"] = str(run_dir)
    (run_dir / "trainer_config.json").write_text(json.dumps(serialized_config, indent=2), encoding="utf-8")

    run_config = {
        "dataset_name": "fixture",
        "train_text_path": str(train_text),
        "valid_text_path": str(valid_text),
        "vocab_path": str(FIXTURES / "gpt2_vocab.json"),
        "merges_path": str(FIXTURES / "gpt2_merges.txt"),
        "train_cache_path": str(train_cache),
        "valid_cache_path": str(valid_cache),
        "train_document_ids_path": str(train_doc_ids_path),
        "valid_document_ids_path": str(valid_doc_ids_path),
        "special_tokens": ["<|endoftext|>"],
        "max_document_length": None,
        "model_config": model.get_config(),
        "trainer_config": serialized_config,
    }
    (run_dir / "run_config.json").write_text(json.dumps(run_config), encoding="utf-8")
    return run_dir, tokenizer, valid_text


def test_sample_script_handles_empty_prompt_with_special_token_seed(tmp_path):
    run_dir = build_sample_run_dir(tmp_path)

    result = run_sample_script(run_dir, "--max-new-tokens", "2")

    assert "<|endoftext|>" not in result.stdout


def test_sample_script_respects_num_samples_without_kv_cache(tmp_path):
    run_dir = build_sample_run_dir(tmp_path)

    result = run_sample_script(
        run_dir,
        "--prompt",
        "hi",
        "--max-new-tokens",
        "2",
        "--num-samples",
        "3",
        "--no-kv-cache",
    )

    assert result.stdout.count("=== Sample ") == 3


def test_eval_script_reports_ce_bpb_without_mutating_saved_config(tmp_path):
    run_dir, tokenizer, valid_text = build_eval_run_dir(tmp_path)
    trainer_config_path = run_dir / "trainer_config.json"
    original_trainer_config = trainer_config_path.read_text(encoding="utf-8")

    result = subprocess.run(
        [sys.executable, str(EVAL_SCRIPT), "--run-dir", str(run_dir), "--device", "cpu", "--eval-iters", "7"],
        cwd=ROOT,
        capture_output=True,
        text=True,
        check=True,
    )

    assert trainer_config_path.read_text(encoding="utf-8") == original_trainer_config

    output_lines = result.stdout.strip().splitlines()
    metrics = json.loads("\n".join(output_lines[:-1]))
    reported_bpb = float(output_lines[-1].split("=", maxsplit=1)[1])
    raw_lines = valid_text.read_text(encoding="utf-8").splitlines(keepends=True)
    raw_token_count = sum(len(tokenizer.encode(line)) for line in raw_lines)
    raw_byte_count = sum(len(line.encode("utf-8")) for line in raw_lines)
    expected_bpb = metrics["valid_ce_loss"] / math.log(2) * (raw_token_count / raw_byte_count)

    assert metrics["valid_loss"] > metrics["valid_ce_loss"]
    assert math.isclose(reported_bpb, expected_bpb, abs_tol=5e-5)


def test_maybe_init_distributed_falls_back_to_cpu_without_cuda(monkeypatch):
    pretrain = load_script_module(PRETRAIN_SCRIPT, "test_pretrain_script")
    backend_calls: list[str] = []

    monkeypatch.setenv("WORLD_SIZE", "2")
    monkeypatch.setenv("LOCAL_RANK", "0")
    monkeypatch.setattr(pretrain.torch.cuda, "is_available", lambda: False)
    monkeypatch.setattr(pretrain.dist, "init_process_group", lambda backend: backend_calls.append(backend))
    monkeypatch.setattr(pretrain.dist, "get_rank", lambda: 0)
    monkeypatch.setattr(pretrain.dist, "get_world_size", lambda: 2)

    distributed, rank, world_size, master_process, device = pretrain.maybe_init_distributed(
        argparse.Namespace(device="mps", ddp_backend=None)
    )

    assert distributed is True
    assert rank == 0
    assert world_size == 2
    assert master_process is True
    assert device == "cpu"
    assert backend_calls == ["gloo"]


def test_pretrain_schedule_helpers_support_staged_context_growth():
    pretrain = load_script_module(PRETRAIN_SCRIPT, "test_pretrain_schedule_helpers")
    schedule = pretrain.parse_schedule(["0:4", "10:8", "20:16"])

    assert pretrain.schedule_value_at_step(32, schedule, step=0) == 4
    assert pretrain.schedule_value_at_step(32, schedule, step=15) == 8
    assert pretrain.schedule_value_at_step(32, schedule, step=25) == 16
    assert pretrain.max_scheduled_value(32, schedule) == 32
    assert pretrain.max_scheduled_value(2, schedule) == 16


def test_parse_training_stage_can_restore_full_attention():
    pretrain = load_script_module(PRETRAIN_SCRIPT, "test_pretrain_restore_full_attention")
    stage = pretrain.parse_training_stage("start=10,window=full")

    assert stage.attention_window_size is None
    assert stage.override_attention_window_size is True


def test_default_num_kv_heads_returns_valid_divisors():
    pretrain = load_script_module(PRETRAIN_SCRIPT, "test_pretrain_gqa_defaults")

    expected = {
        1: 1,
        2: 1,
        3: 1,
        4: 2,
        6: 3,
        8: 4,
        10: 2,
        12: 4,
        14: 2,
        16: 4,
    }
    for num_heads, num_kv_heads in expected.items():
        assert pretrain.default_num_kv_heads(num_heads) == num_kv_heads