File size: 7,308 Bytes
0693474 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | from __future__ import annotations
from pathlib import Path
import math
import random
import signal
import threading
import torch
import pytest
from offload.dspark.d2.train.data import Block, SyntheticBlocks
from offload.dspark.d2.train.data import sha256
from offload.dspark.d2.train.loop import Run, train
from offload.dspark.d2.train import loop
from offload.dspark.d2.train.model import Drafter
from offload.dspark.d2.train.synthetic import SyntheticSource, tiny_config
def model() -> Drafter:
torch.manual_seed(31)
c = tiny_config()
return Drafter(SyntheticSource(c), c)
@pytest.mark.parametrize("batch_size", (1, 2))
def test_exact_resume_optimizer_rng_schedule_data_order(tmp_path: Path, batch_size: int) -> None:
torch.set_num_threads(2)
data = SyntheticBlocks(tiny_config(), count=3)
run = Run(epochs=2, chunk=32, batch_size=batch_size)
a, b = model(), model()
pa, pb = tmp_path/"full.pt", tmp_path/"resume.pt"
train(a, data, run, pa)
train(b, data, run, pb, max_steps=2)
archive = tmp_path/"resume.checkpoints/step-000000002.pt"
preserved = sha256(archive)
resumed = model()
train(resumed, data, run, pb, resume=True)
assert sha256(archive) == preserved
assert torch.load(archive, weights_only=False)["step"] == 2
sa = torch.load(pa, weights_only=False)
sb = torch.load(pb, weights_only=False)
assert sa["step"] == sb["step"] == 2*math.ceil(3/batch_size)
for n in sa["trainable"]:
assert torch.equal(sa["trainable"][n], sb["trainable"][n]), n
for i, oa in sa["optimizer"]["state"].items():
for n, t in oa.items():
assert torch.equal(t, sb["optimizer"]["state"][i][n])
assert sa["scheduler"] == sb["scheduler"]
assert [r["loss"] for r in sa["curve"]] == [r["loss"] for r in sb["curve"]]
assert torch.equal(sa["torch_rng"], sb["torch_rng"])
assert pa.stat().st_mode & 0o777 == 0o600
def test_stop_writes_boundary_and_changed_contract_refused(tmp_path: Path) -> None:
stop = threading.Event()
stop.set()
p = tmp_path/"term.pt"
data = SyntheticBlocks(tiny_config(), count=2)
out = train(model(), data, Run(epochs=1), p, stop)
assert out["step"] == 0 and p.is_file()
with pytest.raises(ValueError, match="resume contract changed"):
train(model(), data, Run(epochs=2), p, resume=True)
def test_mid_forward_cancellation_discards_partial_step(tmp_path: Path) -> None:
stop = threading.Event()
m = model()
original = {n: p.detach().clone() for n, p in m.named_parameters()}
handle = m.main_proj.register_forward_hook(lambda *_: stop.set())
path = tmp_path/"cancelled.pt"
result = train(m, SyntheticBlocks(m.c, count=2), Run(epochs=1), path, stop)
handle.remove()
assert result["step"] == 0
state = torch.load(path, weights_only=False)
for n, p in state["trainable"].items():
assert torch.equal(p, original[n])
# The cancelled scope must not poison subsequent independent inference.
stop.clear()
m.hidden_from(SyntheticBlocks(m.c)[0]["inputs"])
@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGINT])
@pytest.mark.parametrize("cancel_inside_greedy", [True, False])
def test_reference_observer_signal_saves_and_resumes_boundary(
tmp_path: Path, signum: int, cancel_inside_greedy: bool,
) -> None:
def reference() -> Drafter:
torch.manual_seed(31)
c = tiny_config()
return Drafter(SyntheticSource(c), c, acc=torch.float64, reference_math=True)
stop = threading.Event()
seen: list[int] = []
def observe(m: Drafter, block: Block, h: torch.Tensor, hx: torch.Tensor, metrics: dict[str, float]) -> None:
seen.append(int(metrics["step"]))
if metrics["step"] == 2:
# Both RNG streams must roll back with the uncommitted step.
torch.rand(3)
random.random()
signal.raise_signal(signum)
if cancel_inside_greedy:
m.greedy(h, hx, block["inputs"]["anchor"], chunk=32)
data = SyntheticBlocks(tiny_config(), count=3)
run = Run(epochs=1, chunk=32)
full, boundary, interrupted = [tmp_path/n for n in ("full.pt", "boundary.pt", "interrupted.pt")]
train(reference(), data, run, full)
train(reference(), data, run, boundary, max_steps=1)
m = reference()
old = signal.signal(signum, lambda *_: stop.set())
try:
result = train(m, data, run, interrupted, stop, observer=observe)
finally:
signal.signal(signum, old)
assert seen == [1, 2]
assert result["step"] == 1 and result["complete"] is False
assert all(p.grad is None for p in m.parameters())
saved = torch.load(interrupted, weights_only=False)
expected = torch.load(boundary, weights_only=False)
assert saved["step"] == expected["step"] == 1
assert torch.equal(saved["torch_rng"], expected["torch_rng"])
assert saved["random_rng"] == expected["random_rng"]
for n, p in saved["trainable"].items():
assert torch.equal(p, expected["trainable"][n]), n
train(reference(), data, run, interrupted, resume=True)
saved, expected = [torch.load(p, weights_only=False) for p in (interrupted, full)]
for n, p in saved["trainable"].items():
assert torch.equal(p, expected["trainable"][n]), n
for i, state in saved["optimizer"]["state"].items():
for n, t in state.items():
assert torch.equal(t, expected["optimizer"]["state"][i][n])
assert saved["scheduler"] == expected["scheduler"]
assert [r["loss"] for r in saved["curve"]] == [r["loss"] for r in expected["curve"]]
def test_confidence_stage_refuses_zero_labelled_updates(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(loop, "confidence_loss", lambda *a, **k: (None, 0))
output = tmp_path/"confidence.pt"
with pytest.raises(ValueError, match="no labelled updates"):
train(model(), SyntheticBlocks(tiny_config(), count=1), Run(epochs=1, stage="confidence"), output)
assert output.is_file() # retain the failed stage's inspection/checkpoint evidence
def test_initial_adapter_requires_matching_frozen_source(tmp_path: Path) -> None:
from offload.dspark.d2.train.replay import apply_adapter
m = model()
p = tmp_path/"adapter.pt"
state = {"schema": "dsv41.d2.train-checkpoint.v1", "frozen_identity": m.frozen_identity,
"trainable": {n: v.detach().clone() for n, v in m.named_parameters()}}
torch.save(state, p)
assert apply_adapter(model(), str(p)) == sha256(p)
state["frozen_identity"] = "different-source"
torch.save(state, p)
with pytest.raises(ValueError, match="frozen source"):
apply_adapter(model(), str(p))
def test_reference_resume_refuses_a_fast_backend_checkpoint(tmp_path: Path) -> None:
stop = threading.Event()
stop.set()
p = tmp_path / "reference.pt"
data = SyntheticBlocks(tiny_config(), count=1)
run = Run(epochs=1)
train(model(), data, run, p, stop)
state = torch.load(p, weights_only=False)
assert state["fast_training"] is False and state["fast_backend"] is None
state["fast_training"] = True
loop.atomic_save(p, state)
with pytest.raises(ValueError, match="backend changed"):
train(model(), data, run, p, resume=True)
|