[zyhu-hu]

Title: Crash-safe checkpoint writes, correct best.pth, own-face halo adjoint by default

Summary

1. Checkpoint writes no longer corrupt latest.pth or best.pth

fabric.save uses lightning's _atomic_save. That writes a temp file in the system temp dir and then shutil.moves it onto the target. When the temp dir is on a different filesystem from the run dir (node-local /tmp vs. a shared run dir), the move copies into the existing file in place. Two bugs follow:

• best.pth tracks the latest state, not the best. best.pth is a hard link to latest.pth, so every later save of latest.pth overwrites it too.
• A killed job leaves a checkpoint the next job can't load. A walltime kill during the in-place copy leaves a truncated latest.pth.

New screamcast/checkpoint_io.py:

• Saving: it writes <path>.tmp in the run dir, fsyncs it, CRC-checks it and os.replaces it onto the path. No checkpoint file is ever rewritten in place.
• best.pth: it is linked to the verified .tmp before the rename. A kill between the two steps can no longer leave best.pth on an older state.
• Stale temp files: leftover *.pth.tmp files are removed when a run starts and before each save. fsspec ignores a failed copy onto an existing file, so a stale but complete tmp could otherwise pass the CRC check and replace a newer latest.pth.
• Resume: if latest.pth fails its CRC check, the run falls back with a loud warning to the newest valid steps_*.pth, keeping the lower best_valid_loss recorded in best.pth. With no usable fallback it stops with a clear error. A missing latest.pth next to steps_*.pth prints a warning. The selection logic lives in choose_resume_checkpoint() and apply_best_valid_loss_floor().
• Multi-rank failures: every failure on rank 0 (save, verify, link, resume selection) is broadcast, so all ranks raise instead of hanging in a barrier.

2. Rollout: the halo adjoint stays within each cube face by default

This changes the default behaviour of scripts/ace/run_screamcast_nudged.py.

• --halo-adjoint-cross-face now defaults to False. --halo-adjoint keeps the overlap-add at tile seams inside a face and drops it across cube-face edges.
• The reason: cross-face halo targets are KNN-interpolated off the neighbour face's grid. Mixing them back in over-smooths the outermost pixels of every face (docs/halo-adjoint-cross-face.md).
• --halo-adjoint-cross-face restores the old behaviour. The setting is still recorded in the output zarr attrs.

3. README: recommended training schedule

The production checkpoint was trained in four stages, each resuming the previous stage's best.pth:

1. 64×64 pretraining at 1e-4 for ~10M images.
2. A cosine anneal to 1e-6 over ~3M images (dampened_cosine_with_hard_restarts, cosine_n_cycles=1, warmup_steps=0).
3. A 128×128 fine-tune at 1e-5 (better: tile + 2×halo) for ~0.5–1M images.
4. 2-step, then 4-step fine-tunes at 1e-5.

The README also recommends training with pipeline.do_precip_relu=False and clipping precipitation at inference instead. With the ReLU in the forward pass, a precipitation head can get stuck below zero with no gradient.

Tests

New tests/test_checkpoint_io.py covers:

• the cross-filesystem move and a kill mid-copy (each also run against plain fabric.save as a strict xfail, to show the bug);
• the resume selection: CRC fallback to steps_*, the best_valid_loss floor, the no-fallback error;
• removal of stale tmp files;
• a two-rank gloo run where rank 0 fails and both ranks must raise.

Run on e1b3e0c in a GPU container, with S3 access to the SCREAM data. Each result was compared with main (8db1c74):

• Lint: python ci/check_licenses.py and pre-commit run -a pass.
• pytest, everything except tests/distributed: all pass apart from test_training_launch, the same as on main. That includes the dataloader tests against the real SCREAM store, and test_characterization_inference: PR 10's outputs for all three production configs match reference outputs generated from main.
• test_training_launch: times out at its 180 s limit on both main and this branch, while reading data from S3. It runs the pytest config, which doesn't save checkpoints, so the new write path is covered by tests/test_checkpoint_io.py only.
• Not run: tests/distributed.

🤖 Generated with Claude Code