Add configurable dtype for ZeRO checkpoint export

Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>

Author: gaoxiaomo

deepspeed/runtime/engine.py

@@ -5391,12 +5391,12 @@ def get_layer_state_dict(module, prefix=""):
 
     def _copy_recovery_script(self, save_path):
         base_dir = os.path.dirname(os.path.dirname(__file__))
-        script = "zero_to_fp32.py"
-        src = os.path.join(base_dir, "utils", script)
-        dst = os.path.join(save_path, script)
-        #logger.info(f"creating recovery script {dst}")
-        copyfile(src, dst)
-        self._change_recovery_script_permissions(dst)
+        for script in ("zero_to_fp32.py", "zero_to_torch.py"):
+            src = os.path.join(base_dir, "utils", script)
+            dst = os.path.join(save_path, script)
+            #logger.info(f"creating recovery script {dst}")
+            copyfile(src, dst)
+            self._change_recovery_script_permissions(dst)
 
     def _change_recovery_script_permissions(self, dst):
         # make executable (safeguard for file shares - Azure as example)

deepspeed/utils/zero_to_fp32.py

@@ -53,6 +53,26 @@ class zero_model_state:
 # load to cpu
 device = torch.device('cpu')
 
+OUTPUT_DTYPE_NAMES = {
+    'float32': torch.float32,
+    'fp32': torch.float32,
+    'float16': torch.float16,
+    'fp16': torch.float16,
+    'bfloat16': torch.bfloat16,
+    'bf16': torch.bfloat16,
+}
+
+
+def _resolve_output_dtype(dtype):
+    requested_dtype = dtype
+    if isinstance(dtype, str):
+        dtype_name = dtype[6:] if dtype.startswith('torch.') else dtype
+        dtype = OUTPUT_DTYPE_NAMES.get(dtype_name.lower())
+    if dtype not in set(OUTPUT_DTYPE_NAMES.values()):
+        supported = ', '.join(sorted(OUTPUT_DTYPE_NAMES))
+        raise ValueError(f"Unsupported output dtype {requested_dtype!r}. Choose one of: {supported}")
+    return dtype
+
 
 def atoi(text):
     return int(text) if text.isdigit() else text
@@ -540,10 +560,13 @@ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zer
     return state_dict
 
 
-def to_torch_tensor(state_dict, return_empty_tensor=False):
+def to_torch_tensor(state_dict, return_empty_tensor=False, dtype=None):
     """
     Convert state_dict of GatheredTensor to torch tensor
     """
+    if dtype is not None:
+        dtype = _resolve_output_dtype(dtype)
+
     torch_state_dict = {}
     converted_tensors = {}
     for name, tensor in state_dict.items():
@@ -554,9 +577,10 @@ def to_torch_tensor(state_dict, return_empty_tensor=False):
         else:
             converted_tensors[tensor_id] = name
             if return_empty_tensor:
-                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=dtype or tensor.dtype)
             else:
-                torch_state_dict[name] = tensor.contiguous()
+                contiguous_tensor = tensor.contiguous()
+                torch_state_dict[name] = contiguous_tensor.to(dtype=dtype) if dtype else contiguous_tensor
     return torch_state_dict
 
 
@@ -625,25 +649,29 @@ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
         return to_torch_tensor(state_dict)
 
 
-def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
-                                               output_dir,
-                                               max_shard_size="5GB",
-                                               safe_serialization=False,
-                                               tag=None,
-                                               exclude_frozen_parameters=False):
+def convert_zero_checkpoint_to_state_dict(checkpoint_dir,
+                                          output_dir,
+                                          dtype=torch.float32,
+                                          max_shard_size="5GB",
+                                          safe_serialization=False,
+                                          tag=None,
+                                          exclude_frozen_parameters=False):
     """
-    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    Convert ZeRO 2 or 3 checkpoint into a consolidated ``state_dict`` file that can be
     loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
 
     Args:
         - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
-        - ``output_dir``: directory to the pytorch fp32 state_dict output files
+        - ``output_dir``: directory for the PyTorch state_dict output files
+        - ``dtype``: output tensor dtype. Supports float32, float16, and bfloat16 as strings or torch dtypes.
         - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
         - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
         - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
         - ``exclude_frozen_parameters``: exclude frozen parameters
     """
 
+    dtype = _resolve_output_dtype(dtype)
+
     # Dependency pre-check
     if safe_serialization:
         try:
@@ -669,7 +697,7 @@ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
     if max_shard_size is not None:
         filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
         # an memory-efficient approach for sharding
-        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True, dtype=dtype)
         state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
                                                               filename_pattern=filename_pattern,
                                                               max_shard_size=max_shard_size)
@@ -684,7 +712,7 @@ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
     filename_to_tensors = state_dict_split.filename_to_tensors.items()
     for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
         shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
-        shard_state_dict = to_torch_tensor(shard_state_dict)
+        shard_state_dict = to_torch_tensor(shard_state_dict, dtype=dtype)
         output_path = os.path.join(output_dir, shard_file)
         if safe_serialization:
             save_file(shard_state_dict, output_path, metadata={"format": "pt"})
