diff --git a/param_decomp/configs.py b/param_decomp/configs.py index eea12ab23..b842b2431 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -1109,12 +1109,14 @@ class Cadence(BaseConfig): dense_log_phase: DenseLogPhase | None = None """Optional denser logging for early training; `None` means a flat `train_log_every`.""" save_every: PositiveInt | None = None - keep_last_n_checkpoints: PositiveInt | None = None + keep_last_n_checkpoints: PositiveInt """How many of the most-recent orbax `ckpts//` checkpoints to keep on disk - after each checkpoint write. `None` (the default) keeps all checkpoints — the - conservative choice for research where prior steps may matter. Opt in to e.g. `3` - for long jobs where disk pressure outweighs the value of intermediate checkpoints; - the final-step checkpoint is always included in the retained set.""" + after each checkpoint write; the final-step checkpoint is always in the retained + set. Required, and deliberately not defaulted: retention is the single biggest + lever a run has on shared storage, so every run states its own. An LM checkpoint + is 100-200 GB and ~85% of that is the `training` item nothing but resume reads, so + `40` is ~7 TB of trajectory tail. Keep it small; if you need old steps for their + decompositions, keep them and run `param_decomp_lab.tools.thin_training_state`.""" def should_log_train(self, step: int) -> bool: if self.dense_log_phase is not None and step < self.dense_log_phase.until_step: diff --git a/param_decomp/run.py b/param_decomp/run.py index 6116dd524..c9237cd49 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -501,7 +501,7 @@ def run_decomposition_training( # flash attention under the scan+cond masked forward). Explicit NamedShardings elsewhere # are unaffected. jax.set_mesh(mesh) - assert cadence.save_every is not None and cadence.keep_last_n_checkpoints is not None, cadence + assert cadence.save_every is not None, cadence save_every = cadence.save_every run.run_dir.mkdir(parents=True, exist_ok=True) diff --git a/param_decomp_lab/experiments/config.py b/param_decomp_lab/experiments/config.py index 20d08a6f3..b27d0e484 100644 --- a/param_decomp_lab/experiments/config.py +++ b/param_decomp_lab/experiments/config.py @@ -172,7 +172,7 @@ def assert_canonical_algorithm_config(cfg: ExperimentConfig) -> None: assert vu_opt.grad_clip_norm is not None, "components grad clip is part of the method" cadence = cfg.cadence - assert cadence.save_every is not None and cadence.keep_last_n_checkpoints is not None, cadence + assert cadence.save_every is not None, cadence def run_instance(cfg: ExperimentConfig, run_id: str) -> RunInstance: diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py index 2806790a6..59549f3f0 100644 --- a/param_decomp_lab/experiments/lm/load_run.py +++ b/param_decomp_lab/experiments/lm/load_run.py @@ -172,7 +172,6 @@ def _restore_decomposition( rules = placement.from_config_for_consumer(cfg.runtime.sharding, mesh, model.sites) abstract = jax.eval_shape(lambda: init_decomposition(model, cfg.ci_fn, init_key, mesh, rules)) - assert cfg.cadence.keep_last_n_checkpoints is not None, cfg.cadence manager = make_checkpoint_manager(run_dir / "ckpts", cfg.cadence.keep_last_n_checkpoints) resolved_step = manager.latest_step() if step is None else step assert resolved_step is not None, f"no checkpoints under {run_dir / 'ckpts'}" diff --git a/param_decomp_lab/tests/test_launch.py b/param_decomp_lab/tests/test_launch.py index 53e28aaba..9e418ed02 100644 --- a/param_decomp_lab/tests/test_launch.py +++ b/param_decomp_lab/tests/test_launch.py @@ -41,7 +41,7 @@ "loss_metrics": [{"type": "FaithfulnessLoss", "coeff": 1.0}], }, "runtime": {"device": "cuda:0", "launch": "inline", "dp": 1, "sharding": "zero1"}, - "cadence": {"train_log_every": 1}, + "cadence": {"train_log_every": 1, "keep_last_n_checkpoints": 2}, "target": { "spec": { "kind": "hf", diff --git a/param_decomp_lab/tests/test_thin_training_state.py b/param_decomp_lab/tests/test_thin_training_state.py new file mode 100644 index 000000000..9058bf67a --- /dev/null +++ b/param_decomp_lab/tests/test_thin_training_state.py @@ -0,0 +1,134 @@ +"""The checkpoint thinner: which `training` items a sweep drops, and that dropping one +leaves the step's `decomposition` and the rest of the run untouched.""" + +from pathlib import Path + +import pytest +from pytest import MonkeyPatch + +from param_decomp_lab.tools.thin_training_state import ( + RESUME_WINDOW, + Checkpoint, + Run, + Skip, + Thin, + execute, + plan_run, + read_run, +) + + +def _run(run_id: str, steps: list[int], live: bool) -> Run: + return Run( + run_id, + live, + tuple(Checkpoint(step, Path(f"/{run_id}/{step}/training"), 1) for step in steps), + ) + + +def _write_ckpts(run_dir: Path, steps: list[int], with_training: bool = True) -> None: + for step in steps: + step_dir = run_dir / "ckpts" / str(step) + (step_dir / "decomposition" / "d").mkdir(parents=True) + (step_dir / "decomposition" / "d" / "chunk").write_bytes(b"product") + (step_dir / "_CHECKPOINT_METADATA").write_text("{}") + if with_training: + (step_dir / "training" / "d").mkdir(parents=True) + (step_dir / "training" / "d" / "chunk").write_bytes(b"trajectory tail") + + +def test_keeps_the_resume_window(): + plan = plan_run(_run("p-aaaaaaaa", [1, 2, 3, 4], live=False), keep_newest=2) + assert isinstance(plan, Thin) + assert [c.step for c in plan.victims] == [1, 2] + + +def test_zero_keep_thins_everything_on_a_dead_run(): + plan = plan_run(_run("p-aaaaaaaa", [1, 2], live=False), keep_newest=0) + assert isinstance(plan, Thin) + assert [c.step for c in plan.victims] == [1, 2] + + +def test_live_run_is_floored_at_the_resume_window(): + plan = plan_run(_run("p-aaaaaaaa", list(range(10)), live=True), keep_newest=0) + assert isinstance(plan, Thin) + assert [c.step for c in plan.victims] == list(range(10 - RESUME_WINDOW)) + + +def test_run_shorter_than_the_window_is_skipped(): + assert plan_run(_run("p-aaaaaaaa", [7], live=True), keep_newest=0) == Skip( + "p-aaaaaaaa", "inside-resume-window" + ) + + +def test_pre_split_run_has_no_training_item(tmp_path: Path): + run_dir = tmp_path / "p-aaaaaaaa" + _write_ckpts(run_dir, [100], with_training=False) + assert read_run(run_dir, frozenset()) == Skip("p-aaaaaaaa", "no-training-item") + + +def test_read_run_sizes_and_orders_thinnable_steps(tmp_path: Path): + run_dir = tmp_path / "p-bbbbbbbb" + _write_ckpts(run_dir, [2000, 1000]) + run = read_run(run_dir, frozenset({"p-bbbbbbbb"})) + assert isinstance(run, Run) + assert run.live + assert [c.step for c in run.thinnable] == [1000, 2000] + assert {c.nbytes for c in run.thinnable} == {len(b"trajectory tail")} + + +def test_execute_drops_training_and_spares_the_decomposition(tmp_path: Path): + run_dir = tmp_path / "p-cccccccc" + _write_ckpts(run_dir, [1000, 2000, 3000]) + trash = tmp_path / "trash" + trash.mkdir() + + run = read_run(run_dir, frozenset()) + assert isinstance(run, Run) + plan = plan_run(run, keep_newest=1) + assert isinstance(plan, Thin) + assert [c.step for c in list(execute(plan, trash))] == [1000, 2000] + + ckpts = run_dir / "ckpts" + assert not (ckpts / "1000" / "training").exists() + assert not (ckpts / "2000" / "training").exists() + assert (ckpts / "3000" / "training" / "d" / "chunk").exists() + for step in (1000, 2000, 3000): + assert (ckpts / str(step) / "decomposition" / "d" / "chunk").read_bytes() == b"product" + assert (ckpts / str(step) / "_CHECKPOINT_METADATA").exists() + assert list(trash.iterdir()) == [] + + +def test_a_second_sweep_finds_nothing_left(tmp_path: Path): + run_dir = tmp_path / "p-dddddddd" + _write_ckpts(run_dir, [1000, 2000]) + trash = tmp_path / "trash" + trash.mkdir() + + run = read_run(run_dir, frozenset()) + assert isinstance(run, Run) + plan = plan_run(run, keep_newest=1) + assert isinstance(plan, Thin) + list(execute(plan, trash)) + + again = read_run(run_dir, frozenset()) + assert isinstance(again, Run) + assert plan_run(again, keep_newest=1) == Skip("p-dddddddd", "inside-resume-window") + + +def test_a_run_owned_by_someone_else_is_never_touched(tmp_path: Path, monkeypatch: MonkeyPatch): + run_dir = tmp_path / "p-eeeeeeee" + _write_ckpts(run_dir, [1000, 2000]) + monkeypatch.setattr("os.getuid", lambda: run_dir.stat().st_uid + 1) + assert read_run(run_dir, frozenset()) == Skip("p-eeeeeeee", "not-owned") + + +@pytest.mark.parametrize("keep_newest", [0, 1, RESUME_WINDOW, 5]) +def test_a_live_run_always_keeps_a_resumable_step(keep_newest: int): + run = _run("p-ffffffff", list(range(6)), live=True) + plan = plan_run(run, keep_newest) + kept = ( + set(range(6)) - {c.step for c in plan.victims} if isinstance(plan, Thin) else set(range(6)) + ) + assert max(kept) == 5 + assert len(kept) >= RESUME_WINDOW diff --git a/param_decomp_lab/tools/__init__.py b/param_decomp_lab/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/param_decomp_lab/tools/thin_training_state.py b/param_decomp_lab/tools/thin_training_state.py new file mode 100644 index 000000000..bfed8c2ce --- /dev/null +++ b/param_decomp_lab/tools/thin_training_state.py @@ -0,0 +1,223 @@ +"""Reclaim the trainer-only half of checkpoints that will never be resumed. + +A PD checkpoint step holds two orbax items (see `param_decomp/checkpoint.py`): +`decomposition` — the trained product every consumer reads — and `training` — the +optimizer moments, persistent adversaries and step counter that only trainer resume +touches. On the LM runs `training` is 80-90% of the bytes, and a run keeps one copy +per retained step, so a `keep_last_n_checkpoints: 40` job carries ~40x the trajectory +tail nobody will ever read. Three cluster-wide storage sweeps have gone at this by +hand; this is that operation, committed. + +Thinning removes `ckpts//training` and nothing else. The step keeps its +`decomposition`, so harvest, autointerp, clustering, CI statistics and the offline +PGD anatomy all still load it (`param_decomp_lab.experiments.lm.load_run.open_jax_run` +restores the `decomposition` item alone). What a thinned step loses is resumability: +`restore_latest` asks for both items and raises on the missing one. + +That loss is why `--keep-newest` exists and defaults to `RESUME_WINDOW`. A running +trainer holds its `CheckpointManager` in memory and never re-reads the directory, so +thinning older steps underneath it is invisible to it; what would bite is a SLURM +requeue landing on a newest step whose `training` is gone. Runs referenced by a live +job are therefore floored at `RESUME_WINDOW` however low `--keep-newest` is set. + + python -m param_decomp_lab.tools.thin_training_state # dry run, whole out dir + python -m param_decomp_lab.tools.thin_training_state --delete + python -m param_decomp_lab.tools.thin_training_state --keep_newest=0 --run_ids=p-abc12345 + +`_CHECKPOINT_METADATA` is deliberately left listing both items: it is read only by an +argument-less composite restore, which nothing in this repo issues, and rewriting a +file inside a live run's checkpoint tree buys nothing. +""" + +import os +import re +import shutil +import subprocess +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import fire + +from param_decomp_lab.infra.settings import PARAM_DECOMP_OUT_DIR, SLURM_LOGS_DIR + +RESUME_WINDOW = 2 +"""Newest checkpoints a run keeps resumable: the newest, plus one for the requeue that +lands while a save is in flight.""" + +RUN_ID = re.compile(rb"p-[0-9a-f]{8}") +LIVE_STATES = ("RUNNING", "PENDING", "COMPLETING", "CONFIGURING", "REQUEUED") +TRASH_DIR = ".thin_trash" + + +@dataclass(frozen=True) +class Checkpoint: + step: int + training: Path + nbytes: int + + +@dataclass(frozen=True) +class Run: + run_id: str + live: bool + thinnable: tuple[Checkpoint, ...] + """Ascending by step; steps already missing their `training` item are absent.""" + + +@dataclass(frozen=True) +class Thin: + run: Run + victims: tuple[Checkpoint, ...] + + +@dataclass(frozen=True) +class Skip: + run_id: str + reason: Literal["not-owned", "no-training-item", "inside-resume-window"] + + +RunPlan = Thin | Skip + + +# ---------------------------------------------------------------- pure core + + +def plan_run(run: Run, keep_newest: int) -> RunPlan: + """Which of a run's `training` items are safe to drop.""" + keep = max(keep_newest, RESUME_WINDOW) if run.live else keep_newest + victims = run.thinnable[: len(run.thinnable) - keep] + if not victims: + return Skip(run.run_id, "inside-resume-window") + return Thin(run, tuple(victims)) + + +def reclaimed(plans: list[RunPlan]) -> int: + return sum(c.nbytes for p in plans if isinstance(p, Thin) for c in p.victims) + + +def terabytes(nbytes: int) -> str: + return f"{nbytes / 1e12:.3f} TB" + + +# ---------------------------------------------------------------- the filesystem edge + + +def tree_bytes(root: Path) -> int: + return sum(f.stat().st_size for f in root.rglob("*") if f.is_file()) + + +def read_run(run_dir: Path, live_run_ids: frozenset[str]) -> Run | Skip: + if run_dir.stat().st_uid != os.getuid(): + return Skip(run_dir.name, "not-owned") + thinnable = sorted( + ( + Checkpoint(int(step_dir.name), training, tree_bytes(training)) + for step_dir in run_dir.glob("ckpts/*") + if step_dir.name.isdigit() + if (training := step_dir / "training").is_dir() + ), + key=lambda c: c.step, + ) + if not thinnable: + return Skip(run_dir.name, "no-training-item") + return Run(run_dir.name, run_dir.name in live_run_ids, tuple(thinnable)) + + +def live_run_ids() -> frozenset[str]: + """Run ids named in the stdout of every SLURM job that is not finished. + + The trainer announces its run id into its own log, so the log is the one honest + job -> run link (`scontrol`'s comment field is free-form and often absent). A live + job whose log we cannot read fails the sweep rather than silently freeing its run. + """ + jobs = subprocess.run( + ["scontrol", "show", "jobs", "-o"], capture_output=True, text=True, check=True + ).stdout + logs = { + Path(m.group(1)) + for line in jobs.splitlines() + if any(f"JobState={state}" in line for state in LIVE_STATES) + if (m := re.search(r"StdOut=(\S+)", line)) + if m.group(1).startswith(str(SLURM_LOGS_DIR)) + } + return frozenset( + m.group().decode() + for log in logs + if log.exists() + for m in RUN_ID.finditer(log.read_bytes()) + ) + + +def sweep(out_dir: Path, run_ids: tuple[str, ...] | None, keep_newest: int) -> list[RunPlan]: + live = live_run_ids() + run_dirs = ( + sorted((out_dir / "runs").glob("p-*")) + if run_ids is None + else [out_dir / "runs" / run_id for run_id in run_ids] + ) + return [ + read if isinstance(read := read_run(d, live), Skip) else plan_run(read, keep_newest) + for d in run_dirs + ] + + +def execute(plan: Thin, trash: Path) -> Iterator[Checkpoint]: + """Rename each victim out of the checkpoint tree before deleting it, so a concurrent + reader sees the item either whole or absent — never half-deleted.""" + for victim in plan.victims: + staged = trash / f"{plan.run.run_id}-{victim.step}-training" + victim.training.rename(staged) + shutil.rmtree(staged) + yield victim + + +# ---------------------------------------------------------------- CLI + + +def main( + out_dir: str | Path = PARAM_DECOMP_OUT_DIR, + run_ids: str | tuple[str, ...] | None = None, + keep_newest: int = RESUME_WINDOW, + delete: bool = False, +) -> None: + """Report (or, with --delete, drop) thinnable `training` items. + + out_dir: run root; defaults to $PARAM_DECOMP_OUT_DIR. + run_ids: restrict to these runs; default sweeps every run in `out_dir/runs`. + keep_newest: newest checkpoints per run left resumable. Live runs are floored at + RESUME_WINDOW. 0 thins every checkpoint of every non-live run — that forfeits + resumption for good, so pass it deliberately. + delete: actually delete. Default reports the plan and touches nothing. + """ + root = Path(out_dir) + only = (run_ids,) if isinstance(run_ids, str) else run_ids + plans = sweep(root, only, keep_newest) + + for plan in sorted( + (p for p in plans if isinstance(p, Thin)), + key=lambda p: -sum(c.nbytes for c in p.victims), + ): + steps = [c.step for c in plan.victims] + live = " LIVE" if plan.run.live else "" + size = terabytes(sum(c.nbytes for c in plan.victims)) + print(f"{plan.run.run_id}{live} {size} {len(steps)} steps {steps[0]}..{steps[-1]}") + + total = reclaimed(plans) + print(f"\n{sum(isinstance(p, Thin) for p in plans)} runs, {terabytes(total)} reclaimable") + if not delete: + print("dry run — pass --delete to remove") + return + + trash = root / TRASH_DIR + trash.mkdir(exist_ok=True) + freed = sum( + victim.nbytes for plan in plans if isinstance(plan, Thin) for victim in execute(plan, trash) + ) + trash.rmdir() + print(f"freed {terabytes(freed)}") + + +if __name__ == "__main__": + fire.Fire(main)