File size: 4,345 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 | """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",
)
@pytest.mark.parametrize("name", [n for n in algorithm_names() if n not in SPECIAL | {"gradient_accumulation"}])
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
)
@pytest.mark.parametrize("name", recipe_names())
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
|