@@ -710,6 +738,22 @@ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
             f.write(content)
 
 
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """Backward-compatible fp32 checkpoint conversion."""
+    return convert_zero_checkpoint_to_state_dict(checkpoint_dir,
+                                                 output_dir,
+                                                 dtype=torch.float32,
+                                                 max_shard_size=max_shard_size,
+                                                 safe_serialization=safe_serialization,
+                                                 tag=tag,
+                                                 exclude_frozen_parameters=exclude_frozen_parameters)
+
+
 def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
     """
     1. Put the provided model to cpu

deepspeed/utils/zero_to_torch.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+
+import argparse
+
+import deepspeed.utils.zero_to_fp32 as zero_to_fp32
+
+
+def main(args=None):
+    parser = argparse.ArgumentParser(
+        description="Convert a DeepSpeed ZeRO checkpoint to float32, float16, or bfloat16 PyTorch weights.")
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument("output_dir", type=str, help="directory for the converted PyTorch state_dict files")
+    parser.add_argument("--dtype",
+                        type=str,
+                        choices=sorted(zero_to_fp32.OUTPUT_DTYPE_NAMES),
+                        required=True,
+                        help="output tensor dtype")
+    parser.add_argument("--max_shard_size",
+                        type=str,
+                        default="5GB",
+                        help="maximum size of each checkpoint shard, such as 5GB or 500MB")
+    parser.add_argument("--safe_serialization",
+                        default=False,
+                        action='store_true',
+                        help="save with safetensors instead of PyTorch pickle serialization")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier, e.g., global_step1")
+    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug output")
+    parsed_args = parser.parse_args(args)
+
+    zero_to_fp32.debug = parsed_args.debug
+    zero_to_fp32.convert_zero_checkpoint_to_state_dict(parsed_args.checkpoint_dir,
+                                                       parsed_args.output_dir,
+                                                       dtype=parsed_args.dtype,
+                                                       max_shard_size=parsed_args.max_shard_size,
+                                                       safe_serialization=parsed_args.safe_serialization,
+                                                       tag=parsed_args.tag,
+                                                       exclude_frozen_parameters=parsed_args.exclude_frozen_parameters)
+
+
+if __name__ == "__main__":
+    main()

docs/_tutorials/zero.md

@@ -282,7 +282,35 @@ Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)
 
 The `zero_to_fp32.py` script gets created automatically when you save a checkpoint.
 
-Note: currently this script uses 2x memory (general RAM) of the size of the final checkpoint.
+To reduce the exported checkpoint size, use the accompanying `zero_to_torch.py`
+script and choose `float16` or `bfloat16`. The original `zero_to_fp32.py`
+behavior remains unchanged.
+
+```bash
+$ ./zero_to_torch.py . checkpoint-bf16 --dtype bfloat16
+```
+
+The conversion reconstructs the ZeRO master weights in fp32, casts each output
+shard immediately before it is saved, and plans shard sizes using the selected
+dtype. Shared parameters remain shared in the exported state dict. The same
+functionality is available from Python:
+
+```python
+import torch
+from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_state_dict
+
+convert_zero_checkpoint_to_state_dict(
+    checkpoint_dir,
+    output_dir,
+    dtype=torch.bfloat16,
+)
+```
+
+Questions and maintenance: [@gaoxiaomo](https://github.com/gaoxiaomo).
+
+Note: fp32 conversion currently uses about 2x the final checkpoint size in CPU
+memory. Lower-precision export still reconstructs the fp32 master weights, so
+its peak memory is larger than 2x the final fp16/bf16 checkpoint size.
 {: .notice--info}
 
 Alternatively, if you have plenty of spare CPU memory and instead of getting the file you want your model to be updated to its fp32 weights, you can do the following at the end of the training:

docs/code-docs/source/model-checkpointing.rst

@@ -23,6 +23,8 @@ DeepSpeed provides routines for extracting fp32 weights from the saved ZeRO chec
 
 .. autofunction:: deepspeed.utils.zero_to_fp32.convert_zero_checkpoint_to_fp32_state_dict
 
+.. autofunction:: deepspeed.utils.zero_to_fp32.convert_zero_checkpoint_to_state_dict
+
 
 Avoiding ZeRO Checkpoint Bloat
 ------------------------------

tests/unit/checkpoint/test_convert_checkpoint.py

@@ -5,12 +5,68 @@
 
 import torch
 import torch.nn as nn
+import pytest
 
 import deepspeed
-from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
+import deepspeed.utils.zero_to_fp32 as zero_to_fp32
+from deepspeed.utils.zero_to_fp32 import (convert_zero_checkpoint_to_fp32_state_dict,
+                                          convert_zero_checkpoint_to_state_dict, to_torch_tensor)
+from deepspeed.utils.zero_to_torch import main as zero_to_torch_main
 from unit.common import DistributedTest
 
 
+def test_output_dtype_conversion_preserves_shared_tensors():
+    tensor = torch.arange(16, dtype=torch.float32)
+    state_dict = {"weight": tensor, "shared_weight": tensor}
+
+    converted = to_torch_tensor(state_dict, dtype="bf16")
+    assert converted["weight"].dtype == torch.bfloat16
+    assert id(converted["weight"]) == id(converted["shared_weight"])
+
+    empty = to_torch_tensor(state_dict, return_empty_tensor=True, dtype="fp16")
+    assert empty["weight"].dtype == torch.float16
+    assert id(empty["weight"]) == id(empty["shared_weight"])
+
+    with pytest.raises(ValueError, match="Unsupported output dtype"):
+        to_torch_tensor(state_dict, dtype="int8")
+
+
+def test_checkpoint_file_output_dtype(monkeypatch, tmp_path):
+
+    def make_state_dict(*args, **kwargs):
+        weight = torch.linspace(-1, 1, 4096, dtype=torch.float32)
+        return {"weight": weight, "shared_weight": weight}
+
+    monkeypatch.setattr(zero_to_fp32, "get_fp32_state_dict_from_zero_checkpoint", make_state_dict)
+    fp32_dir = tmp_path / "fp32"
+    bf16_dir = tmp_path / "bf16"
+
+    convert_zero_checkpoint_to_fp32_state_dict("unused", fp32_dir, max_shard_size=None)
+    convert_zero_checkpoint_to_state_dict("unused", bf16_dir, dtype="bf16", max_shard_size=None)
+
+    fp32_state_dict = torch.load(fp32_dir / "pytorch_model.bin")
+    bf16_state_dict = torch.load(bf16_dir / "pytorch_model.bin")
+    assert bf16_state_dict["weight"].dtype == torch.bfloat16
+    assert id(bf16_state_dict["weight"]) == id(bf16_state_dict["shared_weight"])
+    torch.testing.assert_close(bf16_state_dict["weight"].float(), fp32_state_dict["weight"], rtol=5e-3, atol=5e-3)
+    assert (bf16_dir / "pytorch_model.bin").stat().st_size < (fp32_dir / "pytorch_model.bin").stat().st_size * 0.6
+
+
+def test_zero_to_torch_cli_passes_dtype(monkeypatch, tmp_path):
+    call = {}
+
+    def record_conversion(checkpoint_dir, output_dir, **kwargs):
+        call.update(checkpoint_dir=checkpoint_dir, output_dir=output_dir, **kwargs)
+
+    monkeypatch.setattr(zero_to_fp32, "convert_zero_checkpoint_to_state_dict", record_conversion)
+    zero_to_torch_main(["checkpoint", str(tmp_path), "--dtype", "fp16", "--max_shard_size", "1GB"])
+
+    assert call["checkpoint_dir"] == "checkpoint"
+    assert call["output_dir"] == str(tmp_path)
+    assert call["dtype"] == "fp16"
+    assert call["max_shard_size"] == "1GB"
+
+
 class ModelWithSharedWeights(nn.Module):
 
     def __init__(self):
@@ -43,6 +99,8 @@ def test_convert_zero_checkpoint_to_fp32_state_dict(self, tmpdir):
         )
         ds_save_dir = tmpdir / "checkpoint_ds"
         deepspeed_engine.save_checkpoint(ds_save_dir, tag="checkpoint")
+        assert (ds_save_dir / "zero_to_fp32.py").exists()
+        assert (ds_save_dir / "zero_to_torch.py").exists()
 
         model = ModelWithSharedWeights()
 
@@ -58,3 +116,20 @@ def test_convert_zero_checkpoint_to_fp32_state_dict(self, tmpdir):
 
         # load state_dict into model
         model.load_state_dict(state_dict, strict=True)
+
+        # Exporting in bfloat16 uses the target dtype for both shard planning
+        # and serialization. At 300KB this model fits in one bf16 shard but
+        # would be split if shard planning still counted fp32 bytes.
+        bf16_save_dir = tmpdir / "checkpoint_bf16"
+        convert_zero_checkpoint_to_state_dict(ds_save_dir, bf16_save_dir, dtype="bfloat16", max_shard_size="300KB")
+        bf16_state_dict = torch.load(bf16_save_dir / 'pytorch_model.bin')
+        assert not (bf16_save_dir / 'pytorch_model.bin.index.json').exists()
+
+        assert id(bf16_state_dict['layer1.weight']) == id(bf16_state_dict['layer2.weight'])
+        for name, tensor in bf16_state_dict.items():
+            assert tensor.dtype == torch.bfloat16
+            torch.testing.assert_close(tensor.float(), state_dict[name], rtol=5e-3, atol=5e-3)
+
+        fp32_size = (fp32_save_dir / 'pytorch_model.bin').stat().st_size
+        bf16_size = (bf16_save_dir / 'pytorch_model.bin').stat().st_size
+        assert bf16_size < fp32_size * 0.6