[Checkpoint][Test] Add test for optimizer state_dict and resharding to 2d checkpoint test (#91092)
This PR updates the 2d checkpoint model state test to include:
1. optimizer state dict test
2. simple resharding test (pg change)
3. rename test
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91092
Approved by: https://github.com/fduwjj
