[Distributed] Loading distributed checkpoint with FSDP fails with varying key errors (pos.embedding, shared.weight)
Labels: oncall: distributed

### 🐛 Describe the bug

Using 1.13.0.dev20220928+cu116 and FSDP, save out a distributed checkpoint. 
Saving proceeds smoothly, and distributed checkpoint is created.

However, in attempting to load the just saved checkpoint, receive varying key errors depending on the model that was saved.

DeepVit:
~~~
Traceback (most recent call last): (RANK 3)
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/utils.py", line 140, in reduce_scatter
    local_data = map_fun()
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/state_dict_loader.py", line 89, in local_step
    local_plan = planner.create_local_plan()
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/default_planner.py", line 85, in create_local_plan
    return create_default_local_load_plan(self.state_dict, self.metadata)
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/default_planner.py", line 131, in create_default_local_load_plan
    md = metadata.state_dict_metadata[fqn]
KeyError: 'pos_embedding'
~~~

T5:
~~~
Traceback (most recent call last): (RANK 3)
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/utils.py", line 140, in reduce_scatter
    local_data = map_fun()
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/state_dict_loader.py", line 89, in local_step
    local_plan = planner.create_local_plan()
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/default_planner.py", line 85, in create_local_plan
    return create_default_local_load_plan(self.state_dict, self.metadata)
  File "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/distributed/_shard/checkpoint/default_planner.py", line 131, in create_default_local_load_plan
    md = metadata.state_dict_metadata[fqn]
KeyError: 'shared.weight'
~~~

I'm unclear if we need to make changes to properly load, but this code used to work fine two months ago.  Saving works nicely as before, just the loading is erring.

here's the code used to load:
~~~
def load_distributed_model_checkpoint(model, rank, cfg):

    if cfg.checkpoint_type == StateDictType.LOCAL_STATE_DICT:
        print(f"loading distributed checkpoint, rank {rank}...")
        folder_name = (
            cfg.dist_checkpoint_root_folder
            + "/"
            + cfg.dist_checkpoint_folder
            + "-"
            + cfg.model_name
        )

        checkdir = Path.cwd() / folder_name

        if not checkdir.exists():
            if rank == 0:
                print(f"No checkpoint directory found...skipping")
            return

        if rank == 0:
            load_timer = Timer()
            load_timer.start()

        reader = FileSystemReader(checkdir)

        with FSDP.state_dict_type(
            model,
            StateDictType.LOCAL_STATE_DICT,
        ):
            state_dict = model.state_dict()
            load_state_dict(state_dict, reader)
            if rank == 0:
                load_timer.interval(name="load of state dict ")
            model.load_state_dict(state_dict)

        print(f"--> local state loaded on rank {rank}")
        if rank == 0:
            load_timer.stop()

        return
~~~





### Versions

Collecting environment information...
PyTorch version: 1.13.0.dev20220928+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.3
Libc version: glibc-2.31

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-1020-aws-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.6.124
GPU models and configuration: 
GPU 0: NVIDIA A10G
GPU 1: NVIDIA A10G
GPU 2: NVIDIA A10G
GPU 3: NVIDIA A10G

Nvidia driver version: 510.73.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.3
[pip3] torch==1.13.0.dev20220928+cu116
[pip3] torch-model-archiver==0.5.3b20220226
[pip3] torch-workflow-archiver==0.2.4b20220513
[pip3] torchserve==0.6.0b20220513
[pip3] torchtext==0.13.1
[pip3] torchvision==0.14.0.dev20220928+cu116
[pip3] vit-pytorch==0.35.8
[conda] blas                      2.116                       mkl    conda-forge
[conda] blas-devel                3.9.0            16_linux64_mkl    conda-forge
[conda] captum                    0.5.0                         0    pytorch
[conda] cudatoolkit               11.6.0              hecad31d_10    conda-forge
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            16_linux64_mkl    conda-forge
[conda] magma-cuda116             2.6.1                         1    pytorch
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.23.3           py39hba7629e_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.13.0.dev20220928+cu116          pypi_0    pypi
[conda] torch-model-archiver      0.5.3                    py39_0    pytorch
[conda] torch-workflow-archiver   0.2.4                    py39_0    pytorch
[conda] torchserve                0.6.0                    py39_0    pytorch
[conda] torchtext                 0.13.1                     py39    pytorch
[conda] torchvision               0.14.0.dev20220928+cu116          pypi_0    pypi

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