wehub-resource-sync
103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
"""
|
|
Regression tests for the legacy pretraining checkpoint/resume helpers.
|
|
|
|
Run from the repo root:
|
|
PYTHONPATH=. python tests/test_checkpoint_resume.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import os
|
|
from unittest import mock
|
|
|
|
import torch
|
|
|
|
from scripts.train_transformer import (
|
|
checkpoint_path,
|
|
list_checkpoints,
|
|
resolve_resume_path,
|
|
restore_training_checkpoint,
|
|
save_training_checkpoint,
|
|
prune_old_checkpoints,
|
|
)
|
|
from src.models.transformer import Transformer
|
|
|
|
|
|
def _tiny_model():
|
|
torch.manual_seed(0)
|
|
return Transformer(n_head=2, n_embed=8, context_length=8, vocab_size=32, N_BLOCKS=1)
|
|
|
|
|
|
def _tiny_config():
|
|
return {
|
|
"t_lr": 1e-3,
|
|
"t_lr_decayed": 1e-4,
|
|
"t_lr_decay_step": 10,
|
|
"device": "cpu",
|
|
}
|
|
|
|
|
|
def test_checkpoint_round_trip_and_latest_resume():
|
|
cfg = _tiny_config()
|
|
model = _tiny_model()
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["t_lr"])
|
|
|
|
# Populate optimizer state so the round trip proves optimizer resume works too.
|
|
idx = torch.randint(0, 32, (2, 4))
|
|
_, loss = model(idx, idx)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
losses = [3.0, 2.0, 1.0]
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
first = checkpoint_path(tmp, 2)
|
|
second = checkpoint_path(tmp, 3)
|
|
save_training_checkpoint(first, model, optimizer, cfg, losses, step=2)
|
|
save_training_checkpoint(second, model, optimizer, cfg, losses + [0.5], step=3)
|
|
prune_old_checkpoints(tmp, keep_last=1)
|
|
|
|
remaining = list_checkpoints(tmp)
|
|
assert remaining == [second]
|
|
assert resolve_resume_path("latest", tmp) == second
|
|
|
|
restored = _tiny_model()
|
|
restored_optim = torch.optim.AdamW(restored.parameters(), lr=cfg["t_lr"])
|
|
next_step, restored_losses = restore_training_checkpoint(
|
|
second,
|
|
restored,
|
|
restored_optim,
|
|
cfg,
|
|
"cpu",
|
|
)
|
|
|
|
assert next_step == 4
|
|
assert restored_losses == losses + [0.5]
|
|
assert restored_optim.state_dict()["state"]
|
|
for original, loaded in zip(model.parameters(), restored.parameters()):
|
|
assert torch.allclose(original, loaded)
|
|
|
|
|
|
def test_checkpoint_save_failure_does_not_leave_partial_file():
|
|
cfg = _tiny_config()
|
|
model = _tiny_model()
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["t_lr"])
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
target = checkpoint_path(tmp, 7)
|
|
with mock.patch("scripts.train_transformer.torch.save", side_effect=RuntimeError("boom")):
|
|
try:
|
|
save_training_checkpoint(target, model, optimizer, cfg, [1.0], step=7)
|
|
assert False, "save_training_checkpoint should re-raise save errors"
|
|
except RuntimeError:
|
|
pass
|
|
|
|
assert not os.path.exists(target)
|
|
assert not [name for name in os.listdir(tmp) if name.endswith(".tmp")]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_checkpoint_round_trip_and_latest_resume()
|
|
test_checkpoint_save_failure_does_not_leave_partial_file()
|
|
print("ALL CHECKPOINT RESUME TESTS PASSED")
|