Repository: pytorch
Issue: Unable to checkpoint model and optimizer state when using Hybrid Sharding Strategy
Labels: high priority,oncall: distributed,triaged,module: fsdp,module: distributed_checkpoint

### 🐛 Describe the bug

Fully Sharded Data Parallel (FSDP) is a wrapper for sharding module parameters across data parallel workers. It supports various sharding strategies for distributed training including ShardingStrategy.HYBRID_SHARD. In this strategy, the parameters are fully sharded within a node and replicated across nodes. 
Experimenting with ShardingStrategy.HYBRID_SHARD, we were able to successfully complete multiple iterations (forward and backward passes). However, trying to checkpoint the model and optimizer states throws errors. 

Relevant Code Snippet to retrieve the model and optimizer states:

```
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType

with FSDP.state_dict_type(
    model,
    StateDictType.SHARDED_STATE_DICT,
):
    model_state = model.state_dict()
    optim_state = FSDP.optim_state_dict(model, optimizer)
```
Relevant Code Snippet to save the model and optimizer states:
```
from torch.distributed._shard.checkpoint import (
    FileSystemWriter,
    SavePlan,
    save_state_dict,
)
from torch.distributed.checkpoint.default_planner import DefaultSavePlanner

save_state_dict(state_dict=model, storage_writer=writer, planner=DefaultSavePlanner())
if optimizer is not None:
    torch.save(optimizer, os.path.join(save_name, f"optim_{rank}.pth"))
```
Error while retrieving the optimizer state:

```
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2072, in all_gather_object
object_list[i] = _tensor_to_object(tensor, tensor_size)object_list[i] = _tensor_to_object(tensor, tensor_size)dist.all_gather_object(object_list, processed_state)

```

Error while writing model state:

```

save_state_dict(state_dict=model, storage_writer=writer, planner=DefaultSavePlanner())
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/checkpoint/state_dict_saver.py", line 104, in save_state_dict
central_plan = distW.reduce_scatter("plan", local_step, global_step)
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/checkpoint/utils.py", line 175, in reduce_scatter
all_data = self.gather_object(local_data)
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/checkpoint/utils.py", line 106, in gather_object
dist.gather_object(
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1451, in wrapper
return func(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2166, in gather_object
gather(
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1451, in wrapper
return func(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2713, in gather
work = default_pg.gather(output_tensors, input_tensors, opts)

```


### Versions

PyTorch Nightly Build - 05-31-2023 

@lessw2020 @chauhang 

cc @ezyang @gchanan @zou3519 @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu