fully_shard: load optim state not tested

[rohan-varma]

🚀 The feature, motivation and pitch

https://github.com/pytorch/pytorch/blob/master/test/distributed/_composable/fully_shard/test_fully_shard_optim_checkpoint.py are unittests for fully_shard + optim checkpoint, but these do not appear to actually test the load API.

Marking as high pri due to the composability effort requiring this.

Alternatives

No response

Additional context

No response

cc @ezyang @gchanan @zou3519 @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @vincentqb @jbschlosser @albanD @janeyx99

[janeyx99]

@rohan-varma could you specify what type of effort you would want to see here?

[rohan-varma]

@janeyx99 I think this is mostly a task for distributed team, sorry if it wasn't clear. We have some custom optim state dict loading APIs and need to test them in this file: https://github.com/pytorch/pytorch/blob/master/test/distributed/_composable/fully_shard/test_fully_shard_optim_checkpoint.py