pretrain.py

import argparse
from dataclasses import asdict
import json
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

import torch
import torch.distributed as dist

from nano_transformer import TransformerLM
from nano_transformer.training import (
    Trainer,
    TrainerConfig,
    TrainingSchedule,
    TrainingStage,
    available_datasets,
    default_document_ids_path,
    get_dataset_spec,
    prepare_pretraining_data,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Pretrain a Nano-Transformer language model.")
    parser.add_argument("--dataset", choices=available_datasets(), default="tinystories")
    parser.add_argument("--data-dir", type=Path, default=ROOT / "data")
    parser.add_argument("--train-text-path", type=Path)
    parser.add_argument("--valid-text-path", type=Path)
    parser.add_argument("--vocab-path", type=Path)
    parser.add_argument("--merges-path", type=Path)
    parser.add_argument("--train-cache-path", type=Path)
    parser.add_argument("--valid-cache-path", type=Path)
    parser.add_argument("--special-token", action="append", dest="special_tokens")
    parser.add_argument("--force-rebuild-cache", action="store_true")
    parser.add_argument("--out-dir", type=Path, default=ROOT / "out" / "pretrain")
    parser.add_argument("--resume-from", type=Path)

    parser.add_argument("--context-length", type=int, default=128)
    parser.add_argument("--context-stage", action="append", default=[])
    parser.add_argument("--training-stage", action="append", default=[])
    parser.add_argument("--d-model", type=int, default=256)
    parser.add_argument("--d-ff", type=int)
    parser.add_argument("--num-layers", type=int, default=4)
    parser.add_argument("--num-heads", type=int, default=8)
    parser.add_argument("--num-kv-heads", type=int)
    parser.add_argument("--rope-theta", type=float, default=10000.0)
    parser.add_argument("--rope-scaling-type", choices=("linear", "yarn"))
    parser.add_argument("--rope-scaling-factor", type=float, default=1.0)
    parser.add_argument("--rope-base-context-length", type=int)
    parser.add_argument("--attention-window-size", type=int)
    parser.add_argument("--attention-global-tokens", type=int, default=0)
    parser.add_argument("--qk-norm-eps", type=float, default=1e-6)
    parser.add_argument("--no-qk-norm", action="store_false", dest="use_qk_norm")
    parser.add_argument("--attn-logit-softcap", type=float)
    parser.add_argument("--final-logit-softcap", type=float)
    parser.add_argument("--zero-init-residual", action="store_true", dest="zero_init_residual")
    parser.add_argument("--no-zero-init-residual", action="store_false", dest="zero_init_residual")
    parser.add_argument("--num-experts", type=int)
    parser.add_argument("--num-experts-per-token", type=int, default=2)
    parser.add_argument("--num-shared-experts", type=int, default=0)
    parser.add_argument("--moe-layers", type=str)
    parser.add_argument("--moe-capacity-factor", type=float)
    parser.add_argument("--moe-capacity-factor-eval", type=float)
    parser.add_argument("--moe-min-capacity", type=int, default=4)
    parser.add_argument("--moe-jitter-eps", type=float, default=0.0)
    parser.add_argument("--moe-aux-loss-coef", type=float, default=0.0)
    parser.add_argument("--moe-z-loss-coef", type=float, default=0.0)
    parser.add_argument("--no-moe-router-fp32", action="store_false", dest="moe_router_fp32")
    parser.add_argument("--moe-router-balance-strategy", choices=("aux_loss", "bias", "none"), default="aux_loss")
    parser.add_argument("--moe-router-bias-update-rate", type=float, default=0.01)

    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--batch-size-stage", action="append", default=[])
    parser.add_argument("--mtp-weights", type=str, default="1.0")
    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
    parser.add_argument("--document-boundary-sampling-probability", type=float, default=0.5)
    parser.add_argument("--max-document-length", type=int)
    parser.add_argument("--max-iters", type=int, default=1000)
    parser.add_argument("--learning-rate", type=float, default=3e-4)
    parser.add_argument("--min-learning-rate", type=float, default=3e-5)
    parser.add_argument("--warmup-iters", type=int, default=100)
    parser.add_argument("--lr-decay-iters", type=int)
    parser.add_argument("--weight-decay", type=float, default=0.1)
    parser.add_argument("--beta1", type=float, default=0.9)
    parser.add_argument("--beta2", type=float, default=0.95)
    parser.add_argument("--eps", type=float, default=1e-8)
    parser.add_argument("--optimizer", choices=("adamw", "muon"), default="adamw")
    parser.add_argument("--cautious-weight-decay", action="store_true")
    parser.add_argument("--muon-momentum", type=float, default=0.95)
    parser.add_argument("--muon-nesterov", action="store_true", dest="muon_nesterov")
    parser.add_argument("--no-muon-nesterov", action="store_false", dest="muon_nesterov")
    parser.add_argument("--muon-ns-steps", type=int, default=5)
    parser.add_argument("--muon-eps", type=float, default=1e-8)
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--eval-interval", type=int, default=100)
    parser.add_argument("--eval-iters", type=int, default=10)
    parser.add_argument("--log-interval", type=int, default=10)
    parser.add_argument("--checkpoint-interval", type=int, default=500)
    parser.add_argument("--device", type=str, default=default_device())
    parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32")
    parser.add_argument("--seed", type=int, default=1337)
    parser.add_argument("--compile", action="store_true", dest="compile_model")
    parser.add_argument("--pin-memory", action="store_true")
    parser.add_argument("--no-tf32", action="store_false", dest="use_tf32")
    parser.add_argument("--no-fused-optimizer", action="store_false", dest="use_fused_optimizer")
    parser.add_argument("--tie-embeddings", action="store_true", dest="tie_embeddings")
    parser.add_argument("--no-tie-embeddings", action="store_false", dest="tie_embeddings")
    parser.add_argument("--scale-embeddings", action="store_true", dest="scale_embeddings")
    parser.add_argument("--no-scale-embeddings", action="store_false", dest="scale_embeddings")
    parser.add_argument("--embedding-init-std", type=float)
    parser.add_argument("--no-scale-lr-with-batch-size", action="store_false", dest="scale_lr_with_batch_size")
    parser.add_argument("--no-schedule-weight-decay-with-lr", action="store_false", dest="schedule_weight_decay_with_lr")
    parser.add_argument("--loss-spike-threshold", type=float, default=2.0)
    parser.add_argument("--grad-norm-spike-threshold", type=float, default=2.0)
    parser.add_argument("--ddp-backend", choices=("nccl", "gloo"))
    parser.set_defaults(
        moe_router_fp32=True,
        use_tf32=True,
        use_fused_optimizer=True,
        tie_embeddings=True,
        scale_embeddings=True,
        use_qk_norm=True,
        zero_init_residual=True,
        scale_lr_with_batch_size=True,
        schedule_weight_decay_with_lr=True,
        muon_nesterov=True,
    )
    return parser.parse_args()


def default_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def resolve_paths(args: argparse.Namespace) -> tuple[Path, Path, Path, Path, Path, Path, list[str], str]:
    if any(
        value is not None
        for value in (
            args.train_text_path,
            args.valid_text_path,
            args.vocab_path,
            args.merges_path,
            args.train_cache_path,
            args.valid_cache_path,
        )
    ):
        required = {
            "train_text_path": args.train_text_path,
            "valid_text_path": args.valid_text_path,
            "vocab_path": args.vocab_path,
            "merges_path": args.merges_path,
        }
        missing = [name for name, value in required.items() if value is None]
        if missing:
            raise ValueError(f"Custom dataset mode requires {missing}.")

        train_cache_path = args.train_cache_path or default_cache_path(args.train_text_path)
        valid_cache_path = args.valid_cache_path or default_cache_path(args.valid_text_path)
        special_tokens = args.special_tokens or ["<|endoftext|>"]
        dataset_name = args.dataset
        return (
            args.train_text_path,
            args.valid_text_path,
            args.vocab_path,
            args.merges_path,
            train_cache_path,
            valid_cache_path,
            special_tokens,
            dataset_name,
        )

    spec = get_dataset_spec(args.dataset, data_dir=args.data_dir)
    special_tokens = args.special_tokens or list(spec.special_tokens)
    return (
        spec.train_text_path,
        spec.valid_text_path,
        spec.vocab_path,
        spec.merges_path,
        spec.train_cache_path,
        spec.valid_cache_path,
        special_tokens,
        spec.name,
    )


def default_cache_path(text_path: Path) -> Path:
    return text_path.with_name(f"{text_path.stem}-tokens.npy")


def parse_moe_layers(raw: str | None, num_layers: int) -> list[int] | None:
    if raw is None or raw.strip() == "":
        return None
    if raw.strip().lower() == "all":
        return list(range(num_layers))
    return [int(item.strip()) for item in raw.split(",") if item.strip()]


def parse_schedule(entries: list[str]) -> tuple[tuple[int, int], ...]:
    schedule: list[tuple[int, int]] = []
    for entry in entries:
        start, value = entry.split(":", maxsplit=1)
        schedule.append((int(start), int(value)))
    return tuple(sorted(schedule))


def parse_mtp_weights(raw: str | None) -> tuple[float, ...]:
    if raw is None or raw.strip() == "":
        return (1.0,)
    parts = [part.strip() for part in raw.split("|") if part.strip()]
    if not parts:
        return (1.0,)
    return tuple(float(part) for part in parts)


def parse_attention_window_override(raw: str | None) -> tuple[int | None, bool]:
    if raw is None:
        return None, False
    normalized = raw.strip().lower()
    if normalized in {"none", "null", "full", "off", "disabled"}:
        return None, True
    return int(raw), True


def parse_training_stage(entry: str) -> TrainingStage:
    parts = [part.strip() for part in entry.split(",") if part.strip()]
    if not parts:
        raise ValueError("training-stage entries may not be empty.")
    values: dict[str, str] = {}
    for part in parts:
        if "=" not in part:
            raise ValueError(
                "training-stage entries must use key=value pairs, "
                f"got {entry!r}."
            )
        key, value = part.split("=", maxsplit=1)
        values[key.strip().lower()] = value.strip()
    if "start" not in values:
        raise ValueError(f"training-stage entry must define start=..., got {entry!r}.")
    start_iteration = int(values.pop("start"))
    batch_size_raw = values.pop("batch", values.pop("batch_size", None))
    context_length_raw = values.pop("context", values.pop("context_length", None))
    attention_window_raw = values.pop("window", values.pop("attention_window_size", None))
    mtp_raw = values.pop("mtp", values.pop("mtp_weights", None))
    attention_window_size, override_attention_window_size = parse_attention_window_override(attention_window_raw)
    if values:
        raise ValueError(f"Unknown training-stage fields {sorted(values)} in {entry!r}.")
    return TrainingStage(
        start_iteration=start_iteration,
        batch_size=None if batch_size_raw is None else int(batch_size_raw),
        context_length=None if context_length_raw is None else int(context_length_raw),
        attention_window_size=attention_window_size,
        override_attention_window_size=override_attention_window_size,
        mtp_weights=() if mtp_raw is None else parse_mtp_weights(mtp_raw),
    )


def schedule_value_at_step(default: int, schedule: tuple[tuple[int, int], ...], step: int) -> int:
    resolved = default
    for start_iteration, value in schedule:
        if step >= start_iteration:
            resolved = value
        else:
            break
    return resolved


def max_scheduled_value(default: int, schedule: tuple[tuple[int, int], ...]) -> int:
    if not schedule:
        return default
    return max(default, *(value for _, value in schedule))


def merge_training_stages(
    *,
    context_length_schedule: tuple[tuple[int, int], ...],
    batch_size_schedule: tuple[tuple[int, int], ...],
    training_stages: tuple[TrainingStage, ...],
) -> tuple[TrainingStage, ...]:
    stage_updates: dict[int, dict[str, object]] = {}
    for start_iteration, value in context_length_schedule:
        stage_updates.setdefault(start_iteration, {})["context_length"] = value
    for start_iteration, value in batch_size_schedule:
        stage_updates.setdefault(start_iteration, {})["batch_size"] = value
    for stage in training_stages:
        stage_update = stage_updates.setdefault(stage.start_iteration, {})
        if stage.batch_size is not None:
            stage_update["batch_size"] = stage.batch_size
        if stage.context_length is not None:
            stage_update["context_length"] = stage.context_length
        if stage.override_attention_window_size:
            stage_update["override_attention_window_size"] = True
            stage_update["attention_window_size"] = stage.attention_window_size
        if stage.mtp_weights:
            stage_update["mtp_weights"] = stage.mtp_weights
    return tuple(
        TrainingStage(start_iteration=start_iteration, **stage_update)
        for start_iteration, stage_update in sorted(stage_updates.items())
    )


def default_num_kv_heads(num_heads: int) -> int:
    if num_heads <= 1:
        return 1
    for candidate in (4, 3, 2):
        if candidate < num_heads and num_heads % candidate == 0:
            return candidate
    return 1


def maybe_init_distributed(args: argparse.Namespace) -> tuple[bool, int, int, bool, str]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size <= 1:
        return False, 0, 1, True, args.device

    backend = args.ddp_backend or ("nccl" if torch.cuda.is_available() else "gloo")
    dist.init_process_group(backend=backend)
    rank = dist.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    if torch.cuda.is_available():
        torch.cuda.set_device(local_rank)
        device = f"cuda:{local_rank}"
    else:
        if args.device.startswith("cuda"):
            raise ValueError("Distributed training requested a CUDA device, but CUDA is not available.")
        device = "cpu"
        if rank == 0 and args.device != "cpu":
            print(f"Distributed training does not support {args.device} here; falling back to cpu.")
    return True, rank, dist.get_world_size(), rank == 0, device


def save_run_config(
    run_dir: Path,
    *,
    dataset_name: str,
    train_text_path: Path,
    valid_text_path: Path,
    vocab_path: Path,
    merges_path: Path,
    train_cache_path: Path,
    valid_cache_path: Path,
    train_document_ids_path: Path,
    valid_document_ids_path: Path,
    special_tokens: list[str],
    max_document_length: int | None,
    model: TransformerLM,
    trainer_config: TrainerConfig,
) -> None:
    run_config = {
        "dataset_name": dataset_name,
        "train_text_path": str(train_text_path),
        "valid_text_path": str(valid_text_path),
        "vocab_path": str(vocab_path),
        "merges_path": str(merges_path),
        "train_cache_path": str(train_cache_path),
        "valid_cache_path": str(valid_cache_path),
        "train_document_ids_path": str(train_document_ids_path),
        "valid_document_ids_path": str(valid_document_ids_path),
        "special_tokens": special_tokens,
        "max_document_length": max_document_length,
        "model_config": model.get_config(),
        "trainer_config": {
            **as_json_dict(trainer_config),
            "output_dir": str(trainer_config.output_dir),
        },
    }
    run_dir.mkdir(parents=True, exist_ok=True)
    with (run_dir / "run_config.json").open("w", encoding="utf-8") as f:
        json.dump(run_config, f, indent=2)


def as_json_dict(config: TrainerConfig) -> dict[str, object]:
    result = asdict(config)
    result["output_dir"] = str(config.output_dir)
    return result


def main() -> None:
    args = parse_args()
    distributed, rank, _, master_process, resolved_device = maybe_init_distributed(args)
    args.device = resolved_device

    try:
        torch.manual_seed(args.seed + rank)

        (
            train_text_path,
            valid_text_path,
            vocab_path,
            merges_path,
            train_cache_path,
            valid_cache_path,
            special_tokens,
            dataset_name,
        ) = resolve_paths(args)
        train_document_ids_path = default_document_ids_path(train_cache_path)
        valid_document_ids_path = default_document_ids_path(valid_cache_path)
        context_length_schedule = parse_schedule(args.context_stage)
        batch_size_schedule = parse_schedule(args.batch_size_stage)
        mtp_weights = parse_mtp_weights(args.mtp_weights)
        explicit_training_stages = tuple(parse_training_stage(entry) for entry in args.training_stage)
        merged_training_stages = merge_training_stages(
            context_length_schedule=context_length_schedule,
            batch_size_schedule=batch_size_schedule,
            training_stages=explicit_training_stages,
        )
        training_schedule = TrainingSchedule(
            batch_size=args.batch_size,
            context_length=args.context_length,
            attention_window_size=args.attention_window_size,
            mtp_weights=mtp_weights,
            stages=merged_training_stages,
        )
        initial_runtime_state = training_schedule.resolve(0)
        model_context_length = training_schedule.max_context_length()
        initial_context_length = initial_runtime_state.context_length
        initial_batch_size = initial_runtime_state.batch_size
        max_mtp_tokens = training_schedule.max_mtp_tokens()

        prepared = None
        if distributed and master_process:
            prepared = prepare_pretraining_data(
                train_text_path=train_text_path,
                valid_text_path=valid_text_path,
                vocab_path=vocab_path,
                merges_path=merges_path,
                train_cache_path=train_cache_path,
                valid_cache_path=valid_cache_path,
                train_document_ids_path=train_document_ids_path,
                valid_document_ids_path=valid_document_ids_path,
                special_tokens=special_tokens,
                force_rebuild=args.force_rebuild_cache,
                return_document_ids=True,
                max_document_length=args.max_document_length,
            )
        if distributed:
            dist.barrier()
        if prepared is None:
            prepared = prepare_pretraining_data(
                train_text_path=train_text_path,
                valid_text_path=valid_text_path,
                vocab_path=vocab_path,
                merges_path=merges_path,
                train_cache_path=train_cache_path,
                valid_cache_path=valid_cache_path,
                train_document_ids_path=train_document_ids_path,
                valid_document_ids_path=valid_document_ids_path,
                special_tokens=special_tokens,
                force_rebuild=args.force_rebuild_cache if not distributed else False,
                return_document_ids=True,
                max_document_length=args.max_document_length,
            )
        train_dataset, valid_dataset, tokenizer, train_document_ids, valid_document_ids = prepared

        if master_process:
            print(f"Loaded {len(train_dataset):,} train tokens and {len(valid_dataset):,} validation tokens.")

        d_ff = args.d_ff or (4 * args.d_model)
        moe_layers = parse_moe_layers(args.moe_layers, args.num_layers)
        num_kv_heads = args.num_kv_heads or default_num_kv_heads(args.num_heads)
        model = TransformerLM(
            vocab_size=len(tokenizer.vocab),
            context_length=model_context_length,
            d_model=args.d_model,
            num_layers=args.num_layers,
            num_heads=args.num_heads,
            num_kv_heads=num_kv_heads,
            d_ff=d_ff,
            rope_theta=args.rope_theta,
            scale_embeddings=args.scale_embeddings,
            embedding_init_std=args.embedding_init_std,
            attn_logit_softcap=args.attn_logit_softcap,
            final_logit_softcap=args.final_logit_softcap,
            rope_scaling_type=args.rope_scaling_type,
            rope_scaling_factor=args.rope_scaling_factor,
            rope_base_context_length=args.rope_base_context_length,
            attention_window_size=args.attention_window_size,
            attention_global_tokens=args.attention_global_tokens,
            use_qk_norm=args.use_qk_norm,
            qk_norm_eps=args.qk_norm_eps,
            num_experts=args.num_experts,
            num_experts_per_token=args.num_experts_per_token,
            moe_layers=moe_layers,
            tie_embeddings=args.tie_embeddings,
            zero_init_residual=args.zero_init_residual,
            max_mtp_tokens=max_mtp_tokens,
            num_shared_experts=args.num_shared_experts,
            moe_capacity_factor=args.moe_capacity_factor,
            moe_capacity_factor_eval=args.moe_capacity_factor_eval,
            moe_min_capacity=args.moe_min_capacity,
            moe_jitter_eps=args.moe_jitter_eps,
            moe_aux_loss_coef=args.moe_aux_loss_coef,
            moe_z_loss_coef=args.moe_z_loss_coef,
            moe_router_fp32=args.moe_router_fp32,
            moe_router_balance_strategy=args.moe_router_balance_strategy,
            moe_router_bias_update_rate=args.moe_router_bias_update_rate,
        )

        run_dir = args.out_dir / dataset_name
        config = TrainerConfig(
            batch_size=initial_batch_size,
            context_length=initial_context_length,
            max_iters=args.max_iters,
            learning_rate=args.learning_rate,
            min_learning_rate=args.min_learning_rate,
            warmup_iters=args.warmup_iters,
            lr_decay_iters=args.lr_decay_iters,
            weight_decay=args.weight_decay,
            beta1=args.beta1,
            beta2=args.beta2,
            eps=args.eps,
            grad_clip=args.grad_clip,
            eval_interval=args.eval_interval,
            eval_iters=args.eval_iters,
            log_interval=args.log_interval,
            checkpoint_interval=args.checkpoint_interval,
            device=args.device,
            dtype=args.dtype,
            seed=args.seed,
            output_dir=run_dir,
            gradient_accumulation_steps=args.gradient_accumulation_steps,
            compile_model=args.compile_model,
            pin_memory=args.pin_memory,
            use_tf32=args.use_tf32,
            use_fused_optimizer=args.use_fused_optimizer,
            context_length_schedule=context_length_schedule,
            batch_size_schedule=batch_size_schedule,
            mtp_weights=mtp_weights,
            training_stages=merged_training_stages,
            scale_lr_with_batch_size=args.scale_lr_with_batch_size,
            schedule_weight_decay_with_lr=args.schedule_weight_decay_with_lr,
            cautious_weight_decay=args.cautious_weight_decay,
            optimizer_name=args.optimizer,
            document_boundary_sampling_probability=args.document_boundary_sampling_probability,
            muon_momentum=args.muon_momentum,
            muon_nesterov=args.muon_nesterov,
            muon_ns_steps=args.muon_ns_steps,
            muon_eps=args.muon_eps,
            loss_spike_threshold=args.loss_spike_threshold,
            grad_norm_spike_threshold=args.grad_norm_spike_threshold,
        )
        trainer = Trainer(
            model,
            train_dataset,
            valid_dataset,
            config,
            train_document_ids=train_document_ids,
            valid_document_ids=valid_document_ids,
        )

        if master_process:
            save_run_config(
                run_dir,
                dataset_name=dataset_name,
                train_text_path=train_text_path,
                valid_text_path=valid_text_path,
                vocab_path=vocab_path,
                merges_path=merges_path,
                train_cache_path=train_cache_path,
                valid_cache_path=valid_cache_path,
                train_document_ids_path=train_document_ids_path,
                valid_document_ids_path=valid_document_ids_path,
                special_tokens=special_tokens,
                max_document_length=args.max_document_length,
                model=model,
                trainer_config=config,
            )

        if args.resume_from is not None:
            resumed_iteration = trainer.load_checkpoint(args.resume_from)
            if master_process:
                print(f"Resumed checkpoint from iteration {resumed_iteration}.")

        if master_process:
            print(f"Starting training on {args.device} with output dir {run_dir}...")
        history = trainer.train()
        if master_process:
            final_eval = next((entry for entry in reversed(history) if entry.get("split") == "eval"), None)
            if final_eval is not None:
                print(
                    "Finished training at iteration "
                    f"{trainer.iteration} with train_loss={final_eval['train_loss']:.4f}, "
                    f"valid_loss={final_eval['valid_loss']:.4f}."
                )
            else:
                print(f"Finished training at iteration {trainer.iteration}.")
    finally:
        if distributed and dist.is_initialized():
            dist.destroy_process_group()


if __name__ == "__main__":
    main()