Blame for agron_nets/training/engine.py

"""A small, dependency-light training engine for the Agron architecture zoo.

The engine composes four independent axes:

* an **architecture** from :mod:`agron_nets.architectures`,
* an **optimiser** from :mod:`agron_nets.optimizers`,
* a **learning-rate schedule** from :mod:`agron_nets.schedulers`,
* any number of **algorithms** from :mod:`agron_nets.algorithms`.

Example
-------
>>> from agron_nets.training import AgronTrainer, TrainerConfig
>>> trainer = AgronTrainer(
...     model,
...     optimizer="muon",
...     scheduler="wsd",
...     algorithms=["gradient_clip", "ema", "z_loss"],
...     config=TrainerConfig(max_steps=1000, warmup_steps=50),
... )
>>> trainer.fit(train_loader)
"""

from __future__ import annotations

import os
import time
from dataclasses import asdict, dataclass
from typing import Any, Dict, Iterable, List, Optional

import torch
import torch.nn as nn

from ..algorithms import AlgorithmContext, BaseAlgorithm, build_algorithm
from ..algorithms.base import compose_loss
from ..optimizers import OPTIMIZERS, build_optimizer
from ..schedulers import build_scheduler
from ..utils import model_device, set_seed, write_json


@dataclass
class TrainerConfig:
    """Everything the training loop needs that is not a model component."""

    max_steps: int = 1000
    epochs: int = 1
    warmup_steps: int = 0
    log_every: int = 10
    save_every: int = 0
    output_dir: Optional[str] = None
    grad_clip: Optional[float] = 1.0
    seed: int = 42
    weight_decay: Optional[float] = None
    sophia_hessian_every: int = 0
    log_tokens_per_second: bool = True

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


def _batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
    return {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()}


def _truncate_batch(batch: Dict[str, torch.Tensor], max_len: int) -> Dict[str, torch.Tensor]:
    """Keep the last ``max_len`` positions so RoPE never runs out of cache."""
    if max_len <= 0:
        return batch
    out = dict(batch)
    for key in ("input_ids", "labels", "attention_mask", "position_ids"):
        value = out.get(key)
        if torch.is_tensor(value) and value.dim() >= 2 and value.shape[1] > max_len:
            out[key] = value[:, -max_len:]
    return out


def _extract_loss(outputs: Any) -> Optional[torch.Tensor]:
    if isinstance(outputs, dict):
        return outputs.get("loss")
    return getattr(outputs, "loss", None)


def _is_step_driver(algo: BaseAlgorithm) -> bool:
    """An algorithm that owns the forward/backward pass itself."""
    return type(algo).train_step is not BaseAlgorithm.train_step


