Repository: pytorch/pytorch

Issue:
```load_sharded_optimizer_state_dict``` error on multi node
### 🐛 Describe the bug

I am hitting the following error when loading optimizer using FSDP with 2 nodes. No issues when saving/loading optimizer with single node.

```
optim_state = load_sharded_optimizer_state_dict(
    model_state_dict=state_dict["model"],
    optimizer_key="optimizer",
    storage_reader=checkpoint.FileSystemReader(checkpoint_dir),
)
```
Error:
```
optim_state = load_sharded_optimizer_state_dict(
  File "/opt/conda/lib/python3.9/site-packages/torch/distributed/checkpoint/optimizer.py", line 281, in load_sharded_optimizer_state_dict
    state_dict[key] = _shard_tensor(
  File "/opt/conda/lib/python3.9/site-packages/torch/distributed/_shard/api.py", line 69, in _shard_tensor
    st = sharding_spec.shard(tensor, src_rank=src_rank, process_group=process_group)
  File "/opt/conda/lib/python3.9/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py", line 166, in shard
    local_tensor = torch.empty(
RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
```





### Versions

```
PyTorch version: 2.0.0.dev20230227+cu117
```

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


Discussion:
cc @wz337 for distributed checkpoint issues

> RuntimeError: CUDA error: invalid device ordinal

Looks like the checkpoint was saved on a node with more devices than where it is loaded?
@cavdard 
load_sharded_optimizer_state_dict needs to be used in conjunction with FSDP sharded optimizer state.

Question: did you set the state_dict_type to sharded_state_dict? Example: https://github.com/wz337/pytorch/blob/master/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py#L91


Encounter the same problem here.
It seems to try to create a tensor on device like `cuda:31`, which is a global rank index.
@wz337 
Yes I am using `StateDictType.SHARDED_STATE_DICT`
```
    with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
        state_dict = {
            "model": model.state_dict(),
            "scheduler": scheduler.state_dict(),
            .....
            }
 ```
> Encounter the same problem here. It seems to try to create a tensor on device like `cuda:31`, which is a global rank index.

Could you provide a bit context of your setup as well? Thanks!
> > Encounter the same problem here. It seems to try to create a tensor on device like `cuda:31`, which is a global rank index.
> 
> Could you provide a bit context of your setup as well? Thanks!

I used 8 nodes with 8 V100 on each node, and save/load the sharded state dict using the following code.

```python
with FSDP.state_dict_type(
    model,
    StateDictType.SHARDED_STATE_DICT,
):
    state_dict = {
        'model': model.state_dict(),
        'optimizer': FSDP.optim_state_dict(model, optimizer)
    }

    save_state_dict(state_dict, FileSystemWriter(ckpt_path))
```

```python
with FSDP.state_dict_type(
    model,
    StateDictType.SHARDED_STATE_DICT,
):
    storage_reader = FileSystemReader(ckpt_path)
    state_dict = {
        'model': model.state_dict(),
    }
    load_state_dict(state_dict, storage_reader)
    state_dict |= load_sharded_optimizer_state_dict(
        model_state_dict=state_dict['model'],
        optimizer_key='optimizer',
        storage_reader=storage_reader,
    )
```

Error occurred during `load_sharded_optimizer_state_dict` call.

Resolving patch:
diff --git a/torch/distributed/checkpoint/optimizer.py b/torch/distributed/checkpoint/optimizer.py
--- a/torch/distributed/checkpoint/optimizer.py
+++ b/torch/distributed/checkpoint/optimizer.py
@@ -255,7 +255,8 @@ def load_sharded_optimizer_state_dict(
         sharding_spec = ChunkShardingSpec(
             dim=0,
             placements=[
-                f"rank:{i}/cuda:{i}" for i in range(dist.get_world_size())
+                f"rank:{i}/cuda:{i % torch.cuda.device_count()}"
+                for i in range(dist.get_world_size())
             ],
         )
     else:


Test patch: