File size: 4,177 Bytes
6e802a4 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | # test_release_weights.py: Released safetensors load exactly like the training checkpoint.
import json
import subprocess
import sys
from pathlib import Path
import pytest
import torch
import torch.nn as nn
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
safetensors = pytest.importorskip("safetensors")
from safetensors.torch import load_file, save_file
from train import (checkpoint_speaker_rows, checkpoint_tensor_shapes, load_checkpoint,
read_checkpoint)
def _sd():
g = torch.Generator().manual_seed(0)
return {"speaker_encoder.embedding.weight": torch.randn(7, 3, generator=g),
"plan_injector.proj.weight": torch.randn(4, 2, generator=g),
"lin.weight": torch.randn(3, 3, generator=g), "lin.bias": torch.randn(3, generator=g)}
def test_read_checkpoint_safetensors_matches_pt(tmp_path):
sd = _sd()
pt = tmp_path / "c.pt"
torch.save({"model_state_dict": sd, "step": 40000, "optimizer_state_dict": {}}, pt)
st = tmp_path / "model.safetensors"
save_file(sd, str(st), metadata={"step": "40000"})
a, b = read_checkpoint(pt), read_checkpoint(st)
assert b["step"] == 40000 and set(b) == {"model_state_dict", "step"}
assert set(a["model_state_dict"]) == set(b["model_state_dict"])
for k in sd:
assert torch.equal(a["model_state_dict"][k], b["model_state_dict"][k])
assert checkpoint_speaker_rows(pt) == checkpoint_speaker_rows(st) == 7
assert checkpoint_tensor_shapes(pt) == checkpoint_tensor_shapes(st)
assert any(k.startswith("plan_injector.") for k in checkpoint_tensor_shapes(st))
def test_load_checkpoint_accepts_safetensors(tmp_path):
class Tiny(nn.Module):
def __init__(self):
super().__init__()
self.lin = nn.Linear(3, 3)
src = Tiny()
st = tmp_path / "model.safetensors"
save_file({k: v.detach().clone() for k, v in src.state_dict().items()}, str(st), metadata={"step": "7"})
dst = Tiny()
assert load_checkpoint(str(st), dst, device="cpu") == 7
for k, v in src.state_dict().items():
assert torch.equal(dst.state_dict()[k], v)
def _fake_checkpoint(path, **args):
t = torch.randn(5, 4)
sd = _sd()
sd["head.weight"] = t
sd["tied.weight"] = t
targs = {"hybrid_attention_top_k": 0, "backbone": "mamba", "depth_arch": "mamba2", "depth_layers": 4,
"depth_dim": 1024, "depth_cond": "prefix", "depth_feedback": "all"}
targs.update(args)
torch.save({"model_state_dict": sd, "optimizer_state_dict": {"x": torch.zeros(3)}, "step": 40000, "stage": 2,
"args": targs}, path)
return sd
def _export(ck, out, arch):
return subprocess.run([sys.executable, str(ROOT / "scripts" / "export_release_weights.py"), "--checkpoint", str(ck),
"--arch", arch, "--name", "t", "--out", str(out)], capture_output=True, text=True)
def test_export_is_bit_exact_and_keeps_aliases(tmp_path):
ck = tmp_path / "checkpoint_step_40000.pt"
sd = _fake_checkpoint(ck)
r = _export(ck, tmp_path / "out", "pure")
assert r.returncode == 0, r.stdout + r.stderr
back = load_file(str(tmp_path / "out" / "model.safetensors"))
assert set(back) == set(sd)
for k in sd:
assert torch.equal(back[k], sd[k])
cfg = json.loads((tmp_path / "out" / "config.json").read_text())
assert cfg["bit_exact_vs_checkpoint"] and cfg["step"] == 40000 and cfg["speaker_table_rows"] == 7
assert cfg["tied_aliases"] == {"tied.weight": "head.weight"}
assert cfg["generator_flags"][-2:] == ["--num_speakers", "7"] and cfg["has_plan_injector"]
targs = json.loads((tmp_path / "out" / "training_args.json").read_text())
assert targs["depth_cond"] == "prefix"
@pytest.mark.parametrize("arch,args", [("hyb", {}), ("pure", {"hybrid_attention_top_k": 4}),
("tfm", {}), ("pure", {"depth_layers": 2})])
def test_export_refuses_a_mislabelled_checkpoint(tmp_path, arch, args):
ck = tmp_path / "c.pt"
_fake_checkpoint(ck, **args)
r = _export(ck, tmp_path / "out", arch)
assert r.returncode != 0 and "FATAL" in (r.stdout + r.stderr)
|