Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """Training-engine and algorithm-composition tests.""" | |
| from __future__ import annotations | |
| import pytest | |
| import torch | |
| from agron_nets.algorithms import ALGORITHMS, algorithm_names, build_algorithm | |
| from agron_nets.training import ( | |
| RECIPES, | |
| AgronTrainer, | |
| TrainerConfig, | |
| get_recipe, | |
| recipe_names, | |
| ) | |
| from .conftest import make_batch, make_loader, make_model | |
| # Algorithms that need a bespoke batch layout or an external teacher/reference. | |
| SPECIAL = {"dpo", "distillation"} | |
| def _trainer(model, algorithms=None, steps=4, **cfg): | |
| config = TrainerConfig(max_steps=steps, epochs=1, log_every=0, warmup_steps=0, **cfg) | |
| return AgronTrainer( | |
| model, | |
| optimizer="adamw", | |
| scheduler="cosine", | |
| algorithms=algorithms, | |
| config=config, | |
| device="cpu", | |
| ) | |
| def test_algorithm_trains(name): | |
| model, _ = make_model("transformer") | |
| trainer = _trainer(model, [name], steps=6) | |
| report = trainer.fit(make_loader(6)) | |
| assert report["steps"] == 6 | |
| assert trainer.history == [] or any( | |
| torch.isfinite(torch.tensor(h["loss"])) for h in trainer.history | |
| ) | |
| def test_recipe_trains(name): | |
| model, _ = make_model("transformer") | |
| recipe = get_recipe(name) | |
| trainer = AgronTrainer( | |
| model, | |
| optimizer=recipe["optimizer"], | |
| scheduler=recipe["scheduler"], | |
| algorithms=recipe["algorithms"], | |
| config=TrainerConfig(max_steps=8, log_every=0, grad_clip=0.5), | |
| device="cpu", | |
| ) | |
| report = trainer.fit(make_loader(8)) | |
| assert report["steps"] >= 1 | |
| def test_dpo_algorithm_with_preference_batch(): | |
| model, _ = make_model("transformer") | |
| batch = { | |
| "chosen_input_ids": torch.randint(0, 64, (2, 8)), | |
| "rejected_input_ids": torch.randint(0, 64, (2, 8)), | |
| } | |
| trainer = _trainer(model, ["dpo"], steps=2) | |
| report = trainer.fit([batch, batch]) | |
| assert report["steps"] == 2 | |
| def test_distillation_with_teacher(): | |
| teacher, _ = make_model("transformer") | |
| student, _ = make_model("transformer") | |
| algo = build_algorithm("distillation", teacher=teacher) | |
| trainer = _trainer(student, [algo], steps=2) | |
| report = trainer.fit(make_loader(2)) | |
| assert report["steps"] == 2 | |
| def test_gradient_accumulation_only_steps_every_n(): | |
| model, _ = make_model("transformer") | |
| trainer = _trainer(model, [], steps=1) | |
| from agron_nets.algorithms import build_algorithm as build | |
| trainer.algorithms = [build("gradient_accumulation", steps=3)] | |
| trainer._ctx = None | |
| report = trainer.fit(make_loader(3)) | |
| # 3 micro-batches -> exactly one optimiser step. | |
| assert report["steps"] == 1 | |
| def test_schedule_free_optimizer(): | |
| model, _ = make_model("transformer") | |
| trainer = AgronTrainer( | |
| model, | |
| optimizer="schedulefree", | |
| scheduler="none", | |
| config=TrainerConfig(max_steps=3, log_every=0), | |
| device="cpu", | |
| ) | |
| report = trainer.fit(make_loader(3)) | |
| assert report["steps"] == 3 | |
| def test_checkpoint_roundtrip(tmp_path): | |
| model, _ = make_model("transformer") | |
| trainer = _trainer(model, ["ema"], steps=2, output_dir=str(tmp_path)) | |
| trainer.fit(make_loader(2)) | |
| trainer.save_checkpoint(str(tmp_path / "ckpt")) | |
| model2, _ = make_model("transformer") | |
| trainer2 = _trainer(model2, ["ema"], steps=2) | |
| trainer2.load_checkpoint(str(tmp_path / "ckpt")) | |
| assert trainer2.step == trainer.step == 2 | |
| def test_evaluate_returns_finite_loss(): | |
| model, _ = make_model("transformer") | |
| trainer = _trainer(model, [], steps=1) | |
| loss = trainer.evaluate(make_loader(2)) | |
| assert torch.isfinite(torch.tensor(loss)) | |
| def test_cross_architecture_engine(): | |
| for arch in ("llama", "mamba", "conv", "mixer", "gru", "hybrid", "moe", "hyper"): | |
| model, _ = make_model(arch) | |
| trainer = AgronTrainer( | |
| model, | |
| optimizer="lion", | |
| scheduler="linear", | |
| algorithms=["gradient_clip", "z_loss"], | |
| config=TrainerConfig(max_steps=2, log_every=0), | |
| device="cpu", | |
| ) | |
| report = trainer.fit(make_loader(2, vocab_size=model.config.vocab_size)) | |
| assert report["steps"] == 2, arch | |