misc problems with saving the checkpoint

[stas00]

setup - 2 local gpus: with the same default config I use everywhere

I have several issues with saving the model:

1. the original fp32 model gets saved as fp16 model - how do I get back the fp32 model? The user may not proceed with deepspeed and will want to share the model, so it needs to be back in fp32. Perhaps there should be a save_fp32_model method? or an option in the checkpoint?

How does it even work on resume if fp32 weights aren't getting saved?

2. I use deepspeed.save_checkpoint(output_dir)

the first checkpoint gets saved it seems if I look at the filesystem, and then it hangs

  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-div-by-zero/deepspeed/runtime/engine.py", line 1456 in _checkpoint_tag_validation
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-div-by-zero/deepspeed/runtime/engine.py", line 1489 in save_checkpoint
  File "/mnt/nvme1/code/huggingface/transformers-ds-save-model/src/transformers/trainer.py", line 1212 in _save_checkpoint
  File "/mnt/nvme1/code/huggingface/transformers-ds-save-model/src/transformers/trainer.py", line 1185 in _maybe_log_save_evaluate
  File "/mnt/nvme1/code/huggingface/transformers-ds-save-model/src/transformers/trainer.py", line 1094 in train

(Thanks to @jeffra for the tip on py-spy!)

I obviously tried to disable the check and went on to discover the undocumented config option:

    "checkpoint": {
        "tag_validation": "ignore"
    },

which I reverse engineered. perhaps it could be documented? (warn/ignore/fail are the 3 options)

But I also didn't use any tags...

I also tried to save only from the rank 0 process to no avail.

so when I add this, now it gets stuck in:

  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 2427 in barrier
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-div-by-zero/deepspeed/runtime/engine.py", line 1524 in _create_zero_checkpoint_files
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-div-by-zero/deepspeed/runtime/engine.py", line 1496 in save_checkpoint
  File "/mnt/nvme1/code/huggingface/transformers-ds-save-model/src/transformers/trainer.py", line 1212 in _save_checkpoint
  File "/mnt/nvme1/code/huggingface/transformers-ds-save-model/src/transformers/trainer.py", line 1185 in _maybe_log_save_evaluate
  File "/mnt/nvme1/code/huggingface/transformers-ds-save-model/src/transformers/trainer.py", line 1094 in train

if I remove save_checkpoint the program doesn't hang and completes just fine.

If I don't save intermediary checkpoints and save only when the training is finished it hangs too (i.e. on the first call)

[tjruwase]

@stas00, if this is ZeRO, then all ranks needs to call save_checkpoint to ensure the distributed optimizer states are correctly saved. You are hitting this barrier that ensures that ZeRO checkpoint files are correctly created.

[stas00]

That was it for the hanging - thank you very much!

And how do I recover the fp32 model back?

[tjruwase]

DeepSpeed's save_checkpoint and load_checkpoint should save and restore fp32 weights in the optimizer state correctly, ideally. The fp32 weights, momentum, and variance are stored in the per-rank zero checkpoint files. However, we currently have a bug in the save logic. I am not sure whether this bug is the issue you are seeing.

Can you please add the following to your config to workaround the bug?
zero_optimization: {'elastic_checkpoint':=false}

[stas00]

No change, I get the outputs of the same size w/ or w/o that setting you suggested I use (i.e. tried true and false):

-rw-rw-r-- 1 stas stas 347M Feb 26 11:57 zero_pp_rank_1_mp_rank_00optim_states.pt
-rw-rw-r-- 1 stas stas 347M Feb 26 11:57 zero_pp_rank_0_mp_rank_00optim_states.pt
-rw-rw-r-- 1 stas stas 116M Feb 26 11:57 mp_rank_00_model_states.pt

The model should be twice the size

[tjruwase]

Can you explain further on what you mean by the model should be twice the size.

The proportion of the file sizes seem reasonable:
mp_rank_00* stores fp16 weights,
zero_pp_rank_* store fp32 weights, momentum, and variance for ranks 0 and 1 correspondingly.

The combined zero_pp_rank* file size (4Bytes+4Bytes+4Bytes) should be roughly 6X of mp_rank_00* file (2Bytes). Right now the ratio is 5.98.

What is the model size?

[stas00]

The original is 231M

So what you're saying is that the fp32 version is inside zero_pp_rank_* files. That part I understood.

So the question is the same, how do I recover the original fp32 version after the finetuning?

What you currently provide is a way to resume training with DeepSpeed, what we need is a way for a user to take the finetuned fp32 model and use it in any environment. I hope that clarifies the need better.

For example, Deepspeed doesn't do inference at the moment, so we need the full fp32 model to do inference with.

[tjruwase]

Got it. Thanks for the clarification. Unfortunately, there is no easy way to extract the fp32 weights inside the zero_pp_rank_* files. This makes it a TODO. Could please open an issue?

[stas00]

Done: #800

[stas00]

And proposed a documentation addition to the first part of your answer here #801

[stas00]

This has been partly resolved, partly moved to a new issue #800