File size: 888 Bytes
f0b317c | 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 | from safetensors.torch import load_file
from model import GPT, GPTConfig
EXPECTED_PARAMETER_COUNT = 4_782_336
def test_state_dict_keys_match(config_data, model):
state_names = set(model.state_dict().keys())
loaded = load_file(
"model.safetensors",
device="cpu",
)
assert set(loaded.keys()) == state_names
def test_strict_load_succeeds(config_data, model):
loaded = load_file(
"model.safetensors",
device="cpu",
)
result = GPT(config=GPTConfig.from_json(config_data)).load_state_dict(
loaded,
strict=True,
)
assert not result.missing_keys
assert not result.unexpected_keys
def test_parameter_count(model):
parameters = sum(
parameter.numel()
for parameter in model.parameters()
if parameter.requires_grad
)
assert parameters == EXPECTED_PARAMETER_COUNT |