ViperEkura
9d3ae76683
- extract shared helpers for dataset writers, scheduler construction, thread interleaving, hf roundtrips, and moe configs - remove about 20 cases whose only assertions were format checks, restated declarations, fake-taxonomy duplicates, or test-local scaffolding - strengthen weak cases into exact reference comparisons, positional mask checks, and deterministic outcomes - replace two schedule factory smoke tests with cosine/sgdr formula assertions - delete root-level CLI tests whose merge-priority facts are covered by tests/config/test_cli.py - suite shrinks from 857 to 826 items; ruff format, import order, and pytest all green
118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from astrai.trainer.schedule import CosineScheduler, SchedulerFactory, SGDRScheduler
|
|
|
|
|
|
def _stepped_lrs(scheduler, optimizer, n_steps):
|
|
"""Return the lr after construction plus each of *n_steps* steps."""
|
|
lrs = list(scheduler.get_last_lr())
|
|
for _ in range(n_steps):
|
|
optimizer.step()
|
|
scheduler.step()
|
|
lrs.append(scheduler.get_last_lr()[0])
|
|
return lrs
|
|
|
|
|
|
def test_cosine_scheduler_warms_up_then_decays_to_floor():
|
|
"""lr ramps linearly to base_lr during warmup, cosine-decays after it,
|
|
and never drops below min_rate * base_lr."""
|
|
base_lr = 0.001
|
|
model = torch.nn.Linear(10, 2)
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr)
|
|
scheduler = SchedulerFactory.create(
|
|
"cosine", optimizer, warmup_steps=2, lr_decay_steps=4, min_rate=0.1
|
|
)
|
|
|
|
assert isinstance(scheduler, CosineScheduler)
|
|
lrs = _stepped_lrs(scheduler, optimizer, n_steps=7)
|
|
|
|
assert lrs[0] == pytest.approx(0.1 * base_lr) # warmup starts at the floor
|
|
assert lrs[1] == pytest.approx(0.5 * base_lr) # halfway through warmup
|
|
assert lrs[2] == pytest.approx(base_lr) # warmup complete
|
|
expected_mid = base_lr * 0.5 * (1.0 + math.cos(math.pi * 0.25))
|
|
assert lrs[3] == pytest.approx(expected_mid) # quarter into decay
|
|
assert lrs[5] > 0.1 * base_lr # 3/4 into decay: not clamped yet
|
|
assert lrs[6] == pytest.approx(0.1 * base_lr) # clamped at min_rate floor
|
|
assert lrs[7] == pytest.approx(0.1 * base_lr) # stays at the floor
|
|
assert all(lr >= 0.1 * base_lr - 1e-12 for lr in lrs)
|
|
|
|
|
|
def test_cosine_scheduler_decays_to_zero_with_min_rate_zero():
|
|
"""min_rate=0 must reach exactly 0.0 at the end of decay, not NaN."""
|
|
base_lr = 0.001
|
|
model = torch.nn.Linear(10, 2)
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr)
|
|
scheduler = SchedulerFactory.create(
|
|
"cosine", optimizer, warmup_steps=1, lr_decay_steps=9, min_rate=0.0
|
|
)
|
|
|
|
lrs = _stepped_lrs(scheduler, optimizer, n_steps=11)
|
|
|
|
assert lrs[10] == 0.0
|
|
assert lrs[11] == 0.0
|
|
assert all(math.isfinite(lr) for lr in lrs)
|
|
|
|
|
|
def test_sgdr_scheduler_restarts_each_cycle():
|
|
"""lr anneals within a cycle, then jumps back to base_lr on restart."""
|
|
base_lr = 0.001
|
|
model = torch.nn.Linear(10, 2)
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr)
|
|
scheduler = SchedulerFactory.create(
|
|
"sgdr", optimizer, warmup_steps=2, cycle_length=4, t_mult=1, min_rate=0.1
|
|
)
|
|
|
|
assert isinstance(scheduler, SGDRScheduler)
|
|
lrs = _stepped_lrs(scheduler, optimizer, n_steps=7)
|
|
|
|
assert lrs[2] == pytest.approx(base_lr) # cycle start
|
|
expected_mid = base_lr * (0.1 + 0.9 * 0.5) # halfway through the cycle
|
|
assert lrs[4] == pytest.approx(expected_mid)
|
|
assert lrs[5] < lrs[4] # still annealing at the cycle end
|
|
assert lrs[6] == pytest.approx(base_lr) # restart: back to full lr
|
|
|
|
|
|
def test_schedule_factory_state_persistence():
|
|
"""Test scheduler state persistence (save/load)"""
|
|
|
|
model = torch.nn.Linear(10, 2)
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
|
|
|
|
# Create scheduler directly with parameters
|
|
warmup_steps = 100
|
|
total_steps = 1000
|
|
min_rate = 0.1
|
|
lr_decay_steps = total_steps - warmup_steps
|
|
scheduler = SchedulerFactory.create(
|
|
"cosine",
|
|
optimizer,
|
|
warmup_steps=warmup_steps,
|
|
lr_decay_steps=lr_decay_steps,
|
|
min_rate=min_rate,
|
|
)
|
|
|
|
# Take a few steps
|
|
for _ in range(5):
|
|
optimizer.step()
|
|
scheduler.step()
|
|
|
|
# Save state
|
|
state_dict = scheduler.state_dict()
|
|
|
|
# Create new scheduler with same parameters
|
|
new_scheduler = SchedulerFactory.create(
|
|
"cosine",
|
|
optimizer,
|
|
warmup_steps=warmup_steps,
|
|
lr_decay_steps=lr_decay_steps,
|
|
min_rate=min_rate,
|
|
)
|
|
new_scheduler.load_state_dict(state_dict)
|
|
|
|
# Verify states match
|
|
assert scheduler.last_epoch == new_scheduler.last_epoch
|
|
assert scheduler.get_last_lr() == new_scheduler.get_last_lr()
|