Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """Shared command-line plumbing for the Agron training scripts. | |
| Both ``train_from_scratch.py`` and ``train_agron.py`` expose the same | |
| architecture / optimiser / algorithm / scheduler surface through | |
| :func:`add_training_args` and :func:`trainer_from_args`. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import torch | |
| from ..algorithms import algorithm_names | |
| from ..architectures import ARCHITECTURES, build_architecture | |
| from ..optimizers import optimizer_names | |
| from ..schedulers import scheduler_names | |
| from ..utils import count_parameters, format_parameters, read_json | |
| from .engine import AgronTrainer, TrainerConfig | |
| from .recipes import get_recipe, recipe_names | |
| def add_training_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: | |
| """Add the shared model/training flags to ``parser``.""" | |
| g = parser.add_argument_group("agron_nets") | |
| g.add_argument("--arch", choices=sorted(ARCHITECTURES), default=None, | |
| help="architecture from the Agron zoo (default: llama)") | |
| g.add_argument("--optimizer", choices=optimizer_names(), default=None, | |
| help="optimiser (default: from --recipe, else adamw)") | |
| g.add_argument("--scheduler", choices=scheduler_names() + ["none"], default=None, | |
| help="learning-rate schedule (default: from --recipe, else cosine)") | |
| g.add_argument("--algorithm", action="append", default=None, metavar="NAME", | |
| help=f"composable algorithm, repeatable (choices: {', '.join(algorithm_names())})") | |
| g.add_argument("--recipe", choices=recipe_names(), default=None, | |
| help="named optimiser+algorithm bundle (see agron_nets.training.RECIPES)") | |
| g.add_argument("--lr", type=float, default=None, help="peak learning rate") | |
| g.add_argument("--warmup", type=int, default=0, help="warmup steps") | |
| g.add_argument("--grad-accum", type=int, default=1, help="gradient accumulation micro-steps") | |
| g.add_argument("--grad-clip", type=float, default=1.0, help="global grad-norm clip (0 disables)") | |
| g.add_argument("--weight-decay", type=float, default=None) | |
| g.add_argument("--batch-size", type=int, default=8) | |
| g.add_argument("--log-every", type=int, default=25) | |
| g.add_argument("--save-every", type=int, default=0) | |
| g.add_argument("--device", default=None, help="cpu / cuda / cuda:1 / mps") | |
| g.add_argument("--seed", type=int, default=42) | |
| g.add_argument("--config-json", default=None, | |
| help="JSON string or path with architecture overrides") | |
| g.add_argument("--sophia-hessian-every", type=int, default=0, | |
| help="refresh Sophia's diagonal Hessian every N steps") | |
| return parser | |
| def arch_overrides(config_json: Optional[str]) -> Dict[str, Any]: | |
| if not config_json: | |
| return {} | |
| if os.path.exists(config_json): | |
| data = read_json(config_json, default={}) or {} | |
| else: | |
| data = json.loads(config_json) | |
| return {k: v for k, v in data.items() if k not in ("model_type", "architectures", "transformers_version")} | |
| def build_model_from_args( | |
| args: argparse.Namespace, | |
| preset: Optional[Dict[str, Any]] = None, | |
| model: Optional[torch.nn.Module] = None, | |
| ) -> Tuple[torch.nn.Module, Any]: | |
| """Build (or accept) the model described by ``args``.""" | |
| if model is not None: | |
| return model, model.config | |
| cfg = dict(preset or {}) | |
| cfg.update(arch_overrides(getattr(args, "config_json", None))) | |
| model, config = build_architecture(getattr(args, "arch", None) or "llama", **cfg) | |
| print(f"Architecture: {config.model_type} | {format_parameters(count_parameters(model))} params") | |
| return model, config | |
| def trainer_from_args( | |
| args: argparse.Namespace, | |
| model: torch.nn.Module, | |
| output_dir: Optional[str] = None, | |
| ) -> AgronTrainer: | |
| """Resolve recipe + explicit flags into a configured :class:`AgronTrainer`.""" | |
| recipe = get_recipe(args.recipe) if getattr(args, "recipe", None) else {} | |
| optimizer = args.optimizer or recipe.get("optimizer", "adamw") | |
| scheduler = args.scheduler or recipe.get("scheduler", "cosine") | |
| algorithms: List[Any] = list(recipe.get("algorithms", [])) | |
| extra = [] | |
| for item in getattr(args, "algorithm", None) or []: | |
| extra.extend(part.strip() for part in item.split(",") if part.strip()) | |
| for name in extra: | |
| if name not in algorithms: | |
| algorithms.append(name) | |
| if getattr(args, "grad_accum", 1) > 1 and "gradient_accumulation" not in algorithms: | |
| algorithms.append("gradient_accumulation") | |
| if getattr(args, "sophia_hessian_every", 0) > 0 and "sophia" != optimizer: | |
| print("warning: --sophia-hessian-every only applies to --optimizer sophia") | |
| config = TrainerConfig( | |
| max_steps=args.steps, | |
| epochs=getattr(args, "epochs", 1), | |
| warmup_steps=args.warmup, | |
| log_every=args.log_every, | |
| save_every=args.save_every, | |
| output_dir=output_dir or args.output, | |
| grad_clip=(args.grad_clip or None), | |
| seed=args.seed, | |
| weight_decay=args.weight_decay, | |
| sophia_hessian_every=args.sophia_hessian_every, | |
| ) | |
| optimizer_kwargs: Dict[str, Any] = {} | |
| if args.lr is not None: | |
| optimizer_kwargs["lr"] = args.lr | |
| trainer = AgronTrainer( | |
| model, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| algorithms=algorithms, | |
| config=config, | |
| device=args.device, | |
| optimizer_kwargs=optimizer_kwargs, | |
| ) | |
| # ``gradient_accumulation`` must know how many micro-steps to accumulate. | |
| if getattr(args, "grad_accum", 1) > 1: | |
| from ..algorithms.stability import GradientAccumulation | |
| for i, algo in enumerate(trainer.algorithms): | |
| if algo.name == "gradient_accumulation": | |
| trainer.algorithms[i] = GradientAccumulation(steps=args.grad_accum) | |
| print( | |
| f"Optimiser: {type(trainer.optimizer).__name__} | Scheduler: {scheduler} | " | |
| f"Algorithms: {', '.join(a.name for a in trainer.algorithms) or 'none'}" | |
| ) | |
| return trainer | |
| __all__ = [ | |
| "add_training_args", | |
| "arch_overrides", | |
| "build_model_from_args", | |
| "trainer_from_args", | |
| ] | |