68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
"""Alignment tests for lesson 14 training-loop mechanics.
|
|
|
|
Run after implementing:
|
|
|
|
python course/impl/tests/test_train_loop_mechanics.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[3]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from course.impl.core.train_loop import ( # noqa: E402
|
|
cosine_lr,
|
|
scale_loss_for_accumulation,
|
|
should_flush_tail,
|
|
should_step_optimizer,
|
|
)
|
|
from trainer.trainer_utils import get_lr # noqa: E402
|
|
|
|
|
|
def test_cosine_lr() -> None:
|
|
base_lr = 5e-4
|
|
total_steps = 20
|
|
for current_step in [0, 1, 5, 10, 19, 20]:
|
|
expected = get_lr(current_step, total_steps, base_lr)
|
|
actual = cosine_lr(current_step, total_steps, base_lr)
|
|
diff = abs(expected - actual)
|
|
print(f"lr_step_{current_step}_diff={diff:.12f}")
|
|
assert diff < 1e-12
|
|
|
|
|
|
def test_loss_scaling() -> None:
|
|
loss = torch.tensor(8.0)
|
|
scaled = scale_loss_for_accumulation(loss, accumulation_steps=4)
|
|
print(f"scaled_loss={scaled.item():.6f}")
|
|
assert torch.equal(scaled, torch.tensor(2.0))
|
|
|
|
|
|
def test_step_timing() -> None:
|
|
stepped = [step for step in range(1, 11) if should_step_optimizer(step, accumulation_steps=4)]
|
|
print(f"optimizer_step_micro_steps={stepped}")
|
|
assert stepped == [4, 8]
|
|
|
|
assert should_flush_tail(last_step=10, start_step=0, accumulation_steps=4)
|
|
assert not should_flush_tail(last_step=8, start_step=0, accumulation_steps=4)
|
|
assert not should_flush_tail(last_step=0, start_step=0, accumulation_steps=4)
|
|
|
|
|
|
def main() -> None:
|
|
try:
|
|
test_cosine_lr()
|
|
test_loss_scaling()
|
|
test_step_timing()
|
|
except NotImplementedError as exc:
|
|
raise SystemExit(f"TODO not implemented yet: {exc}") from exc
|
|
print("train_loop_mechanics=passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|