save/load deepspeed checkpoint

[Jiaxin-Wen]

I've trained a T5 model with deepspeed stage2 and pytorch-lightning have automatically saved the checkpoints as usual.
However, when I try to load the checkpoints, I got the following error

import torch
path = 'xxxx/mp_rank_00_model_states.pt'
torch.load(path, map_location='cpu')

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/torch/serialization.py", line 594, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/torch/serialization.py", line 853, in _load
    result = unpickler.load()

[rohitgr7]

maybe try this:
https://pytorch-lightning.readthedocs.io/en/latest/advanced/advanced_gpu.html#collating-single-file-checkpoint-for-deepspeed-zero-stage-3

I think it applies to stage 2 as well.

[Jiaxin-Wen]

seems that it also calls torch.load

Detected checkpoint of type zero stage 2, world_size: 8
Traceback (most recent call last):
  File "tools/convert.py", line 12, in <module>
    convert_zero_checkpoint_to_fp32_state_dict(save_path, output_path) 
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/pytorch_lightning/utilities/deepspeed.py", line 81, in convert_zero_checkpoint_to_fp32_state_dict
    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/deepspeed/utils/zero_to_fp32.py", line 377, in get_fp32_state_dict_from_zero_checkpoint
    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/deepspeed/utils/zero_to_fp32.py", line 142, in _get_fp32_state_dict_from_zero_checkpoint
    buffers = parse_model_state(model_file)
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/deepspeed/utils/zero_to_fp32.py", line 56, in parse_model_state
    state_dict = torch.load(file, map_location=device)
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/torch/serialization.py", line 594, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/wenjiaxin/anaconda3/envs/aug/lib/python3.8/site-packages/torch/serialization.py", line 853, in _load
    result = unpickler.load()
ModuleNotFoundError: No module named 'config'

[ShaneTian]

maybe try this: https://pytorch-lightning.readthedocs.io/en/latest/advanced/advanced_gpu.html#collating-single-file-checkpoint-for-deepspeed-zero-stage-3

I think it applies to stage 2 as well.

This page does not exist yet.

[Jiaxin-Wen]

see #10964 for solutions.