File size: 3,875 Bytes
bcc8d35 | 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 | """G3 audit regressions that need two real ranks (CPU, gloo): U01-3.
Reuses the 2-rank harness pieces of test_trainer_dist (`_build_trainer`,
`_free_port`); only the child entry point is local, because spawn resolves the
body by module and `test_trainer_dist._BODIES` cannot see bodies defined here.
"""
from __future__ import annotations
import json
import os
import sys
import traceback
import torch.multiprocessing as mp
from test_trainer_dist import JOIN_TIMEOUT_S, WS, _build_trainer, _free_port
def _body_pin_on_rank1_only(rank: int, run_dir: str) -> dict:
"""The R5 pin flag is set on RANK 1 ONLY (a rank-local tripwire fired
there), then both ranks train through saves at micro 4 and 8 with keep-N=1."""
trainer, _model = _build_trainer(run_dir, init_seed=0)
trainer.cfg.train.keep_checkpoints = 1
if rank == 1:
trainer._pin_next_save = True
trainer.train()
root = os.path.join(run_dir, "checkpoints")
return {
"dirs": sorted(d for d in os.listdir(root) if not d.endswith(".partial")),
"pinned": sorted(trainer.pinned_checkpoints),
"flag_left": bool(trainer._pin_next_save),
}
_BODIES = {"pin_rank1": _body_pin_on_rank1_only}
def _child(rank: int, port: int, mode: str, run_dir: str) -> None:
import torch.distributed as dist
os.environ.update(
RANK=str(rank),
WORLD_SIZE=str(WS),
LOCAL_RANK=str(rank),
MASTER_ADDR="127.0.0.1",
MASTER_PORT=str(port),
)
from align.utils.dist import init_distributed
init_distributed(backend="gloo")
code = 0
try:
result = _BODIES[mode](rank, run_dir)
with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f:
json.dump(result, f)
except BaseException:
with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f:
f.write(traceback.format_exc())
code = 1
finally:
try:
dist.destroy_process_group()
except Exception:
pass
if code:
sys.exit(code)
def _run(mode: str, run_dir: str, monkeypatch) -> list[dict]:
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "")
ctx = mp.get_context("spawn")
port = _free_port()
procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)]
for p in procs:
p.start()
for p in procs:
p.join(JOIN_TIMEOUT_S)
hung = [p for p in procs if p.is_alive()]
for p in hung:
p.kill()
p.join()
errs = []
for r in range(WS):
ef = os.path.join(run_dir, f"{mode}_{r}.err")
if os.path.isfile(ef):
errs.append(f"rank {r}:\n{open(ef).read()}")
assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs)
assert not errs, "\n".join(errs)
return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)]
def test_u01_3_pin_set_on_a_non_zero_rank_protects_the_checkpoint_on_rank_0(
tmp_workdir, monkeypatch
):
"""U01-3. `_pin_next_save` is raised by a RANK-LOCAL tripwire verdict, but
pruning and resume.pt are rank 0's. A pin that exists on rank 1 only left
rank 0's pinned set empty: keep-N deleted the very checkpoint meant as the
rewind point and a resume forgot the pin. The pin must be collective.
Mutant: drop the all_reduce in save_checkpoint -> dirs == [micro_00000008]."""
rows = _run("pin_rank1", tmp_workdir, monkeypatch)
for r in rows:
assert r["pinned"] == [4], rows
assert r["flag_left"] is False, rows
assert rows[0]["dirs"] == ["micro_00000004", "micro_00000008"], rows
import torch
state = torch.load(
os.path.join(tmp_workdir, "checkpoints", "micro_00000004", "resume.pt"),
weights_only=True,
map_location="cpu",
)
assert state["pinned_checkpoints"] == [4]
|