[ckpt-rewr] Get Optim State Dict Util API

[eracah]

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
• specify 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

[mvpatel2000]

@eracah is this a refactor or is it adding anything new anywhere? makes it a bit easier to review if I know what parts I need to carefully read through. It seems mostly copy paste but as helper fn?

[eracah]

@eracah is this a refactor or is it adding anything new anywhere? makes it a bit easier to review if I know what parts I need to carefully read through. It seems mostly copy paste but as helper fn?

a lot if it is a refactor, but it adds ignore, include, precision, and explicit cpu_offload control

[mvpatel2000]

LGTM!