integrity of checkpoint pre-file v1

[chango9543]

Implement per-file SHA-256 integrity verification for checkpoints

Related issue: #6

Summary

Add an opt-in integrity verification mechanism for checkpoint directories. On save, every file in the checkpoint directory is hashed with SHA-256 and the hashes are written to _integrity_manifest.json. On load, the hashes are recomputed and compared before any weight is materialised into the model. Any mismatch indicates silent corruption (bit rot, truncated writes, transfer errors) and the load is rejected with CheckpointingException.

Background

See #6 — silent data corruption (bit rot) is a known failure mode on TB-scale checkpoint stores. torch_checkpointing currently has no content-level verification before load, so corrupted weights silently reach the model without loss blowing up immediately.

Design

Granularity: Per-file (v1). One hash per checkpoint file, computed by streaming the already-serialized bytes. Layout-agnostic — works for TorchSerialization, SafetensorsSerialization, or any serialization format.

Algorithm: Hardcoded SHA-256 (matches Megatron-LM's choice). The _compute_hash_dispatch + _collect_hashes_dispatch pattern allows adding future algorithms (e.g. blake3, xxhash64) with zero changes to call sites — just implement a new function and add one branch. Only "sha256" is implemented now; other values raise NotImplementedError.

Manifest format (_integrity_manifest.json):

{
  "version": 1,
  "granularity": "file",
  "algorithm": "sha256",
  "files": {
    "checkpoint_0.pt": "e3b0c4...",
    ".metadata": "a1b2c3..."
  }
}

• algorithm and granularity are recorded in the manifest so the verifier reads them from the manifest, not the current config — an old manifest can still be verified even if the default has changed.
• Files are sorted by name for byte-identical manifests across runs/ranks.
• The manifest itself is excluded from hashing (chicken-and-egg).

Manifest write location: After finalize_callback in CheckpointWriter.write(). At this point files are written, the barrier has passed, and the directory has been renamed to the final path — so the manifest hashes exactly the on-disk layout a verifier will see.

Manifest verify location: Before CheckpointReader.read() in CheckpointLoader.load(). A corrupt checkpoint is rejected before any weight reaches the model.

Distributed-safe: Only rank 0 does the file I/O; the verdict (success or error string) is broadcast via broadcast_object_list(src=0) so all ranks raise in lock-step or proceed together.

Exception type preservation: FileNotFoundError (manifest absent) and CheckpointingException (hash mismatch) are broadcast separately via a type-tagged payload so callers can catch them independently. The loader catches FileNotFoundError and silently skips — old checkpoints remain loadable.

Hardcoded policies (no per-run knobs):

• Mismatch -> always raise CheckpointingException (loading known-corrupt weights is never legitimate).
• Missing manifest -> silently skip (old checkpoints remain loadable without ignorable WARNING noise).
• Extra-file detection -> disabled (would surprise users who legitimately place README / auxiliary files alongside a published checkpoint).

Config surface

A single boolean flag on both config classes:

# Save side
CheckpointWriterConfig(verify_integrity=True)

# Load side
CheckpointLoaderConfig(verify_integrity=True)

# Or via CheckpointManager:
CheckpointManagerConfig(
    save=SyncCheckpointSaverConfig(
        writer_config=CheckpointWriterConfig(verify_integrity=True),
    ),
    load=CheckpointLoaderConfig(verify_integrity=True),
)

Defaults to False — existing checkpoints and CI loops are unaffected.

Files changed

File	Type	Description
torch_checkpointing/integrity.py	New (~285 lines)	Core module: hash dispatch, manifest read/write, distributed verify
torch_checkpointing/checkpoint_writer.py	Modified (+21 lines)	CheckpointWriterConfig.verify_integrity; write() integration
torch_checkpointing/checkpoint_loader.py	Modified (+20 lines)	__init__ accepts config; load() integration
torch_checkpointing/config.py	Modified (+7 lines)	CheckpointLoaderConfig.verify_integrity
torch_checkpointing/checkpoint_manager.py	Modified (+1 line)	Passes config.load to CheckpointLoader
torch_checkpointing/__init__.py	Modified (+1 line)	Export CheckpointingException, MANIFEST_FILENAME
tests/test_integrity.py	New	pytest unit tests (16 cases)

Test coverage

16 unit tests covering:

• Hash computation correctness (known SHA-256 vector, large-file streaming)
• Dispatch routing (sha256 supported, unsupported algorithm raises NotImplementedError)
• Manifest write/load round-trip (format contract, self-exclusion, sorted keys)
• Verify paths (happy, corruption, missing file, collect-all-mismatches, missing manifest)
• Algorithm read from manifest (not from config)
• End-to-end save→corrupt→restore cycle
• Write with explicit/unsupported algorithm and granularity