[BUG] Size of saved model checkpoint becomes much larger after deepspeed.initialize when using ZeRO-2

[ArvinZhuang]

Describe the bug
Originally reported here. @stas00 @tjruwase

For some models, the size of model checkpoints saved by model.save_prtrained() becomes much larger after calling deepspeed.initialize. See examples below.

To Reproduce

from transformers import AutoModelForCausalLM
import deepspeed

ds_config = {
    "optimizer": {
        "type": "AdamW",
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "allgather_partitions": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": True
    },
    "offload_optimizer": {
        "device": "cpu",
        "pin_memory": True
    },
    "train_batch_size": 1,
    "train_micro_batch_size_per_gpu": 1
}

model = AutoModelForCausalLM.from_pretrained("decapoda-research/llama-7b-hf")
model.save_pretrained("before")
deepspeed_engine, _, _, _ = deepspeed.initialize(model=model, config_params=ds_config)
deepspeed_engine.module.save_pretrained("after")

File sizes:

du -a -h --max-depth=1 before/
512     before/config.json
32K     before/pytorch_model.bin.index.json
9.2G    before/pytorch_model-00001-of-00003.bin
9.3G    before/pytorch_model-00002-of-00003.bin
6.7G    before/pytorch_model-00003-of-00003.bin
512     before/generation_config.json
26G     before/

du -a -h --max-depth=1 after/
512     after/config.json
32K     after/pytorch_model.bin.index.json
26G     after/pytorch_model-00001-of-00003.bin
26G     after/pytorch_model-00002-of-00003.bin
26G     after/pytorch_model-00003-of-00003.bin
512     after/generation_config.json
76G     after/

This issue is not always occurred, for example, gpt2 does not have this problem. But I tested decapoda-research/llama-7b-hf, and decapoda-research/llama-13b-hf have this issue.

This can be fixed by re-clone states before the saving:

state_dict = deepspeed_engine.module.state_dict()
state_dict = type(state_dict)(
    {k: v.clone()
     for k,
     v in state_dict.items()})
deepspeed_engine.module.save_pretrained("after_fixed", state_dict=state_dict)

Expected behavior
The saved model size should be unchanged after deepspeed.initialize

System info (please complete the following information):

• deepspeed: 0.8.3
• transformers version: 4.28.0.dev0
• Platform: Linux-4.18.0-372.32.1.el8_6.x86_64-x86_64-with-glibc2.17
• Python version: 3.8.16
• Huggingface_hub version: 0.13.3
• Safetensors version: not installed
• PyTorch version (GPU?): 1.12.1+cu116 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: yes

[stas00]

You can see the original model weights here: https://huggingface.co/decapoda-research/llama-7b-hf/tree/main
There are a total of about 13GB in float16 (in chunks of 400MB in 33 files) - but the code above loads in fp32, since the user hasn't asked for fp16, so it ends up 26GB total. and suspiciously the sharded in save_pretrained files are also 26GB large but each of the 3 files. Something strange is going on.

Suspecting something similar to this fix: #1324 - but it could be something else as well.

[tjruwase]

Need to expose utility for trainers to clone tensors and undo effect of combining zero flattening and https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing

def clone_tensors_for_torch_save(item, device=torch.device('cpu'):

[stas00]

We discussed that and so the tentative solution is to provide a helper util that will clone the z1 and z2 models to cpu which will undo the bloating. The assumption is that the user should have enough free CPU memory to allocate the model, since they had to have it to load it in the first place.

We don't want to do in place shrinking since it will break the flattened buffers should the training continue.

Once the shrunken state_dict is returned to the user they save that model instead.

Then we integrate this solution into HF Trainer and Accelerate for Z1 and Z2 use cases.

It'll be documented how to do it for users who write their own Trainer.

[stas00]

Now to explain what happened.

The problem happens because torch saves the full datastorage and if multiple tensors share it via view() it'll have to save the full storage. So if you take a 10GB storage and split it into 2 shards, it'll save each storage of 10GB and end up with 20GB final file size.

The problem didn't emerge until save_pretrained started sharding large checkpoints. When a single checkpoint file is saved, the flattened tensors correlate 100% to individual tensors and therefore there was 0 bloat.

Now if save_pretrained splits the checkpoint into 10 chunks, it will now consume 10x more disc space, completely defeating the purpose of sharding the checkpoint.

So this needs to be fixed.

[tjruwase]

Below is demo of a fix using PR #3348 and state_dict argument of save_pretrained().

Code:

from transformers import AutoModelForCausalLM
import deepspeed
from deepspeed.runtime.utils import clone_tensors_for_torch_save
import torch 

ds_config = {
    "optimizer": {
        "type": "AdamW",
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "allgather_partitions": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": True
    },
    "offload_optimizer": {
        "device": "cpu",
        "pin_memory": True
    },
    "train_batch_size": 1,
    "train_micro_batch_size_per_gpu": 1
}

model = AutoModelForCausalLM.from_pretrained("decapoda-research/llama-7b-hf")
model.save_pretrained("before")
deepspeed_engine, _, _, _ = deepspeed.initialize(model=model, config_params=ds_config)
deepspeed_engine.module.save_pretrained("after")

lean_state_dict = clone_tensors_for_torch_save(deepspeed_engine.module.state_dict())
deepspeed_engine.module.save_pretrained("lean_after", state_dict=lean_state_dict)

File sizes:

du -a -h --max-depth=1 before/ after/ lean_after/
4.0K    before/config.json
6.7G    before/pytorch_model-00003-of-00003.bin
28K     before/pytorch_model.bin.index.json
9.2G    before/pytorch_model-00001-of-00003.bin
9.3G    before/pytorch_model-00002-of-00003.bin
4.0K    before/generation_config.json
26G     before/
4.0K    after/config.json
26G     after/pytorch_model-00003-of-00003.bin
28K     after/pytorch_model.bin.index.json
26G     after/pytorch_model-00001-of-00003.bin
26G     after/pytorch_model-00002-of-00003.bin
4.0K    after/generation_config.json
76G     after/
4.0K    lean_after/config.json
6.7G    lean_after/pytorch_model-00003-of-00003.bin
28K     lean_after/pytorch_model.bin.index.json
9.2G    lean_after/pytorch_model-00001-of-00003.bin
9.3G    lean_after/pytorch_model-00002-of-00003.bin
4.0K    lean_after/generation_config.json
26G     lean_after/