File size: 1,479 Bytes
d7228c8 | 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 | """
Unit tests for CheckpointManager saving and loading functionality.
"""
import os
import torch
import pytest
from slm.config.model_config import ModelConfig
from slm.config.train_config import TrainConfig
from slm.model.transformer_lm import SLMForCausalLM
from slm.checkpoint.manager import CheckpointManager
def test_checkpoint_manager_save_and_load(tmp_path):
ckpt_dir = str(tmp_path / "checkpoints")
manager = CheckpointManager(output_dir=ckpt_dir)
model_config = ModelConfig(vocab_size=100, d_model=32, n_heads=2, n_layers=1)
train_config = TrainConfig(output_dir=ckpt_dir)
model = SLMForCausalLM(model_config)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
ckpt_path = manager.save_checkpoint(
model=model,
optimizer=optimizer,
scheduler=None,
step=10,
epoch=1,
loss=2.5,
model_config=model_config,
train_config=train_config,
is_best=True
)
assert os.path.exists(ckpt_path)
# Test loading into a new model instance
new_model = SLMForCausalLM(model_config)
loaded_meta = manager.load_checkpoint(ckpt_path, new_model)
assert loaded_meta["step"] == 10
assert loaded_meta["epoch"] == 1
assert loaded_meta["loss"] == 2.5
assert loaded_meta["model_config"]["vocab_size"] == 100
# Verify weight restoration match
for p1, p2 in zip(model.parameters(), new_model.parameters()):
assert torch.equal(p1, p2)
|