zero_to_fp32.py still imports a wrong model after fix

[xycforgithub]

Same as #1165; the issue is closed now so I'm opening a new one.
It looks like the zero_to_fp32.py still cannot load the correct weights under multi-gpu and ZERO-2 setting.

Loading 'mp_rank_00_model_states.pt':

states['module']['deberta.encoder.layer.0.output.dense.weight']
tensor([[-0.0211, 0.0068, 0.0206, ..., 0.0057, 0.0316, 0.0256],
[ 0.0273, 0.0141, 0.0118, ..., -0.0122, 0.0054, 0.0010],
[ 0.0479, -0.0237, -0.0604, ..., -0.0340, -0.0183, 0.0691],
...,
[ 0.0270, -0.0231, 0.0218, ..., 0.0563, 0.0641, -0.0094],
[-0.0563, -0.0837, -0.0427, ..., 0.0242, -0.0132, -0.0512],
[-0.0012, 0.0064, 0.0465, ..., 0.0219, 0.0259, -0.0281]],
device='cuda:0', dtype=torch.float16)

Loading the exported weights using load_state_dict_from_zero_checkpoint:
(Pdb) self.deberta.encoder.layer[0].output.dense.weight
Parameter containing:
tensor([[ 0.0207, -0.0448, 0.0022, ..., 0.0406, -0.0338, -0.0174],
[-0.0577, -0.0648, 0.0404, ..., 0.0108, -0.0167, -0.0100],
[ 0.0548, 0.0063, 0.0024, ..., 0.0311, 0.0249, 0.0167],
...,
[-0.0081, 0.0194, -0.0266, ..., -0.0269, -0.0002, 0.0257],
[ 0.0202, -0.0002, 0.0831, ..., -0.0008, -0.0094, 0.0258],
[-0.0320, 0.0529, -0.0259, ..., 0.0117, -0.0292, -0.0064]],

I'm using Deepspeed 0.4.5. I can confirm the problem happens on 2gpus, and not on 1gpu.

[stas00]

@tjruwase, any chance you guys can take over the maintenance of zero_to_fp32, as I can't keep up with the changes in the core. It's the best that any changes you do get synced with the weights re-consolidation code in tandem.

I created the initial tests here:
https://github.com/microsoft/DeepSpeed/blob/10b48405ab0e69a1caf455f445526c4b39b0dbb8/tests/unit/test_zero.py#L135

So it's probably extending those to include the new variations that Deepspeed is extended to support.

As I suggested ealier perhaps the reconstruction code should be moved along side the partitioning code for each stage, so that they are in the same file and thus are easier to maintain in sync.

So here is how the model gets partitioned, and here is how it gets unpartitioned, one next to each other.

Thank you.

[tjruwase]

@jeffra, FYI

[tjruwase]

@xycforgithub, apologies for the delayed response. Can you please check if this is still an issue as some checkpointing related fixes were recent pushed? If it is still an issue can you please share more details to help repro? Thanks for your patience.

[xycforgithub]

@tjruwase , I can verify the problem is still there in Deepspeed 0.5.2. I saved a model checkpoint in the middle of training. I then used load_state_dict_from_zero_checkpoint(model, ds_dir, tag=str(self.config.start_model_step)) to load the checkpoint (it is a DeBERTa model), and I get
(Pdb) model.deberta.encoder.layer[0].output.dense.weight
Parameter containing:
tensor([[ 0.0207, -0.0448, 0.0022, ..., 0.0406, -0.0338, -0.0174],
[-0.0577, -0.0648, 0.0404, ..., 0.0108, -0.0167, -0.0100],
[ 0.0548, 0.0063, 0.0024, ..., 0.0311, 0.0249, 0.0167],
...,
[-0.0081, 0.0194, -0.0266, ..., -0.0269, -0.0002, 0.0257],
[ 0.0202, -0.0002, 0.0831, ..., -0.0008, -0.0094, 0.0258],
[-0.0320, 0.0529, -0.0259, ..., 0.0117, -0.0292, -0.0064]],
requires_grad=True)

But if I use state_dict = torch.load('mp_rank_00_model_states.pt', map_location=torch.device('cpu')); model.load_state_dict(state_dict['module']), I get
(Pdb) model.deberta.encoder.layer[0].output.dense.weight
Parameter containing:
tensor([[-0.0211, 0.0068, 0.0206, ..., 0.0057, 0.0316, 0.0256],
[ 0.0273, 0.0141, 0.0118, ..., -0.0122, 0.0054, 0.0010],
[ 0.0479, -0.0237, -0.0604, ..., -0.0340, -0.0183, 0.0691],
...,
[ 0.0270, -0.0231, 0.0218, ..., 0.0563, 0.0641, -0.0094],
[-0.0563, -0.0837, -0.0427, ..., 0.0242, -0.0132, -0.0512],
[-0.0012, 0.0064, 0.0465, ..., 0.0219, 0.0259, -0.0281]],
requires_grad=True)

I can verify the latter is correct (same as in the huggingface checkpoint file).

[skpig]

@xycforgithub I have the same issue about loading checkpoint under multi-gpu and ZERO-2 setting. Do you have any methods to load it correctly ?

[stas00]

I'm curious if you have tried with master, I was mainly doing fixes on zero3, but who knows perhaps something else got fixed on the way.

[xycforgithub]

@xycforgithub I have the same issue about loading checkpoint under multi-gpu and ZERO-2 setting. Do you have any methods to load it correctly ?

The only method under zero2 I find is to load the fp16 weights in mp_rank_00_model_states.pt.

[skpig]

@stas00 I just updated my deepspeed to version 0.5.3 and tested the load_state_dict_from_zero_checkpoint(), but something went wrong.

Traceback (most recent call last):
  File "zero_to_fp32.py", line 375, in <module>
    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)
  File "zero_to_fp32.py", line 313, in convert_zero_checkpoint_to_fp32_state_dict
    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
  File "zero_to_fp32.py", line 299, in get_fp32_state_dict_from_zero_checkpoint
    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
  File "zero_to_fp32.py", line 244, in _get_fp32_state_dict_from_zero_checkpoint
    f"consumed {offset} numels out of {avail_numel} - something is wrong")
