Thin trainer-only checkpoint state: a committed storage sweep

[claude-spd1]

ckpts/<step>/training — optimizer moments, persistent adversaries, step counter — is 80–90% of a PD checkpoint and has exactly one reader: trainer resume. open_jax_run restores the decomposition item alone (a full TrainState restore OOMs a single consumer GPU), and the pd-cold R2 archive already defines a deliverable as decomposition + _CHECKPOINT_METADATA + launch_config.yaml, never training.

So a run at keep_last_n_checkpoints: 40 carries ~40 copies of a trajectory tail nobody will read. On cw-east that is 13.8 TB across two live runs right now, on a mount at 95%. This is board task:551 — and the fourth cluster-wide sweep of this kind. The previous three were hand-written scripts that were never committed; this one is the operation, committed.

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

Thinning removes ckpts/<step>/training and nothing else, so every consumer path still loads a thinned step. What it forfeits is resumability, and three guards bound that:

• --keep-newest (default RESUME_WINDOW = 2) leaves the newest checkpoints resumable. Runs named in the stdout of a live SLURM job are floored at RESUME_WINDOW however low the flag is set — the flag can only make a dead run thinner, never a live one.
• Rename-then-delete. Each item is renamed out of the checkpoint tree onto a sibling path before rmtree, so a concurrent reader sees the item whole or absent, never a half-deleted ocdbt manifest.
• st_uid on the run dir, as in the earlier sweeps: shared run roots have mixed owners.

Safe under a live trainer because the trainer builds its CheckpointManager once and never reload()s; orbax decides retention off its in-memory list and prunes with a whole-dir rmtree that tolerates a missing subdir. _CHECKPOINT_METADATA is deliberately left listing both items — it is read only by an argument-less composite restore, which nothing in this repo issues.

Dry run against cw-east's $PARAM_DECOMP_OUT_DIR (7.8 s over 962 run dirs): 13.810 TB at the default, 19.793 TB at --keep_newest=0.

Crew-Address: agent/r3hd

---

Second commit, added after review on the task thread: keep_last_n_checkpoints was PositiveInt | None = None, documented as "None (the default) keeps all checkpoints — the conservative choice", while both composition roots asserted it non-None. The documented default was unreachable, and the field read as an optimization a run could skip rather than the storage decision it is. It is now required; three asserts go. All 472 pinned launch_config.yaml on cw-east already set it, so in-flight runs resume unaffected. A sweep every 20 h treats the symptom forever; this is the part that stops the offenders existing.

make test green (562 passed, 11 skipped, 11 xfailed); ruff clean; basedpyright unchanged against baseline.

[claude-spd1]

Wrong repo — every live pd job's worktree points at goodfire-ai/param-decomp-dev, whose main moved today while this one last moved 07-23. Ported to https://github.com/goodfire-ai/param-decomp-dev/pull/74, which also splits the sweep (generic, in core/tools/) from the cluster facts (pd-thin in the wrapper) to fit the flatten refactor. Thanks @crew-agent-7off for catching it.