train.py

"""
train.py — Main Pre-Training Entrypoint
=========================================

Launches GPT-2 pre-training (single-GPU or multi-GPU via torchrun).

Single GPU::

    python train.py

Multi-GPU (8×A100 on one node)::

    torchrun --standalone --nproc_per_node=8 train.py

Custom config::

    python train.py --scale medium --max_steps 10000 --compile

Argument Conventions:
    All arguments map 1-to-1 to fields in GPTConfig, TrainingConfig,
    DataConfig, or SystemConfig. This means the full training run is
    reproducible from the printed config alone.
"""

from __future__ import annotations

import argparse
import logging

from src.config.model_config import GPTConfig, ModelScale
from src.config.training_config import TrainingConfig
from src.config.data_config import DataConfig
from src.config.system_config import SystemConfig
from src.model.gpt import GPT
from src.data.dataloader import DataLoaderFactory
from src.training.trainer import Trainer
from src.training.mlflow_logger import MLflowLogger
from src.training.mlflow_logger import MLflowLogger
from src.training.checkpoint import CheckpointManager
from src.evaluation.hellaswag import HellaSwagEvaluator
from src.utils import setup_logging, VRAMProfiler


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="GPT-2 Pre-Training (NanoGPT-Production)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Model
    parser.add_argument(
        "--scale",
        type=str,
        default="small",
        choices=["small", "medium", "large", "xl"],
        help="Model scale preset.",
    )
    parser.add_argument("--n_layer", type=int, default=None)
    parser.add_argument("--n_head",  type=int, default=None)
    parser.add_argument("--n_embd",  type=int, default=None)
    parser.add_argument("--use_rope",    action="store_true", help="Rotary positional encoding.")
    parser.add_argument("--use_swiglu",  action="store_true", help="SwiGLU activation.")
    parser.add_argument("--use_rmsnorm", action="store_true", help="RMSNorm instead of LayerNorm.")

    # Fine-tune
    parser.add_argument(
        "--finetune", type=str, default=None,
        help="Initial checkpoint path. Only loads the model weights; the optimizer is reset."
    )

    # Training
    parser.add_argument("--max_steps",    type=int,   default=19073)
    parser.add_argument("--warmup_steps", type=int,   default=None,
                        help="Override warmup steps. Default: 715.")
    parser.add_argument("--max_lr",      type=float, default=6e-4)
    parser.add_argument("--min_lr",      type=float, default=None,
                        help="Min LR at end of cosine decay. Default: max_lr / 10.")
    parser.add_argument("--micro_batch", type=int,   default=16,     dest="micro_batch_size")
    parser.add_argument("--seq_len",     type=int,   default=1024,   dest="sequence_length")
    parser.add_argument("--total_batch", type=int,   default=524288, dest="total_batch_size")
    parser.add_argument("--dtype",       type=str,   default="bfloat16",
                        choices=["float32", "float16", "bfloat16"])
    parser.add_argument("--no_compile",  action="store_true")
    parser.add_argument("--no_resume",   action="store_true")

    # Data & paths
    parser.add_argument("--data_dir", type=str, default="data/fineweb_edu")
    parser.add_argument("--ckpt_dir", type=str, default="checkpoints")
    parser.add_argument("--log_dir",  type=str, default="logs")

    # Evaluation
    parser.add_argument("--mlflow",         action="store_true",
                        help="Enable MLflow experiment tracking.")
    parser.add_argument("--mlflow_exp",     type=str, default="nanogpt-production",
                        help="MLflow experiment name.")
    parser.add_argument("--eval_hellaswag", action="store_true")
    parser.add_argument("--hellaswag_max",  type=int, default=1000)

    return parser.parse_args()