ValueError: consumed 255381772 numels out of 255381776 - something is wrong

Here is my code.

1. First I saved the model under ZeRO2 and 4 gpus settings . I used my own model class MyBart (subclass of transformers.BartPretrainedModel).

    model = MyBart.from_pretrained("facebook/bart-base")
    param_optimizer = list(model.named_parameters())
    optimizer_grouped_parameters = [{
        'params':
            [p for n, p in param_optimizer],
        'weight_decay':
            0.01
    }]

    model_engine, optimizer, _, scheduler = deepspeed.initialize(args=args, model=model,
                                                                 model_parameters=optimizer_grouped_parameters)
    model_engine.save_checkpoint("Model/tmp", tag="ckpt")
    path, state = model_engine.load_checkpoint("Model/tmp", tag="ckpt")

2. Than I load the checkpoint without deepspeed, and the Error Message above happened.

    model = MyBart.from_pretrained("facebook/bart-base")
    print(model.lm_head.weight)
    print(model.si_lm_head.weight)
    print(model.base_model.encoder.layers[0].fc1.weight)
    model = load_state_dict_from_zero_checkpoint(model, "Model/tmp")
    print(model.lm_head.weight)
    print(model.si_lm_head.weight)
    print(model.base_model.encoder.layers[0].fc1.weight)

[stas00]

please re-test with deepspeed@master - a few bug fixes were merged this week. Please report back if the problem has been resolved or not - thank you!

of course you will need to either re-train or you can just copy over the improved script from the git source :)

[stas00]

ValueError: consumed 255381772 numels out of 255381776 - something is wrong

The diff is 4 bytes, so this looks like the new padding introduced by Z2 recently. this code is a suspect:

https://github.com/microsoft/DeepSpeed/blob/c1829c45851eb1988684c9a2f4fd054ebd446b95/deepspeed/utils/zero_to_fp32.py#L226-L234

Perhaps I got the new padding wrong? Does the model work if you remove the assert?

