prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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
  
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `deepspeed/checkpoint/utils.py`
  Content:
  ```
  1 # Copyright (c) Microsoft Corporation.
  2 # SPDX-License-Identifier: Apache-2.0
  3 
  4 # DeepSpeed Team
  5 
  6 import os
  7 from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX)
  8 
  9 
  10 def get_model_ckpt_name_for_rank(base_folder, mp_rank_str):
  11     ckpt_name = os.path.join(
  12         base_folder,
  13         MODEL_FILE_PREFIX + mp_rank_str + MODEL_FILE_SUFFIX,
  14     )
  15     return ckpt_name
  16 
  17 
  18 def get_zero_ckpt_name_for_rank(base_folder, dp_rank, mp_rank):
  19     zero_prefix = f'{ZERO_FILE_PREFIX}{dp_rank}'
  20     mp_rank_string = f'_{MODEL_FILE_PREFIX}{mp_rank:02d}'
  21     zero_ckpt_name = os.path.join(
  22         base_folder,
  23         zero_prefix + mp_rank_string + OPTIM_FILE_SUFFIX,
  24     )
  25     return zero_ckpt_name
  26 
  27 
  28 def get_layer_ckpt_name_for_rank(base_folder, layer_id, tp_rank):
  29     ckpt_file = f'{layer_id}-model_{tp_rank:02d}{MODEL_FILE_SUFFIX}'
  30     ckpt_path = os.path.join(base_folder, ckpt_file)
  31     return ckpt_path
  32
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
       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
  
verification_info: {"golden_diff": "diff --git a/deepspeed/checkpoint/utils.py b/deepspeed/checkpoint/utils.py\n--- a/deepspeed/checkpoint/utils.py\n+++ b/deepspeed/checkpoint/utils.py\n@@ -4,6 +4,7 @@\n # DeepSpeed Team\n \n import os\n+import torch\n from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX)\n \n \n@@ -29,3 +30,33 @@\n     ckpt_file = f'{layer_id}-model_{tp_rank:02d}{MODEL_FILE_SUFFIX}'\n     ckpt_path = os.path.join(base_folder, ckpt_file)\n     return ckpt_path\n+\n+\n+# We pass cloned tensors to torch.save() to avoid checkpoint bloat that occurs when torch.save()\n+# saves the underlying storage rather than the slice of the storage corresponding to individual tensors.\n+# This is a problem in DeepSpeed because we often allocate tensors using slices of large flattened buffers.\n+# Tensor cloning helps to avoid this problem because the storage of cloned tensors are closer to the true size.\n+# It is expected that the garbage collector will reclaim the cloned tensor storage to avoid memory bloat.\n+# See https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing\n+def clone_tensors_for_torch_save(item, device=torch.device('cpu')):\n+    \"\"\"\n+    Returns a copy of ``item`` with all enclosed tensors replaced by clones on a specified device.\n+    Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.\n+\n+    Parameters:\n+        - ``item``: tensor to clone or (possibly nested) container of tensors to clone.\n+        - ``device``: target device (defaults to 'cpu')\n+\n+    Returns:\n+        - copy of ``item`` with cloned tensors on target device\n+    \"\"\"\n+    if torch.is_tensor(item):\n+        return item.detach().clone().to(device)\n+    elif isinstance(item, list):\n+        return [clone_tensors_for_torch_save(v, device) for v in item]\n+    elif isinstance(item, tuple):\n+        return tuple([clone_tensors_for_torch_save(v, device) for v in item])\n+    elif isinstance(item, dict):\n+        return type(item)({k: clone_tensors_for_torch_save(v, device) for k, v in item.items()})\n+    else:\n+        return item\n", "issue": "[BUG] Size of saved model checkpoint becomes much larger after deepspeed.initialize when using ZeRO-2\n**Describe the bug**\r\nOriginally reported [here](https://github.com/huggingface/transformers/issues/22822). @stas00 @tjruwase\r\n\r\nFor some models, the size of model checkpoints saved by `model.save_prtrained()` becomes much larger after calling `deepspeed.initialize`. See examples below.\r\n\r\n\r\n**To Reproduce**\r\n```python\r\nfrom transformers import AutoModelForCausalLM\r\nimport deepspeed\r\n\r\nds_config = {\r\n    \"optimizer\": {\r\n        \"type\": \"AdamW\",\r\n    },\r\n    \"zero_optimization\": {\r\n        \"stage\": 2,\r\n        \"offload_optimizer\": {\r\n            \"device\": \"cpu\",\r\n            \"pin_memory\": True\r\n        },\r\n        \"allgather_partitions\": True,\r\n        \"allgather_bucket_size\": 2e8,\r\n        \"overlap_comm\": True,\r\n        \"reduce_scatter\": True,\r\n        \"reduce_bucket_size\": 2e8,\r\n        \"contiguous_gradients\": True\r\n    },\r\n    \"offload_optimizer\": {\r\n        \"device\": \"cpu\",\r\n        \"pin_memory\": True\r\n    },\r\n    \"train_batch_size\": 1,\r\n    \"train_micro_batch_size_per_gpu\": 1\r\n}\r\n\r\nmodel = AutoModelForCausalLM.from_pretrained(\"decapoda-research/llama-7b-hf\")\r\nmodel.save_pretrained(\"before\")\r\ndeepspeed_engine, _, _, _ = deepspeed.initialize(model=model, config_params=ds_config)\r\ndeepspeed_engine.module.save_pretrained(\"after\")\r\n```\r\n\r\nFile sizes:\r\n\r\n```bash\r\ndu -a -h --max-depth=1 before/\r\n512     before/config.json\r\n32K     before/pytorch_model.bin.index.json\r\n9.2G    before/pytorch_model-00001-of-00003.bin\r\n9.3G    before/pytorch_model-00002-of-00003.bin\r\n6.7G    before/pytorch_model-00003-of-00003.bin\r\n512     before/generation_config.json\r\n26G     before/\r\n\r\ndu -a -h --max-depth=1 after/\r\n512     after/config.json\r\n32K     after/pytorch_model.bin.index.json\r\n26G     after/pytorch_model-00001-of-00003.bin\r\n26G     after/pytorch_model-00002-of-00003.bin\r\n26G     after/pytorch_model-00003-of-00003.bin\r\n512     after/generation_config.json\r\n76G     after/\r\n```\r\n\r\nThis 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.\r\n\r\nThis can be fixed by re-clone states before the saving:\r\n```python\r\nstate_dict = deepspeed_engine.module.state_dict()\r\nstate_dict = type(state_dict)(\r\n    {k: v.clone()\r\n     for k,\r\n     v in state_dict.items()})\r\ndeepspeed_engine.module.save_pretrained(\"after_fixed\", state_dict=state_dict)\r\n```\r\n\r\n**Expected behavior**\r\nThe saved model size should be unchanged after `deepspeed.initialize`\r\n\r\n**System info (please complete the following information):**\r\n- deepspeed: 0.8.3\r\n- transformers version: 4.28.0.dev0\r\n- Platform: Linux-4.18.0-372.32.1.el8_6.x86_64-x86_64-with-glibc2.17\r\n- Python version: 3.8.16\r\n- Huggingface_hub version: 0.13.3\r\n- Safetensors version: not installed\r\n- PyTorch version (GPU?): 1.12.1+cu116 (True)\r\n- Tensorflow version (GPU?): not installed (NA)\r\n- Flax version (CPU?/GPU?/TPU?): not installed (NA)\r\n- Jax version: not installed\r\n- JaxLib version: not installed\r\n- Using GPU in script?: yes\r\n- Using distributed or parallel set-up in script?: yes\r\n\r\n\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nimport os\nfrom .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX)\n\n\ndef get_model_ckpt_name_for_rank(base_folder, mp_rank_str):\n    ckpt_name = os.path.join(\n        base_folder,\n        MODEL_FILE_PREFIX + mp_rank_str + MODEL_FILE_SUFFIX,\n    )\n    return ckpt_name\n\n\ndef get_zero_ckpt_name_for_rank(base_folder, dp_rank, mp_rank):\n    zero_prefix = f'{ZERO_FILE_PREFIX}{dp_rank}'\n    mp_rank_string = f'_{MODEL_FILE_PREFIX}{mp_rank:02d}'\n    zero_ckpt_name = os.path.join(\n        base_folder,\n        zero_prefix + mp_rank_string + OPTIM_FILE_SUFFIX,\n    )\n    return zero_ckpt_name\n\n\ndef get_layer_ckpt_name_for_rank(base_folder, layer_id, tp_rank):\n    ckpt_file = f'{layer_id}-model_{tp_rank:02d}{MODEL_FILE_SUFFIX}'\n    ckpt_path = os.path.join(base_folder, ckpt_file)\n    return ckpt_path\n", "path": "deepspeed/checkpoint/utils.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nimport os\nimport torch\nfrom .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX)\n\n\ndef get_model_ckpt_name_for_rank(base_folder, mp_rank_str):\n    ckpt_name = os.path.join(\n        base_folder,\n        MODEL_FILE_PREFIX + mp_rank_str + MODEL_FILE_SUFFIX,\n    )\n    return ckpt_name\n\n\ndef get_zero_ckpt_name_for_rank(base_folder, dp_rank, mp_rank):\n    zero_prefix = f'{ZERO_FILE_PREFIX}{dp_rank}'\n    mp_rank_string = f'_{MODEL_FILE_PREFIX}{mp_rank:02d}'\n    zero_ckpt_name = os.path.join(\n        base_folder,\n        zero_prefix + mp_rank_string + OPTIM_FILE_SUFFIX,\n    )\n    return zero_ckpt_name\n\n\ndef get_layer_ckpt_name_for_rank(base_folder, layer_id, tp_rank):\n    ckpt_file = f'{layer_id}-model_{tp_rank:02d}{MODEL_FILE_SUFFIX}'\n    ckpt_path = os.path.join(base_folder, ckpt_file)\n    return ckpt_path\n\n\n# We pass cloned tensors to torch.save() to avoid checkpoint bloat that occurs when torch.save()\n# saves the underlying storage rather than the slice of the storage corresponding to individual tensors.\n# This is a problem in DeepSpeed because we often allocate tensors using slices of large flattened buffers.\n# Tensor cloning helps to avoid this problem because the storage of cloned tensors are closer to the true size.\n# It is expected that the garbage collector will reclaim the cloned tensor storage to avoid memory bloat.\n# See https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing\ndef clone_tensors_for_torch_save(item, device=torch.device('cpu')):\n    \"\"\"\n    Returns a copy of ``item`` with all enclosed tensors replaced by clones on a specified device.\n    Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.\n\n    Parameters:\n        - ``item``: tensor to clone or (possibly nested) container of tensors to clone.\n        - ``device``: target device (defaults to 'cpu')\n\n    Returns:\n        - copy of ``item`` with cloned tensors on target device\n    \"\"\"\n    if torch.is_tensor(item):\n        return item.detach().clone().to(device)\n    elif isinstance(item, list):\n        return [clone_tensors_for_torch_save(v, device) for v in item]\n    elif isinstance(item, tuple):\n        return tuple([clone_tensors_for_torch_save(v, device) for v in item])\n    elif isinstance(item, dict):\n        return type(item)({k: clone_tensors_for_torch_save(v, device) for k, v in item.items()})\n    else:\n        return item\n", "path": "deepspeed/checkpoint/utils.py"}]}