Improved checkpoint load/save and state_dict interfaces

I'm trying to interface Deepspeed with a new lightweight training code abstraction for my model collection (https://github.com/rwightman/pytorch-image-models)

I'm finding it quite challenging to write reasonable abstractions with how tighlty intertwined the code around checkpoints is with DeepSpeed. The checkpoint naming, disk io (save/load), and DS specific state dict handling is all wound together across a number of public/private functions in Engine.

I have a class that manages checkpoints, allows controlling the number of checkpoint files in the history via top metrics, etc. It also handles period recovery checkpoints vs epoch checkpoints depending on epoch length and the reliability of your training machines. There are ways to make this all work with the DS checkpoints but it'd be a giant hack and I'd be accessing internal DS state. I'd prefer to control all filenames myself, be able to track them, etc so I can delete when needed.

I'm also working on a filesystem abstraction that lets one choose between writing directly to a local disk, or cloud storage buckets, etc. Again not possible if DS handles all of the lower level IO.

I do realize there are specific requirements for saving state across the ranks. 100% understood. Would there not be a better interface that both clients can use AND the current save_checkpoint/load_checkpoint?

Something like a serialize/deserialze layer. Like existing save/load needs to be called per rank, and returns or takes dicts of state_dicts (ie, of https://github.com/microsoft/DeepSpeed/blob/ed3de0c21b1fea330de9c1a78a23ca33f340ef20/deepspeed/runtime/engine.py#L1905-L1918) and some identification metdata. The keys would correspond to the 'sub'-checkpoints like the zero_pp_rank_*, *_model_states, *_optim_states, etc. Identification would be any extra data (ie mp_rank) needed for the client to track the dicts so that they can get restored to the correct ranks, or this could be embedded in the keys in a parsable fashion or stored in the serialized dict.

Then the client can more easily decide how/when/where to save the files so long as they deal with matching the ranks correctly. For me at least, this would be much easier than a brittle hack job accessing DS internals or writing parallel code that will be quickly outdated.

Start with the checkpoint handling in deepspeed/runtime/engine.py around lines 1905-1918, including the existing save_checkpoint and load_checkpoint paths. Define the scope of a public serialization/deserialization layer that preserves rank and sub-checkpoint metadata while letting clients control filenames and storage. Done means the proposed interface works alongside the current checkpoint APIs.