Unable to load deepspeed checkpoint

[lanx7]

Hi.
I found an error while testing Deep Speed Stage 3. Ask for help.

I've saved the checkpoint file and collated it into a single file according to the guide.

python zero_to_fp32.py model_chp/last.ckpt/ model_chp/last.ckpt/final.pt

When I loaded the file (final.pt) with load_from_checkpoint() function, the following error occurred.

-------------------------------------------------------------
Traceback (most recent call last):
  File "deepspeed_test.py", line 76, in <module>
    model = TestModel.load_from_checkpoint(args.model)
  File "/opt/conda/envs/gpt3/lib/python3.8/site-packages/pytorch_lightning/core/saving.py", line 156, in load_from_checkpoint
    model = cls._load_model_state(checkpoint, strict=strict, **kwargs)
  File "/opt/conda/envs/gpt3/lib/python3.8/site-packages/pytorch_lightning/core/saving.py", line 198, in _load_model_state
    model = cls(**_cls_kwargs)
  File "/root/gpt/test/trinity_test.py", line 28, in __init__
    self.save_hyperparameters()
  File "/opt/conda/envs/gpt3/lib/python3.8/site-packages/pytorch_lightning/core/mixins/hparams_mixin.py", line 105, in save_hyperparameters
    save_hyperparameters(self, *args, ignore=ignore, frame=frame)
  File "/opt/conda/envs/gpt3/lib/python3.8/site-packages/pytorch_lightning/utilities/parsing.py", line 223, in save_hyperparameters
    assert init_args, "failed to inspect the obj init"
AssertionError: failed to inspect the obj init
------------------------------------------------------------

Please refer to my code below.

Model Initialization Part:
------------------------------------------
class TestModel(LightningModule):
    def __init__(self, **kwargs):
        super(TestModel, self).__init__()
        self.save_hyperparameters()
        ....
----------------------------------------

Training and Checkpointing: 
--------------------------------------
checkpoint_callback = ModelCheckpoint(
        dirpath=args.dir,
        filename=args.name + '_' + '{epoch:02d}-{train_loss:.2f}-{val_loss:.2f}',
        verbose=True,
        save_last=True,
        monitor='val_loss',
        mode='min',
  )
    dict_args = vars(args)
    model = TestModel(**dict_args)
    model.train()
    trainer = Trainer.from_argparse_args(
        args, callbacks=[checkpoint_callback], gradient_clip_val=1.0)
   trainer.fit(model)
------------------------------------------

Loading Checkpoint for Testing:
------------------------------------------------------
model = TestModel.load_from_checkpoint("model_chp/last.ckpt/final.pt")
-----------------------------------------------------

Library Version:
pytorch-ignite==0.4.8
pytorch-lightning==1.5.10
torch==1.10.2+cu113
torchaudio==0.10.2+cu113
torchmetrics==0.7.2
torchtext==0.11.2
torchvision==0.11.3+cu113

[rohitgr7]

hey @lanx7 !

is it possible to share a reproducible issue using https://colab.research.google.com/github/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/bug_report/bug_report_model.ipynb

also, can you try with master too? this might have been fixed recently and next release is soon

[lanx7]

Hi @rohitgr7

I'm not familiar with Colab, but I tried my best to write it.
Please refer to https://colab.research.google.com/drive/1vflOpwQF-5WQiuy9ypskdxuZid3omeL9

Let me know if you need anything more to clarify the issue.

Also, I'll try to try with master too.

[rohitgr7]

can you make it public? thank you

[lanx7]

I've changed it to the public.

Thank you.

[lanx7]

I tested it with master (i.e 1.6.0-dev).
The error I mentioned above doesn't appear. But, I am experiencing another issue.

Traceback (most recent call last):
  File "deepspeed_test.py", line 77, in <module>
    model = TestModel.load_from_checkpoint(args.model)
  File "/opt/conda/envs/pl16/lib/python3.8/site-packages/pytorch_lightning-1.6.0.dev0-py3.8.egg/pytorch_lightning/core/saving.py", line 161, in load_from_checkpoint
    model = cls._load_model_state(checkpoint, strict=strict, **kwargs)
  File "/opt/conda/envs/pl16/lib/python3.8/site-packages/pytorch_lightning-1.6.0.dev0-py3.8.egg/pytorch_lightning/core/saving.py", line 209, in _load_model_state
    keys = model.load_state_dict(checkpoint["state_dict"], strict=strict)
KeyError: 'state_dict'

[lanx7]

Hi @rohitgr7 and other colloborator.

After some observation, I found out that the checkpoint file produced by 'zero_to_fp32.py' or convert_zero_checkpoint_to_fp32_state_dict() is a pytorch standard checkpoint file rather than the one of pytorch-lightning model.

I have succeeded in loading the model using the standardized pytorch API.

		model = TestModel(**dict_args)
		#model.load_state_dict(torch.load(args.model), strict=False)	
		checkpoint = torch.load(args.model, map_location=torch.device('cuda'))
		for key in list(checkpoint.keys()):
			if 'module.' in key:
				checkpoint[key.replace('module.', '')] = checkpoint[key]
				print (key)
				del checkpoint[key]
		model.load_state_dict(checkpoint, strict=False)

But, I still have two issues.

1. All keys in the state_dict have the prefix 'module' attached, so it has been removed.
2. State 'lm_head.weight' is missing, so I set 'load_state_dict' with 'strict=False' to avoid the issue

The outputs of model are similar to ones when I did not use Deepspeed, but it is still suspicious due to some issues.

Can you confirm if this method is correct?
And, even if this method is correct, it would be better to have 'load_from_checkpoint()' supporting the single collated checkpoint from Deepspeed, which are produced by 'zero_to_fp32.py' or convert_zero_checkpoint_to_fp32_state_dict().

[ShaneTian]

@lanx7 I met the same issue.
The default zero_to_fp32.py also does not work for me.
So I use the following code to convert checkpoint:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Convert zero checkpoint to FP32 checkpoint.
"""

import argparse
import os

import torch


def setup_args():
    """Setup arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--zero_ckpt", type=str, required=True,
                        help="The path to the DeepSpeed ZERO checkpoint.")
    parser.add_argument("--fp32_ckpt", type=str, required=True,
                        help="The path to the FP32 checkpoint.")
    parser.add_argument("--delete_zero_ckpt", action="store_true", default=False,
                        help="Whether to delete the DeepSpeed ZERO checkpoint.")
    args = parser.parse_args()
    return args


def main(args):
    """Main function."""
    zero_params = torch.load(args.zero_ckpt)
    fp32_state_dict = {"state_dict": {}}
    for k, v in zero_params.items():
        if k == "module":
            for name, tensor in zero_params["module"].items():
                name = name.replace("module.", "")
                fp32_state_dict["state_dict"][name] = tensor
        else:
            fp32_state_dict[k] = v
    torch.save(fp32_state_dict, args.fp32_ckpt)
    if args.delete_zero_ckpt:
        os.remove(args.zero_ckpt)


if __name__ == "__main__":
    args = setup_args()
    main(args)

But the best way to solve this would be pytorch-lightning. @rohitgr7 Thank you.