File size: 6,252 Bytes
a311e94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """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",
]
|