fix: publish checkpoints atomically

- Write checkpoint payloads to a hidden sibling staging directory, add a versioned checksum manifest, fsync the completed payload, and publish it with an atomic rename
- Republishing an existing step retires the old payload under a hidden sibling name before the atomic rename, so re-runs into the same output directory replace the previous checkpoint instead of raising FileExistsError
- Keep legacy checkpoints loadable, add optional checksum verification, and align metric flushing with checkpoint publication

Co-authored-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

docs/developer/architecture.md

@@ -259,8 +259,8 @@ classDiagram
             +dict meta
             +dict config
             +save(save_dir)
-            +load(save_dir, broadcast) Checkpoint
-            +load_any(save_dir, broadcast) Optional[Checkpoint]
+            +load(save_dir, broadcast, verify_checksums) Checkpoint
+            +load_any(save_dir, broadcast, verify_checksums) Optional[Checkpoint]
         }
     }
 

docs/developer/docker-training.md

@@ -183,8 +183,14 @@ config.json
 model.safetensors
 optimizer.pt
 scheduler.pt
+manifest.json
 ```
 
+New checkpoints write `manifest.json` after every payload file, sync the complete
+staging directory, and then atomically rename that directory into place. Legacy
+checkpoints without a manifest remain resumable when the original required files
+are complete.
+
 `start` resumes the latest complete checkpoint and ignores partial writes. If no
 complete checkpoint exists, `/models/base/config.json` and
 `/models/base/model.safetensors` are required. `stop` sends `SIGTERM`; the

docs/guides/distributed.md

@@ -165,7 +165,9 @@ Checkpoints are saved by **rank-0 only**. The flow:
    - `ddp`: `model.module.state_dict()`
    - `fsdp`: `unshard()` → `full_tensor()` → `reshard()` (collective on all ranks, result kept only on rank-0)
 3. Non-rank-0 ranks get `None` — the save is skipped.
-4. Rank-0 writes `meta.json`, `config.json`, `model.safetensors`, and optional `{key}.pt` (optimizer/scheduler state).
+4. Rank-0 writes metadata, weights, optional optimizer/scheduler state, and a
+   checksum manifest to a hidden sibling directory, then atomically renames the
+   complete checkpoint into place.
 
 > **FSDP note**: Even though only rank-0 saves, all ranks must participate in `unwrap_model` because `unshard()` and `full_tensor()` are collective operations. The barriers in `checkpoint_context` keep all ranks in lockstep.
 

docs/guides/training.md

@@ -196,11 +196,19 @@ Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoi
 
 ```
 Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config)
-  ├── save(save_dir)    meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
-  └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
+  ├── save(save_dir)    atomically publishes manifest.json + metadata + weights + optional {key}.pt
+  └── load(save_dir, broadcast=False, verify_checksums=False)    loads locally or broadcasts from rank-0
 ```
 
-`Checkpoint.save()` writes whenever it is called. During training, `CheckpointCallback` uses the executor checkpoint context so only rank 0 receives a state dict and calls `save()`.
+`Checkpoint.save()` writes to a hidden sibling staging directory, records file
+sizes and SHA-256 checksums in `manifest.json`, flushes the files, and atomically
+renames the completed directory into place. Published checkpoint directories are
+immutable: saving to an existing non-empty path raises `FileExistsError`. Legacy
+checkpoints without a manifest remain loadable. Pass `verify_checksums=True` when
+loading to hash every published file.
+
+During training, `CheckpointCallback` uses the executor checkpoint context so
+only rank 0 receives a state dict and calls `save()`.
 
 Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
 Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.