| 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]) |
| |
| 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: |
| |
| 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() |
|
|
|
|
| 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) |
|
|
