Repository: microsoft/DeepSpeed

Issue:
[BUG] Size of saved model checkpoint becomes much larger after deepspeed.initialize when using ZeRO-2
**Describe the bug**
Originally reported [here](https://github.com/huggingface/transformers/issues/22822). @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**
```python
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:

```bash
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:
```python
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




Discussion:
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: https://github.com/microsoft/DeepSpeed/pull/1324 - but it could be something else as well.
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'):`
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.

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.


Resolving patch:
diff --git a/deepspeed/checkpoint/utils.py b/deepspeed/checkpoint/utils.py
--- a/deepspeed/checkpoint/utils.py
+++ b/deepspeed/checkpoint/utils.py
@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import os
+import torch
 from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX)
 
 
@@ -29,3 +30,33 @@ def get_layer_ckpt_name_for_rank(base_folder, layer_id, tp_rank):
     ckpt_file = f'{layer_id}-model_{tp_rank:02d}{MODEL_FILE_SUFFIX}'
     ckpt_path = os.path.join(base_folder, ckpt_file)
     return ckpt_path
+
+
+# We pass cloned tensors to torch.save() to avoid checkpoint bloat that occurs when torch.save()
+# saves the underlying storage rather than the slice of the storage corresponding to individual tensors.
+# This is a problem in DeepSpeed because we often allocate tensors using slices of large flattened buffers.
+# Tensor cloning helps to avoid this problem because the storage of cloned tensors are closer to the true size.
+# It is expected that the garbage collector will reclaim the cloned tensor storage to avoid memory bloat.
+# See https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing
+def clone_tensors_for_torch_save(item, device=torch.device('cpu')):
+    """
+    Returns a copy of ``item`` with all enclosed tensors replaced by clones on a specified device.
+    Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.
+
+    Parameters:
+        - ``item``: tensor to clone or (possibly nested) container of tensors to clone.
+        - ``device``: target device (defaults to 'cpu')
+
+    Returns:
+        - copy of ``item`` with cloned tensors on target device
+    """
+    if torch.is_tensor(item):
+        return item.detach().clone().to(device)
+    elif isinstance(item, list):
+        return [clone_tensors_for_torch_save(v, device) for v in item]
+    elif isinstance(item, tuple):
+        return tuple([clone_tensors_for_torch_save(v, device) for v in item])
+    elif isinstance(item, dict):
+        return type(item)({k: clone_tensors_for_torch_save(v, device) for k, v in item.items()})
+    else:
+        return item


Test patch:
diff --git a/tests/unit/checkpoint/test_zero_optimizer.py b/tests/unit/checkpoint/test_zero_optimizer.py
--- a/tests/unit/checkpoint/test_zero_optimizer.py
+++ b/tests/unit/checkpoint/test_zero_optimizer.py
@@ -5,6 +5,8 @@
 
 import deepspeed
 from deepspeed.ops.op_builder import CPUAdamBuilder
+from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
+from deepspeed.accelerator import get_accelerator
 
 from unit.common import DistributedTest, DistributedFixture
 from unit.simple_model import *
@@ -469,3 +471,37 @@ def test_load_module_only(self, tmpdir, zero_stage):
             models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
 
         checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
+
+
+class TestSaveTensorClone(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize('zero_stage', [1, 2])
+    @pytest.mark.parametrize('use_cpu_device', [True, False])
+    def test_save_tensor_clone(self, tmpdir, zero_stage, use_cpu_device):
+
+        ds_config = {
+            "optimizer": {
+                "type": "AdamW",
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            },
+            "train_batch_size": 1,
+            "train_micro_batch_size_per_gpu": 1
+        }
+        hidden_dim = 1024
+        model = SimpleModel(hidden_dim, nlayers=4).half()
+        ref_model_state_dict = model.state_dict()
+
+        ds_engine, _, _, _ = deepspeed.initialize(model=model, config_params=ds_config)
+        clone_device = torch.device('cpu') if use_cpu_device else get_accelerator().current_device()
+        clone_state_dict = clone_tensors_for_torch_save(ds_engine.module.state_dict())
+        compare_state_dicts(ref_model_state_dict, clone_state_dict)
+
+        ref_ckpt_file = os.path.join(tmpdir, 'ref_ckpt.pt')
+        torch.save(ref_model_state_dict, ref_ckpt_file)
+        clone_ckpt_file = os.path.join(tmpdir, 'clone_ckpt.pt')
+        torch.save(clone_state_dict, clone_ckpt_file)
+
+        compare_state_dicts(torch.load(ref_ckpt_file), torch.load(clone_ckpt_file))