| """Codex audit U01 (distributed training control plane) -- confirmed-open findings. |
| |
| Each test is named after the incident it prevents and FAILED on the pre-fix code. |
| |
| Run: CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=2 PYTHONPATH=$PWD .venv/bin/python -m pytest -q tests/test_audit_u01_ctrlplane.py |
| """ |
|
|
| from __future__ import annotations |
|
|
| import gc |
| import glob |
| import json |
| import os |
| import struct |
| import weakref |
| import zipfile |
|
|
| import pytest |
| import torch |
| from test_model_parity import tiny_config |
| from test_trainer import _tiny_trainer |
|
|
| import align.train.trainer as trainer_mod |
| from align.exit_codes import read_halt_marker |
| from align.train.probes import GradNormTripwire, LengthDriftTripwire, LogpsChosenTripwire |
| from align.train.trainer import HaltError |
|
|
|
|
| @pytest.fixture |
| def tiny(tiny_model_constants): |
| return tiny_config(tiny_model_constants) |
|
|
|
|
| def _ckpts(run_dir): |
| return sorted(glob.glob(os.path.join(run_dir, "checkpoints", "micro_*"))) |
|
|
|
|
| def _flip_payload_byte(path: str) -> None: |
| """Same-size bit flip inside a stored tensor of a torch.save zip (the audit's r8).""" |
| b = bytearray(open(path, "rb").read()) |
| with zipfile.ZipFile(path) as z: |
| info = [i for i in z.infolist() if "/data/" in i.filename and i.file_size > 8][0] |
| off = info.header_offset |
| n, m = struct.unpack("<HH", bytes(b[off + 26 : off + 30])) |
| b[off + 30 + n + m + 3] ^= 0x40 |
| open(path, "wb").write(bytes(b)) |
|
|
|
|
| |
| |
| |
|
|
|
|
| class _SpyGradNorm(GradNormTripwire): |
| def __init__(self, **kw): |
| super().__init__(**kw) |
| self.seen: list[float] = [] |
|
|
| def update(self, grad_norm): |
| self.seen.append(float(grad_norm)) |
| return super().update(grad_norm) |
|
|
|
|
| def test_grad_norm_halt_wire_sees_every_optimizer_step_not_every_logged_one(tiny, tmp_workdir): |
| """THE 80%-BLIND HALT WIRE (U01-8). The grad_norm wire was fed only from `_log`; with |
| log_every=10 and grad_accum=2 one optimizer step in five was inspected.""" |
| t, _ = _tiny_trainer( |
| tiny, os.path.join(tmp_workdir, "a"), n_batches=10, grad_accum=2, armed=("grad_norm",), |
| log_every_micro_steps=10, save_every_micro_steps=10, |
| ) |
| spy = _SpyGradNorm() |
| t.tripwires["grad_norm"] = spy |
| t.train() |
| assert t.optimizer_step == 5 |
| finite = [g for g in spy.seen if g == g] |
| assert len(finite) == 5, f"wire inspected {len(finite)} of 5 optimizer steps: {spy.seen}" |
|
|
|
|
| def test_grad_norm_explosion_halts_at_the_boundary_even_between_log_steps(tiny, tmp_workdir): |
| t, _ = _tiny_trainer( |
| tiny, os.path.join(tmp_workdir, "a"), n_batches=10, grad_accum=2, armed=("grad_norm",), |
| log_every_micro_steps=10, save_every_micro_steps=10, |
| ) |
| t.tripwires["grad_norm"] = GradNormTripwire(absolute_ceiling=1e-12) |
| with pytest.raises(HaltError, match="grad_norm") as ei: |
| t.train() |
| assert t.micro_step == 2, "the first optimizer boundary must halt, not the first log step" |
| assert ei.value.exit_code == HaltError.exit_code |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _capture_step_grad_norms(t): |
| norms = [] |
| orig = t.train_module.optim_step |
|
|
| def spy(): |
| norms.append( |
| sum(float(p.grad.float().norm()) ** 2 for p in t.model.parameters() if p.grad is not None) ** 0.5 |
| ) |
| return orig() |
|
|
| t.train_module.optim_step = spy |
| return norms |
|
|
|
|
| def test_tail_partial_accumulation_window_is_stepped_and_leaves_no_stale_grads(tiny, tmp_workdir): |
| """THE DROPPED-TAIL DEFECT (U01-1): 3 batches at accum 2 planned ceil(3/2)=2 optimizer steps, |
| took 1, and left 21 params holding the pending .grad of the unstepped micro-batch.""" |
| t, _ = _tiny_trainer( |
| tiny, os.path.join(tmp_workdir, "a"), n_batches=3, grad_accum=2, save_every_micro_steps=2 |
| ) |
| s = t.train() |
| assert t.optimizer_step == t.scheduler.total_steps == 2 |
| assert all(p.grad is None or float(p.grad.abs().sum()) == 0.0 for p in t.model.parameters()) |
| assert s["tail_window"]["micros"] == 1 and s["tail_window"]["action"] == "stepped" |
|
|
|
|
| def test_tail_window_gradient_is_rescaled_to_a_mean_over_the_micros_it_holds(tiny, tmp_workdir): |
| """A 1-micro tail at accum 2 carries grad/2; stepping it as-is under-weights it. It must |
| equal the gradient a grad_accum=1 trainer computes on the same batch.""" |
| a, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=1, grad_accum=2, save_every_micro_steps=2) |
| b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=1, grad_accum=1, save_every_micro_steps=1) |
| na, nb = _capture_step_grad_norms(a), _capture_step_grad_norms(b) |
| a.train() |
| b.train() |
| assert len(na) == len(nb) == 1 |
| assert na[0] == pytest.approx(nb[0], rel=1e-4) |
|
|
|
|
| def test_a_capped_run_stopping_mid_window_zeroes_the_pending_grads(tiny, tmp_workdir): |
| t, _ = _tiny_trainer( |
| tiny, os.path.join(tmp_workdir, "a"), n_batches=6, grad_accum=2, max_micro_steps=3, |
| save_every_micro_steps=2, |
| ) |
| s = t.train() |
| assert all(p.grad is None or float(p.grad.abs().sum()) == 0.0 for p in t.model.parameters()) |
| assert s["tail_window"]["action"] == "dropped" |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_resume_without_optimizer_state_reseeds_the_fp32_masters_from_the_loaded_weights(tiny, tmp_workdir): |
| """THE REVERTING-FIRST-STEP DEFECT (U01-2): OffloadAdamW masters stay at launch weights when |
| the checkpoint carries no optimizer state; the first step then overwrites params from them.""" |
| a, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| a.train() |
| ck = _ckpts(os.path.join(tmp_workdir, "a"))[-1] |
| st = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu") |
| for k in ("optimizer_state", "optimizer_per_rank"): |
| st.pop(k, None) |
| torch.save(st, os.path.join(ck, "resume.pt")) |
| os.remove(os.path.join(ck, "integrity.json")) if os.path.exists(os.path.join(ck, "integrity.json")) else None |
| b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| b.resume_from(ck) |
| o = b.train_module.optimizer |
| d = max( |
| float((p.detach().float() - m.detach().float()).abs().max()) for p, m in zip(o.params, o._master) |
| ) |
| assert d == 0.0, f"masters diverge from the restored params by {d}" |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _saved(tiny, tmp_workdir, name="a", **kw): |
| run = os.path.join(tmp_workdir, name) |
| t, _ = _tiny_trainer(tiny, run, n_batches=4, grad_accum=2, save_every_micro_steps=4, **kw) |
| t.train() |
| return run, _ckpts(run)[-1] |
|
|
|
|
| def test_a_same_size_bit_flip_in_resume_pt_is_refused_on_load(tiny, tmp_workdir): |
| """THE UNHASHED-PAYLOAD DEFECT (U01-3): resume.pt / optimizer files had no sha256; a flipped |
| byte in a stored tensor loaded as garbage.""" |
| _, ck = _saved(tiny, tmp_workdir) |
| _flip_payload_byte(os.path.join(ck, "resume.pt")) |
| b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| with pytest.raises(Exception, match="sha256"): |
| b.resume_from(ck) |
|
|
|
|
| def test_a_bit_flip_in_a_per_rank_optimizer_file_is_refused_on_load(tiny, tmp_workdir): |
| def hook(cfg): |
| cfg.train.per_rank_optimizer_state = True |
|
|
| _, ck = _saved(tiny, tmp_workdir, cfg_hook=hook) |
| _flip_payload_byte(os.path.join(ck, "optimizer_rank0.pt")) |
| b, _ = _tiny_trainer( |
| tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4, cfg_hook=hook |
| ) |
| with pytest.raises(Exception, match="sha256"): |
| b.resume_from(ck) |
|
|
|
|
| def test_a_checkpoint_without_the_hash_manifest_still_loads_with_a_warning(tiny, tmp_workdir, caplog): |
| _, ck = _saved(tiny, tmp_workdir) |
| assert os.path.isfile(os.path.join(ck, "integrity.json")), "new checkpoints must record hashes" |
| os.remove(os.path.join(ck, "integrity.json")) |
| b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| with caplog.at_level("WARNING"): |
| b.resume_from(ck) |
| assert any("integrity" in r.message for r in caplog.records) |
|
|
|
|
| |
| |
| |
|
|
|
|
| @pytest.mark.parametrize( |
| "kw,match", |
| [ |
| (dict(seed=7), "seed"), |
| (dict(grad_accum=4), "grad_accum"), |
| (dict(n_batches=6), "data"), |
| ], |
| ) |
| def test_resume_refuses_a_changed_seed_accum_or_data_unless_overridden(tiny, tmp_workdir, kw, match): |
| """THE SILENT-RECONFIGURATION RESUME (U01-5): seed / grad_accum / data length changed between |
| save and resume, nothing noticed.""" |
| _, ck = _saved(tiny, tmp_workdir) |
| args = dict(n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| seed = kw.pop("seed", None) |
| args.update(kw) |
| if seed is not None: |
| args["seed"] = seed |
| b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), **args) |
| with pytest.raises(RuntimeError, match=match): |
| b.resume_from(ck) |
| b.allow_resume_mismatch = True |
| b.resume_from(ck) |
|
|
|
|
| def test_resume_of_a_checkpoint_that_predates_the_fingerprint_warns_and_loads(tiny, tmp_workdir, caplog): |
| _, ck = _saved(tiny, tmp_workdir) |
| p = os.path.join(ck, "resume.pt") |
| st = torch.load(p, weights_only=True, map_location="cpu") |
| for k in ("seed", "grad_accum", "data_fingerprint"): |
| st.pop(k, None) |
| torch.save(st, p) |
| os.remove(os.path.join(ck, "integrity.json")) |
| b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| with caplog.at_level("WARNING"): |
| b.resume_from(ck) |
| assert any("fingerprint" in r.message or "predates" in r.message for r in caplog.records) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_tripwire_state_roundtrips_through_resume_pt(tiny, tmp_workdir): |
| """THE RESET-ON-RELAUNCH DEFECT (U01-6): baselines and latches died with the process.""" |
| _, ck = _saved(tiny, tmp_workdir, armed=("grad_norm",)) |
| st = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu") |
| assert "tripwire_per_rank" in st |
| b, _ = _tiny_trainer( |
| tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4, |
| armed=("grad_norm",), |
| ) |
| assert b.tripwires["grad_norm"]._history == [] |
| b.resume_from(ck) |
| assert len(b.tripwires["grad_norm"]._history) == 2 |
|
|
|
|
| def test_wire_state_dicts_roundtrip_baselines_and_latches(): |
| w = LogpsChosenTripwire(window=3) |
| for i in range(5): |
| w.update(-10.0 - i, 10.0, 0.1 * i) |
| w.fired, w._fired_at_row = True, 4 |
| w2 = LogpsChosenTripwire(window=3) |
| w2.load_state_dict(json.loads(json.dumps(w.state_dict()))) |
| assert w2.fired and w2._baseline_nll == w._baseline_nll and w2._nll == w._nll and w2._rows == w._rows |
|
|
| g = GradNormTripwire() |
| g.update(1.0) |
| g.fired = True |
| g2 = GradNormTripwire() |
| g2.load_state_dict(g.state_dict()) |
| assert g2.fired and g2._history == [1.0] |
|
|
| d = LengthDriftTripwire(baseline=0.5, min_pairs=10) |
| d.update(0.5, n_pairs=7) |
| d2 = LengthDriftTripwire(baseline=0.5, min_pairs=10) |
| d2.load_state_dict(d.state_dict()) |
| assert d2._pairs_seen == 7 |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _golden_report(**cfg_extra): |
| return { |
| "per_set": {"math500": {"n": 500, "pass@1": 0.3}, "gsm8k": {"n": 500, "pass@1": 0.55}}, |
| "config": {"mde": 0.05, "n_problems": 1000, "eval_set_hash": "g", **cfg_extra}, |
| } |
|
|
|
|
| def test_golden_report_produced_by_a_different_eval_source_is_not_recertified(tmp_workdir): |
| """THE STALE-REPORT RE-CERTIFICATION (U01-9): the ledger drops the golden pass when |
| align/eval changes, and `record_golden_from_report` then re-recorded it from the OLD report.""" |
| from align.train.stage_gates import GateLedger, eval_source_hash, record_golden_from_report |
|
|
| cur = eval_source_hash() |
| ledger = GateLedger(os.path.join(tmp_workdir, "a"), eval_source_hash=cur) |
| stale = os.path.join(tmp_workdir, "stale.json") |
| rep = _golden_report() |
| rep["eval_source_hash"] = "0" * 16 |
| json.dump(rep, open(stale, "w")) |
| out = record_golden_from_report(ledger, report_path=stale) |
| assert out["recorded"] is False and out["passed"] is None and "eval" in out["reason"] |
| assert not ledger.is_passed("golden_regression") |
|
|
| ok = os.path.join(tmp_workdir, "ok.json") |
| rep["eval_source_hash"] = cur |
| json.dump(rep, open(ok, "w")) |
| assert record_golden_from_report(ledger, report_path=ok)["passed"] is True |
|
|
|
|
| def test_an_unstamped_legacy_golden_report_still_records_with_a_warning(tmp_workdir, caplog): |
| from align.train.stage_gates import GateLedger, record_golden_from_report |
|
|
| ledger = GateLedger(os.path.join(tmp_workdir, "a")) |
| p = os.path.join(tmp_workdir, "legacy.json") |
| json.dump(_golden_report(), open(p, "w")) |
| with caplog.at_level("WARNING"): |
| assert record_golden_from_report(ledger, report_path=p)["passed"] is True |
| assert any("eval_source_hash" in r.message for r in caplog.records) |
|
|
|
|
| def test_stamp_golden_report_writes_the_producing_hash(tmp_workdir): |
| from align.train.stage_gates import eval_source_hash, stamp_golden_report |
|
|
| p = os.path.join(tmp_workdir, "r.json") |
| json.dump(_golden_report(), open(p, "w")) |
| stamp_golden_report(p) |
| assert json.load(open(p))["eval_source_hash"] == eval_source_hash() |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_model_is_collectable_after_the_post_train_handoff(tiny, tmp_workdir): |
| """THE HANDOFF LEAK (U01-11): `_probe_target` / `trunk_probe._model` kept the training replica |
| alive after scripts/sft.py set `trainer.model = None`, starving the post-train eval build.""" |
| t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| t.train() |
| w = weakref.ref(t.model) |
| t.model = None |
| t.optimizer = None |
| gc.collect() |
| assert w() is None |
|
|
|
|
| def test_trainer_release_drops_every_model_reference(tiny, tmp_workdir): |
| t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=4, grad_accum=2, save_every_micro_steps=4) |
| t.train() |
| w = weakref.ref(t.model) |
| t.release() |
| gc.collect() |
| assert w() is None and t._probe_target is None and t.trunk_probe._model is None |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_a_failed_arm_time_self_test_is_a_halt_with_a_marker(tiny, tmp_workdir): |
| """THE RETRIED-DETERMINISTIC-FAILURE DEFECT (U01-21): a rank-0 self-test/probe failure was a |
| plain RuntimeError, so no HALT.json was written and the supervisor relaunched it.""" |
| run = os.path.join(tmp_workdir, "a") |
| t, _ = _tiny_trainer(tiny, run, n_batches=4, grad_accum=2, armed=("logps_chosen",), save_every_micro_steps=4) |
| with pytest.raises(HaltError, match="SELF-TEST FAILED"): |
| t.train() |
| assert read_halt_marker(run) is not None |
|
|
|
|
| def test_a_failed_trunk_probe_is_a_halt_with_a_marker(tiny, tmp_workdir, monkeypatch): |
| run = os.path.join(tmp_workdir, "a") |
| t, _ = _tiny_trainer(tiny, run, n_batches=4, grad_accum=2, save_every_micro_steps=4) |
|
|
| def boom(step, **kw): |
| raise RuntimeError("TRUNK MOVEMENT PROBE FAILED at step 2") |
|
|
| monkeypatch.setattr(t.trunk_probe, "check", boom) |
| with pytest.raises(HaltError, match="TRUNK MOVEMENT PROBE FAILED"): |
| t.train() |
| assert read_halt_marker(run) is not None |
|
|
|
|
| def test_a_missing_tripwire_input_is_still_the_original_keyerror(tiny, tmp_workdir): |
| """Scope guard: only probe / self-test failures are converted; the KeyError contract of the |
| per-step tripwire check (test_control_plane) is unchanged.""" |
| t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), armed=("logps_chosen",)) |
| t.tripwires["logps_chosen"].INPUT_KEYS = () |
| with pytest.raises(KeyError): |
| t.train() |
| assert trainer_mod.HaltError is HaltError |
|
|
