fail to load checkpoints after zero3 initialize

[Stick-To]

  File "/home/lxr/workspace/t5_train/esm5.py", line 66, in <module>
    engine.load_module_state_dict(torch.load("weight/mt5-small/pytorch_model.bin"))
  File "/home/lxr/workspace/t5_train/deepspeed/runtime/engine.py", line 2421, in load_module_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for MT5ForConditionalGeneration:
        size mismatch for shared.weight: copying a param with shape torch.Size([250112, 512]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.embed_tokens.weight: copying a param with shape torch.Size([250112, 512]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.block.0.layer.0.SelfAttention.q.weight: copying a param with shape torch.Size([384, 512]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.block.0.layer.0.SelfAttention.k.weight: copying a param with shape torch.Size([384, 512]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.block.0.layer.0.SelfAttention.v.weight: copying a param with shape torch.Size([384, 512]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.block.0.layer.0.SelfAttention.o.weight: copying a param with shape torch.Size([512, 384]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight: copying a param with shape torch.Size([32, 6]) from checkpoint, the shape in current model is torch.Size([0]).
        size mismatch for encoder.block.0.layer.0.layer_norm.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([0]).

[dittops]

@Stick-To I'm also facing this issue. Could you please share how you resolved it?

[tjruwase]

@dittops, please re-open and share repro steps including a stack trace. Thanks!

[Alchemy5]

im also facing this issue, any help would be great!

[iamsile]

I'm also facing this issue too, anyone who can help ?

[Alchemy5]

I actually found a solution to my problem in the Saving and Loading section of this article: https://huggingface.co/docs/accelerate/usage_guides/deepspeed#saving-and-loading

[yix-chen]

@Stick-To I'm also facing this issue. Could you please share how you resolved it?

[Stick-To]

I have not solve it

[Zx55]

One possible reason could be the conflict of multiple initialization between hf deepspeed integration and explicit call "deepspeed.zero.Init()". I solve this following here.

[ZeyuLing]

I met the same problem and solved it by reinitialization. U can deepcopy an original model before u wrap it with deepspeed.initialize. Load your checkpoint to the original model, and run deepspeed.initialize to the loaded one.