Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """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 | |
| 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) | |
| 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 | |
| # ------------------------------------------------------------------ # | |
| 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", | |
| ] | |