Repository: pytorch
Issue: Unable to resume job using FSDP with 64 nodes, errors appeared during loading sharded optimizer state dict 
Labels: triaged,module: fsdp

### 🐛 Describe the bug

I am following this example code to save and load optimizer states using load_sharded_optimizer_state_dict: https://github.com/pytorch/pytorch/blob/c75e064dd6a2f800476bc84d4f27d7f49cedd055/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py

The resuming went well on 1 node and 8 nodes jobs. But when I tested it on 64 nodes (512 workers) with large model size, two errors appeared on two workers:
Here is first error:

> optim_state = dist_cp_op.load_sharded_optimizer_state_dict(
    File "/opt/conda/lib/python3.9/site-packages/torch/distributed/checkpoint/optimizer.py", line 282, 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 70, 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 181, in shard
      dist.scatter(
    File "/opt/conda/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 1451, in wrapper
      return func(*args, **kwargs)
    File "/opt/conda/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2774, in scatter
      _check_tensor_list(scatter_list, "scatter_list")
    File "/opt/conda/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 599, in _check_tensor_list
      raise RuntimeError(
  RuntimeError: Invalid function argument. Expected parameter `scatter_list` to be of type List[torch.Tensor].

And second error:

> File "/src/shopqa/generative/workflows/seq2seq/train/snapshot.py", line 363, in load_model_optim_final
      optim_state = dist_cp_op.load_sharded_optimizer_state_dict(
    File "/opt/conda/lib/python3.9/site-packages/torch/distributed/checkpoint/optimizer.py", line 282, 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 70, 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 172, in shard
      assert local_tensor is not None
  AssertionError

Also to confirm the saved checkpoint is not corrupted, I ran the same resuming job on 8 nodes with 64 workers. The checkpoint was loaded without errors.
I am wondering is there anything missing when I used load_sharded_optimizer_state_dict in large-scale training? 


### Versions

PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Amazon Linux 2 (x86_64)
GCC version: (GCC) 7.3.1 20180712 (Red Hat 7.3.1-15)
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.26

Python version: 3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:36:39) [GCC 10.4.0] (64-bit runtime)
Python platform: Linux-4.14.301-224.520.amzn2.x86_64-x86_64-with-glibc2.26
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 515.65.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              96
On-line CPU(s) list: 0-95
Thread(s) per core:  2
Core(s) per socket:  24
Socket(s):           2
NUMA node(s):        2
Vendor ID:           GenuineIntel
CPU family:          6
Model:               85
Model name:          Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
Stepping:            7
CPU MHz:             1522.122
BogoMIPS:            5999.99
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            1024K
L3 cache:            36608K
NUMA node0 CPU(s):   0-23,48-71
NUMA node1 CPU(s):   24-47,72-95
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology
nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch in
vpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke

Versions of relevant libraries:
[pip3] botorch==0.6.0
[pip3] flake8==6.0.0
[pip3] gpytorch==1.9.1
[pip3] mypy==0.991
[pip3] mypy-boto3-batch==1.26.34
[pip3] mypy-boto3-ec2==1.26.96
[pip3] mypy-boto3-iam==1.26.91
[pip3] mypy-boto3-s3==1.26.62
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==1.5.10
[pip3] pytorch-triton==2.1.0+7d1a95b046
[pip3] torch==2.0.1
[pip3] torch-model-archiver==0.7.1
[pip3] torch-tb-profiler==0.4.1
[pip3] torchaudio==2.1.0.dev20230523+cu117
[pip3] torchdata==0.6.1
[pip3] torchmetrics==0.10.3
[pip3] torchserve==0.7.1
[pip3] torchsnapshot-nightly==2023.3.15
[pip3] torchtext==0.15.1
[pip3] torchvision==0.16.0.dev20230523+cu117
[pip3] torchx-nightly==2023.4.21
[pip3] triton==2.0.0
[conda] pytorch-triton            2.1.0+7d1a95b046          pypi_0    pypi
[conda] torch                     2.0.1                    pypi_0    pypi
[conda] torch-tb-profiler         0.4.1                    pypi_0    pypi
[conda] torchaudio                2.1.0.dev20230523+cu117          pypi_0    pypi
[conda] torchdata                 0.6.1                    pypi_0    pypi
[conda] torchsnapshot-nightly     2023.3.15                pypi_0    pypi
[conda] torchvision               0.16.0.dev20230523+cu117          pypi_0    pypi
[conda] torchx-nightly            2023.4.21                pypi_0    pypi

cc @zhaojuanmao @mrshenli @rohan-varma @awgu