Checkpoint Hashing

Hashing should be a first order integrity check in a checkpointing library. Hashing's primary use case is the verification that 2 checkpoints are exactly the same, and if not, to pinpoint exactly what tensors changed.

A natural question here would be, why not just check the bitwise equivalence of all of the files listed in a checkpoint instead then? Two reasons:

• It's annoying, because these checkpoints are usually stored on distributed filesystems (slow I/O), and also these checkpoints can be on the order of TBs, so I don't want to have to load every bit to check them
• If we were comparing checkpoints that should be the same but came from 2 runs with different parallelisms configured, we would first need to consolidate the checkpoints so that they're in the same parallelism to compare their files bitwise, which could be annoying and error-prone

This is a very common feature that users ask for, some use cases:

• I've consolidated my checkpoint from M to N files and I'm worried about if I've done it correctly
• I'm working on some new checkpointing feature (e.g. a new resharding algorithm) and I'm not sure if I've done it correctly
• I'm worried about bit flips and partially written checkpoints (e.g. https://github.com/pytorch/torchtitan/issues/4011)

Hashing has a write step and a verify step. One of the very important design questions will be whether hash write should live in checkpoint save (or in an offline util) and whether hash verify should live in checkpoint load (or in an offline util). The important thing to remember is that in general, hashing must be very efficient to live in the actual save/load APIs by default or it must be turned on only optionally.

For a first implementation I doubt simple will be efficient, and IMHO it should be an optional verification we can turn on in the checkpoint save and checkpoint load. I'm also not going to cover hashing resharding workflows here, because selection of the hash algorithm there is very difficult (but I still think it's very important and we should cover this functionality in future!).

Off the top of my head, here's how I'm thinking about the hashing:

1. On save, every rank writes a hash of their tensors by FQN (fully qualified name) e.g. {"model.layer.0.weight.a": ABC, "mode.layer.1.weight.b": DEF} in a separate file alongside the actual checkpoint, e.g. rank_1_checkpoint.json.
2. On non-resharding load, every rank loads rank_1_model.pt for instance and also rank_1_checkpoint.json. It checks the FQNs that it has in its state dict against what the expected hash should be from the JSON file.

thoughts? @meetv18 @aelavender