[BUG] save_checkpoint() missing some params

Describe the bug
I'm using deepspeed to train GNN model from pgn (https://github.com/pyg-team/pytorch_geometric). But after I save the model with engine.save_checkpoint(), I couldn't load the checkpoint because of missing parameters.

To Reproduce
I provide a code snippet.
I launch save() function with deepspeed --include="localhost:2,3" train.py --deepspeed --deepspeed_config config/ds_config.json (I've tried no-ZeRO, ZeRO1 and ZeRO2 settings, all failed).

def save():
    from torch_geometric.nn import GATConv
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", default=0, type=int,
                        help="local_rank for distributed training on gpus")
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    os.environ['MASTER_ADDR'] = 'localhost'  #
    os.environ['MASTER_PORT'] = '12345'  #
    args.local_rank = int(os.environ['LOCAL_RANK'])

    model = GATConv(4, 2, heads=2, edge_dim=3)
    engine, _, _, _ = deepspeed.initialize(args=args, model=model, model_parameters=model.parameters())
    engine.save_checkpoint('Model/tmp_model', tag='checkpoint-0')

And load checkpoint with python launcher.

def load():
    from torch_geometric.nn import GATConv
    model = GATConv(4, 2, heads=2, edge_dim=3)
    model = load_state_dict_from_zero_checkpoint(model, "Model/tmp_model", 'checkpoint-0')

output

[2022-02-24 02:43:46,824] [INFO] [zero_to_fp32.py:425:load_state_dict_from_zero_checkpoint] Extracting fp32 weights
Processing zero checkpoint 'Model/tmp_model/checkpoint-0'
Detected checkpoint of type zero stage 2, world_size: 2
Reconstructed fp32 state dict with 6 params 44 elements
[2022-02-24 02:43:46,826] [INFO] [zero_to_fp32.py:428:load_state_dict_from_zero_checkpoint] Overwriting model with fp32 weights
Traceback (most recent call last):
  File "/home/huangbz/_project2/train.py", line 344, in <module>
    model = load_state_dict_from_zero_checkpoint(model, "Model/tmp_model", 'checkpoint-0')
  File "/home/huangbz/.conda/envs/Graph/lib/python3.6/site-packages/deepspeed/utils/zero_to_fp32.py", line 430, in load_state_dict_from_zero_checkpoint
    model.load_state_dict(state_dict, strict=False)
  File "/home/huangbz/.conda/envs/Graph/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1468, in load_state_dict
    load(self)
  File "/home/huangbz/.conda/envs/Graph/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1466, in load
    load(child, prefix + name + '.')
  File "/home/huangbz/.conda/envs/Graph/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1463, in load
    state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
  File "/home/huangbz/.conda/envs/Graph/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1372, in _load_from_state_dict
    hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
  File "/home/huangbz/.conda/envs/Graph/lib/python3.6/site-packages/torch_geometric/nn/dense/linear.py", line 131, in _lazy_load_hook
    weight = state_dict[prefix + 'weight']
KeyError: 'lin_dst.weight'

System info (please complete the following information):

• GPU count and types [e.g. two machines with x8 A100s each]
• Python version: 3.6
• PGN version: 2.0.3

Start with the save_checkpoint() and load_state_dict_from_zero_checkpoint() entry points and reproduce the GATConv example using the stated DeepSpeed and Python versions. Compare the six reconstructed parameters with the GATConv state dict; done means the saved checkpoint reloads without the missing lin_dst.weight error.