[ckpt-rewr] Get Optim State Dict Util API

[eracah]

trafficstars

What does this PR do?

Adds an API for extracting optimizer state dict from a model and optimizer object.

State dict generation is a necessary operation before the save AND load of a checkpoint. Currently in composer it is coupled with the State, and not very readable, hard to extend, hard to test, and hard for users to harness to do custom things. As such, we present a function to generate state_dict for the optimizer decoupled from State as a standalone function. By making an explicit function for the optimizer, it’s easier to test because we have a standalone function (we don’t have to make a dummy State function). Moreover, it’s easier to save each state dict as a separate file Also, an advanced user can just call these functions themselves if they have a custom, advanced script or callback.

This state dict generation function enables:

• generating sharded or full state dicts
• generating state dicts of different precision
• specify keys to include
• sprecify keys to exclude

These are all options that will be useful for save and load. Because save and load require state dict generation, we need these options in state dict generation as well

GRT-2903