Also what numbers do you get before they are aligned to 4? Could be some edge case in the logic.

i.e. if you print out f"consumed {offset} numels out of {avail_numel}" before this line:

https://github.com/microsoft/DeepSpeed/blob/c1829c45851eb1988684c9a2f4fd054ebd446b95/deepspeed/utils/zero_to_fp32.py#L226

[stas00]

Here is the PR where I was trying to fix it #1372
and here is the test that was reproducing the 2 bytes off problem:
https://github.com/microsoft/DeepSpeed/blob/c1829c45851eb1988684c9a2f4fd054ebd446b95/tests/unit/test_zero.py#L164

[stas00]

Finally please run the script with -d so that we can see which param shapes you got (or enable debug=1 hardcoded in the script) if you're using the API.

[skpig]

@stas00 Thanks for your reply. I'm so sorry for my misunderstanding about master & release version. I'm new to GitHub, please forgive me.

please re-test with deepspeed@master - a few bug fixes were merged this week. Please report back if the problem has been resolved or not - thank you!

I retested the code with master branch. The Value Error (consumed 255381772 numels out of 255381776 - something is wrong) still happened. Then I do the testing following your instruction.

1. I removed the assert, and the loading process went well. I also randomly sampled some parameters to compare with original model. Only found 2 different values with error less than 0.0001.
2. Then I added the code in the script as you told, and output consumed 255381771 numels out of 255381775

i.e. if you print out f"consumed {offset} numels out of {avail_numel}" before this line:

3. I ran the script with -d, and here is the output ( I skip the shapes of my model params because it is too long)

Found buffers: []
fp32_flat_groups[i].shape=torch.Size([63845444])
fp32_flat_groups[i].shape=torch.Size([63845444])
fp32_flat_groups[i].shape=torch.Size([63845444])
fp32_flat_groups[i].shape=torch.Size([63845443])
Have 255381775 numels to process.
Need 255381771 numels in 265 params.
added 0 buffers

[stas00]

@stas00 Thanks for your reply. I'm so sorry for my misunderstanding about master & release version. I'm new to GitHub, please forgive me.

Well, it's I who should ask for forgiveness since I made the assumption are you an old time user without checking in with you first :)

I'm glad you have figured it out on your own, @skpig

please re-test with deepspeed@master - a few bug fixes were merged this week. Please report back if the problem has been resolved or not - thank you!

I retested the code with master branch. The Value Error (consumed 255381772 numels out of 255381776 - something is wrong) still happened. Then I do the testing following your instruction.

1. I removed the assert, and the loading process went well. I also randomly sampled some parameters to compare with original model. Only found 2 different values with error less than 0.0001.

2. Then I added the code in the script as you told, and output `consumed 255381771 numels out of 255381775 `

i.e. if you print out f"consumed {offset} numels out of {avail_numel}" before this line:

1. I ran the script with `-d`, and here is the output ( I skip the shapes of my model params because it is too long)

Found buffers: []
fp32_flat_groups[i].shape=torch.Size([63845444])
fp32_flat_groups[i].shape=torch.Size([63845444])
fp32_flat_groups[i].shape=torch.Size([63845444])
fp32_flat_groups[i].shape=torch.Size([63845443])
Have 255381775 numels to process.
Need 255381771 numels in 265 params.
added 0 buffers

Thank you for running this check, and I actually wanted to see the param shapes. You can always attach any long outputs to your comment as a file, so that it doesn't make the comment too long.

But even from the report above it's clear that the issue is of a padding and it's strange that your saved model is uneven.

So the next step is to be able to reproduce this for me. Could you try to dramatically reduce this model to say 2 layers and see if you could still reproduce the issue?

And ideally we want a public model since I don't know what is inside your MyBart - do you get the same issue if you use Bart instead.

To clarify this is clearly some edge case, so we need to reproduce and understand what's on the edge and then we should be able to fix if easily.

[stas00]

p.s. there is another trick to including a huge log file in the comment and making it collapsible, see: https://gist.github.com/pierrejoubert73/902cc94d79424356a8d20be2b382e1ab#gistcomment-3899835