class AgronTrainer:
    """Compose an architecture with an optimiser, schedule and algorithms."""

    def __init__(
        self,
        model: nn.Module,
        optimizer: Any = "adamw",
        scheduler: Any = "cosine",
        algorithms: Optional[Iterable[Any]] = None,
        config: Optional[TrainerConfig] = None,
        device: Optional[Any] = None,
        optimizer_kwargs: Optional[Dict[str, Any]] = None,
        scheduler_kwargs: Optional[Dict[str, Any]] = None,
        algorithm_kwargs: Optional[Dict[str, Any]] = None,
    ):
        self.model = model
        self.config = config or TrainerConfig()
        self.device = torch.device(device or model_device(model))

        params = [p for p in model.parameters() if p.requires_grad]
        if isinstance(optimizer, str):
            overrides = dict(optimizer_kwargs or {})
            lr = overrides.pop("lr", None)
            if lr is None:
                lr = OPTIMIZERS.get(optimizer.lower(), {}).get("defaults", {}).get("lr", 1e-3)
            self.optimizer = build_optimizer(optimizer, params, lr=lr, **overrides)
        else:
            self.optimizer = optimizer

        if self.config.weight_decay is not None:
            for group in self.optimizer.param_groups:
                group["weight_decay"] = self.config.weight_decay

        if isinstance(scheduler, str) and scheduler.lower() not in ("none", "constant_none"):
            self.scheduler = build_scheduler(
                self.optimizer,
                scheduler,
                total_steps=max(1, self.config.max_steps),
                warmup_steps=self.config.warmup_steps,
                **(scheduler_kwargs or {}),
            )
        else:
            self.scheduler = scheduler if scheduler not in ("none", "constant_none") else None

        if algorithms is None:
            algorithms = []
        elif isinstance(algorithms, str):
            algorithms = [a.strip() for a in algorithms.split(",") if a.strip()]
        self.algorithms: List[BaseAlgorithm] = [
            a if isinstance(a, BaseAlgorithm) else build_algorithm(a, **(algorithm_kwargs or {}))
            for a in algorithms
        ]

        if self.config.grad_clip is not None and not any(
            a.name == "gradient_clip" for a in self.algorithms
        ):
            from ..algorithms.stability import GradientClip

            self.algorithms.append(GradientClip(max_norm=self.config.grad_clip))

        self.step = 0
        self.history: List[Dict[str, float]] = []
        self._ctx: Optional[AlgorithmContext] = None

    # ------------------------------------------------------------------ #
    # Setup helpers
    # ------------------------------------------------------------------ #
    def _make_context(self) -> AlgorithmContext:
        ctx = AlgorithmContext(
            model=self.model,
            optimizer=self.optimizer,
            scheduler=self.scheduler,
            device=self.device,
            step=0,
            epoch=0,
            max_steps=self.config.max_steps,
        )
        ctx.set("loss_hooks", self.algorithms)
        ctx.set("algorithms", self.algorithms)
        return ctx

    def _call(self, name: str, *args) -> None:
        for algo in self.algorithms:
            fn = getattr(algo, name, None)
            if fn is not None:
                fn(self._ctx, *args)

    @property
    def current_lr(self) -> float:
        return float(self.optimizer.param_groups[0]["lr"])

    # ------------------------------------------------------------------ #
    # Training
    # ------------------------------------------------------------------ #
    def fit(
        self,
        train_dataloader,
        eval_dataloader=None,
        callbacks: Optional[List[Any]] = None,
    ) -> Dict[str, Any]:
        set_seed(self.config.seed)
        self.model.to(self.device)
        self.model.train()
        self._ctx = self._make_context()
        self._ctx.step = self.step
        self._call("on_train_begin")

        max_len = int(getattr(self.model.config, "max_position_embeddings", 0) or 0)
        drivers = [a for a in self.algorithms if _is_step_driver(a)]
        if len(drivers) > 1:
            raise ValueError(
                "only one step-owning algorithm can be active at a time, got: "
                + ", ".join(a.name for a in drivers)
            )
        driver = drivers[0] if drivers else None
        schedule_free = hasattr(self.optimizer, "train") and self.optimizer.__class__.__name__ == "ScheduleFreeAdamW"

        started = time.time()
        tokens_seen = 0
        epoch = 0
        stopped = False

        while epoch < self.config.epochs and not stopped:
            self._ctx.epoch = epoch
            self._call("on_epoch_begin")
            for batch in train_dataloader:
                if self.step >= self.config.max_steps:
                    stopped = True
                    break

                batch = _truncate_batch(_batch_to_device(batch, self.device), max_len)
                self._ctx.step = self.step
                self._ctx.set("batch", batch)

                if schedule_free:
                    self.optimizer.train()

                for algo in self.algorithms:
                    batch = algo.on_step_begin(self._ctx, batch)

                if driver is not None:
                    loss = driver.train_step(self._ctx, batch)
                else:
                    outputs = self.model(**batch)
                    loss = _extract_loss(outputs)
                    loss = compose_loss(self._ctx, batch, outputs, loss)
                    if loss is not None:
                        loss.backward()

                tokens_seen += int(batch["input_ids"].numel()) if "input_ids" in batch else 0

                if not all(algo.should_step(self._ctx) for algo in self.algorithms):
                    continue

                self._sophia_hessian_step(batch)

                for algo in self.algorithms:
                    algo.on_before_optimizer_step(self._ctx)

                grad_norm = self._grad_norm()
                self.optimizer.step()
                self.optimizer.zero_grad(set_to_none=True)
                if self.scheduler is not None:
                    self.scheduler.step()

                for algo in self.algorithms:
                    algo.on_step_end(self._ctx, loss)

                self.step += 1

                if self.config.log_every and self.step % self.config.log_every == 0:
                    self._log(loss, grad_norm, tokens_seen, started)
                for cb in callbacks or []:
                    cb(self)
                if self.config.save_every and self.step % self.config.save_every == 0:
                    self.save_checkpoint(os.path.join(self.config.output_dir or "agron-checkpoints", f"step-{self.step}"))

            if not stopped:
                epoch += 1

        self._call("on_train_end")
        self.model.eval()
        if self.config.output_dir:
            self.save_checkpoint(os.path.join(self.config.output_dir, "final"))
            write_json(os.path.join(self.config.output_dir, "train_report.json"), {
                "steps": self.step,
                "tokens": tokens_seen,
                "seconds": round(time.time() - started, 2),
                "optimizer": self.optimizer.__class__.__name__,
                "scheduler": self.scheduler.__class__.__name__ if self.scheduler is not None else None,
                "algorithms": [a.name for a in self.algorithms],
                "config": self.config.to_dict(),
                "history": self.history,
            })

        final_eval = None
        if eval_dataloader is not None:
            final_eval = self.evaluate(eval_dataloader)

        return {
            "steps": self.step,
            "tokens": tokens_seen,
            "history": self.history,
            "eval_loss": final_eval,
            "seconds": round(time.time() - started, 2),
        }

    # ------------------------------------------------------------------ #
    # Evaluation / inference helpers
    # ------------------------------------------------------------------ #
    @torch.no_grad()
    def evaluate(self, dataloader, max_batches: Optional[int] = None) -> float:
        was_training = self.model.training
        self.model.eval()
        if hasattr(self.optimizer, "eval") and self.optimizer.__class__.__name__ == "ScheduleFreeAdamW":
            self.optimizer.eval()
        total_loss, total_tokens = 0.0, 0
        max_len = int(getattr(self.model.config, "max_position_embeddings", 0) or 0)
        for i, batch in enumerate(dataloader):
            if max_batches is not None and i >= max_batches:
                break
            batch = _truncate_batch(_batch_to_device(batch, self.device), max_len)
            outputs = self.model(**batch)
            loss = _extract_loss(outputs)
            if loss is None:
                continue
            n = int(batch["labels"].numel()) if "labels" in batch else 1
            total_loss += float(loss.detach()) * n
            total_tokens += n
        if was_training:
            self.model.train()
            if hasattr(self.optimizer, "train") and self.optimizer.__class__.__name__ == "ScheduleFreeAdamW":
                self.optimizer.train()
        return total_loss / max(1, total_tokens)

    def _grad_norm(self) -> float:
        grads = [p.grad.detach() for p in self.model.parameters() if p.grad is not None]
        if not grads:
            return 0.0
        return float(torch.norm(torch.stack([g.norm(2) for g in grads]), 2))

    def _log(self, loss, grad_norm: float, tokens_seen: int, started: float) -> None:
        entry = {
            "step": self.step,
            "loss": float(loss.detach()) if torch.is_tensor(loss) else float("nan"),
            "lr": self.current_lr,
            "grad_norm": grad_norm,
            "tokens": tokens_seen,
        }
        elapsed = max(1e-9, time.time() - started)
        if self.config.log_tokens_per_second:
            entry["tokens_per_s"] = round(tokens_seen / elapsed, 1)
        self.history.append(entry)
        print(
            f"step {self.step:>6} | loss {entry['loss']:.4f} | lr {entry['lr']:.3e} "
            f"| grad {grad_norm:.3f} | {entry.get('tokens_per_s', 0):.0f} tok/s",
            flush=True,
        )

    def _sophia_hessian_step(self, batch) -> None:
        every = self.config.sophia_hessian_every
        if every <= 0 or not hasattr(self.optimizer, "update_hessian"):
            return
        if self.step % every != 0:
            return
        self.optimizer.zero_grad(set_to_none=True)
        with torch.enable_grad():
            outputs = self.model(**batch)
            loss = _extract_loss(outputs)
            if loss is not None and loss.requires_grad:
                self.optimizer.update_hessian(loss, seed=self.config.seed + self.step)
        self.optimizer.zero_grad(set_to_none=True)

    # ------------------------------------------------------------------ #
    # Checkpointing
    # ------------------------------------------------------------------ #
    def save_checkpoint(self, directory: str) -> str:
        os.makedirs(directory, exist_ok=True)
        self.model.save_pretrained(directory)
        state = {
            "step": self.step,
            "optimizer": self.optimizer.state_dict(),
            "scheduler": self.scheduler.state_dict() if self.scheduler is not None else None,
            "algorithms": [
                {"name": a.name, "state": a.state_dict()} for a in self.algorithms
            ],
            "config": self.config.to_dict(),
            "history": self.history,
        }
        torch.save(state, os.path.join(directory, "trainer_state.pt"))
        return directory

    def load_checkpoint(self, directory: str) -> None:
        path = os.path.join(directory, "trainer_state.pt")
        state = torch.load(path, map_location=self.device, weights_only=False)
        self.step = int(state.get("step", 0))
        if state.get("optimizer") is not None:
            self.optimizer.load_state_dict(state["optimizer"])
        if self.scheduler is not None and state.get("scheduler") is not None:
            self.scheduler.load_state_dict(state["scheduler"])
        algo_state = {item["name"]: item["state"] for item in state.get("algorithms", [])}
        for algo in self.algorithms:
            if algo.name in algo_state:
                algo.load_state_dict(algo_state[algo.name])
        self.history = state.get("history", [])


def train_architecture(
    architecture: str = "llama",
    config: Optional[Dict[str, Any]] = None,
    train_dataloader=None,
    eval_dataloader=None,
    optimizer: str = "adamw",
    scheduler: str = "cosine",
    algorithms: Optional[Iterable[str]] = None,
    trainer_config: Optional[TrainerConfig] = None,
    device: Optional[Any] = None,
) -> Dict[str, Any]:
    """One-liner: build an architecture, train it, return the report."""
    from ..architectures import build_architecture

    model, _ = build_architecture(architecture, **(config or {}))
    trainer = AgronTrainer(
        model,
        optimizer=optimizer,
        scheduler=scheduler,
        algorithms=algorithms,
        config=trainer_config,
        device=device,
    )
    return trainer.fit(train_dataloader, eval_dataloader)


__all__ = [
    "AgronTrainer",
    "TrainerConfig",
    "train_architecture",
]