def main() -> None:
    args = parse_args()

    # ── System config (auto-detect DDP environment) ───────────────────────────
    sys_cfg = SystemConfig.auto_detect()

    # ── Logging ──────────────────────────────────────────────────────────────
    setup_logging(log_dir=args.log_dir, rank=sys_cfg.global_rank)
    logger = logging.getLogger(__name__)

    # ── Model config ─────────────────────────────────────────────────────────
    scale_map = {
        "small":  ModelScale.SMALL,
        "medium": ModelScale.MEDIUM,
        "large":  ModelScale.LARGE,
        "xl":     ModelScale.XL,
    }
    scale = scale_map[args.scale]
    model_overrides = {
        k: v for k, v in [
            ("n_layer",      args.n_layer),
            ("n_head",       args.n_head),
            ("n_embd",       args.n_embd),
            ("use_rope",     True if args.use_rope    else None),
            ("use_swiglu",   True if args.use_swiglu  else None),
            ("use_rmsnorm",  True if args.use_rmsnorm else None),
            ("use_rope",     args.use_rope    or None),
            ("use_swiglu",   args.use_swiglu  or None),
            ("use_rmsnorm",  args.use_rmsnorm or None),
        ]
        if v is not None
    }
    model_cfg = GPTConfig.from_scale(scale, **model_overrides)

    # ── Training config ───────────────────────────────────────────────────────
    train_cfg = TrainingConfig(
        max_steps=args.max_steps,
        min_lr=args.min_lr if args.min_lr is not None else args.max_lr / 10,
        warmup_steps=args.warmup_steps if args.warmup_steps is not None else 715,
        max_lr=args.max_lr,
        micro_batch_size=args.micro_batch_size,
        sequence_length=args.sequence_length,
        total_batch_size=args.total_batch_size,
        dtype=args.dtype,
        compile_model=not args.no_compile,
        checkpoint_dir=args.ckpt_dir,
        log_dir=args.log_dir,
    )

    # ── Data config ───────────────────────────────────────────────────────────
    data_cfg = DataConfig(data_dir=args.data_dir)

    if sys_cfg.is_master:
        logger.info("=" * 70)
        logger.info("NanoGPT-Production — Pre-Training")
        logger.info("=" * 70)
        logger.info(f"Model:   {model_cfg}")
        logger.info(f"Device:  {sys_cfg.device} | GPUs: {sys_cfg.world_size}")
        logger.info(f"Flash:   {sys_cfg.flash_attention} | DDP: {sys_cfg.use_ddp}")
        logger.info(f"Compile: {train_cfg.compile_model} | dtype: {train_cfg.dtype}")
        logger.info(f"Data:    {data_cfg.data_dir}")
        logger.info(
            f"Schedule: {train_cfg.max_steps} steps, "
            f"warmup={train_cfg.warmup_steps}, "
            f"lr={train_cfg.max_lr:.1e}→{train_cfg.min_lr:.1e}"
        )
        vram = VRAMProfiler(sys_cfg.device)
        logger.info(f"VRAM before model init: {vram.summary()}")
        logger.info("=" * 70)

    # ── Component assembly ────────────────────────────────────────────────────
    data_factory = DataLoaderFactory(
        data_config=data_cfg,
        system_config=sys_cfg,
        B=train_cfg.micro_batch_size,
        T=train_cfg.sequence_length,
    )

    ckpt_manager = CheckpointManager(
        checkpoint_dir=train_cfg.checkpoint_dir,
        keep_last_n=3,
    )

    evaluator = None
    if args.eval_hellaswag and sys_cfg.is_master:
        evaluator = HellaSwagEvaluator(
            device=sys_cfg.device,
            max_examples=args.hellaswag_max,
        )

    # Pass fine-tune checkpoint path directly to trainer instance
    if args.finetune:
        import os
        _trainer_finetune = args.finetune
    else:
        _trainer_finetune = None

    mlflow_logger = MLflowLogger(
        experiment_name=args.mlflow_exp,
        enabled=args.mlflow,
    )
    mlflow_logger = MLflowLogger(
        experiment_name=args.mlflow_exp,
        enabled=args.mlflow,
    )
    trainer = Trainer(
        model_config=model_cfg,
        train_config=train_cfg,
        system_config=sys_cfg,
        data_factory=data_factory,
        checkpoint_manager=ckpt_manager,
        evaluator=evaluator,
        resume=not args.no_resume,
        mlflow_logger=mlflow_logger,
    )

    import torch
    if args.finetune:
        print(f"[Fine-tune] Initial weights: {args.finetune}")
        ckpt = torch.load(args.finetune, map_location="cpu", weights_only=False)
        # Install directly, not through CheckpointManager
        # We are bypassing the resume mechanism inside Trainer._setup()
        import os
        os.environ["NANOGPT_FINETUNE_CKPT"] = args.finetune

    # ── Run ───────────────────────────────────────────────────────────────────
    # Attach fine-tune path before training starts
    if _trainer_finetune:
        trainer._finetune_path = _trainer_finetune

    trainer.train()


if __name__ == "__main__":
    main()