chore: import upstream snapshot with attribution

deepspeed/utils/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .logging import logger, log_dist, log_dist_once, set_log_level_from_string
+from .comms_logging import get_caller_func
+#from .distributed import init_distributed
+from .init_on_device import OnDevice
+from .groups import *
+from .nvtx import instrument_w_nvtx
+# TODO: Move tensor fragment and mixed precision to zero utils
+from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address, get_full_hp_grad, map_to_flat_opt_states
+from .tensor_fragment import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
+from .tensor_fragment import set_full_hp_param, set_full_hp_grad
+from .tensor_fragment import safe_set_full_fp32_param, safe_set_full_optimizer_state, safe_set_full_grad
+from .tensor_fragment import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state
+from .tensor_fragment import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state
+from .tensor_fragment import safe_update_full_grad_vectorized
+from .z3_leaf_module import set_z3_leaf_modules, unset_z3_leaf_modules, get_z3_leaf_modules, z3_leaf_module, z3_leaf_parameter, set_z3_leaf_module, set_z3_leaf_modules_by_name, set_z3_leaf_modules_by_suffix
+from .mixed_precision_linkage import link_hp_params, lazy_init_hp_params_optimizer_state
+from deepspeed.runtime.dataloader import RepeatingLoader
+from .numa import get_numactl_cmd

deepspeed/utils/bwc.py

@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+
+def bwc_tensor_model_parallel_rank(mpu=None):
+    """Backwards-compatible way of querying the tensor model parallel rank from
+    an ``mpu`` object.
+
+    *Tensor* model parallelism means that tensors are physically split across
+    processes. This contrasts with *pipeline* model parallelism, in which the
+    layers are partitioned but tensors left intact.
+
+    The API for tensor model parallelism has changed across versions and this
+    helper provides a best-effort implementation across versions of ``mpu``
+    objects.  The preferred mechanism is
+    ``mpu.get_tensor_model_parallel_rank()``.
+
+    This should "just work" with both Megatron-LM and DeepSpeed's pipeline
+    parallelism.
+
+    Args:
+        mpu (model parallel unit, optional): The tensor model parallel rank.
+            If ``mpu=None``, returns 0. Defaults to ``None``.
+
+    Returns:
+        int: the rank
+    """
+    if mpu is None:
+        # No model parallelism in easy :)
+        return 0
+
+    if hasattr(mpu, 'get_tensor_model_parallel_rank'):
+        # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+        return mpu.get_tensor_model_parallel_rank()
+    elif hasattr(mpu, 'get_slice_parallel_rank'):
+        # Some DeepSpeed + pipeline parallelism versions
+        return mpu.get_slice_parallel_rank()
+    else:
+        # Deprecated Megatron and DeepSpeed convention
+        return mpu.get_model_parallel_rank()
+
+
+def bwc_tensor_model_parallel_world_size(mpu=None):
+    """Backwards-compatible way of querying the tensor model parallel world size.
+       Similar to bwc_tensor_model_parallel_rank.
+    """
+    if mpu is None:
+        return 1
+
+    if hasattr(mpu, 'get_tensor_model_parallel_world_size'):
+        # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+        return mpu.get_tensor_model_parallel_world_size()
+    elif hasattr(mpu, 'get_slice_parallel_world_size'):
+        # Some DeepSpeed + pipeline parallelism versions
+        return mpu.get_slice_parallel_world_size()
+    else:
+        # Deprecated Megatron and DeepSpeed convention
+        return mpu.get_model_parallel_world_size()
+
+
+def bwc_tensor_model_parallel_group(mpu=None):
+    """Backwards-compatible way of querying the tensor model parallel group.
+       Similar to bwc_tensor_model_parallel_rank.
+    """
+    if mpu is None:
+        return None
+
+    if hasattr(mpu, 'get_tensor_model_parallel_group'):
+        # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+        return mpu.get_tensor_model_parallel_group()
+    elif hasattr(mpu, 'get_slice_parallel_group'):
+        # Some DeepSpeed + pipeline parallelism versions
+        return mpu.get_slice_parallel_group()
+    else:
+        # Deprecated Megatron and DeepSpeed convention
+        return mpu.get_model_parallel_group()
+
+
+def bwc_pipeline_parallel_world_size(mpu=None):
+    """Backwards-compatible way of querying the pipeline parallel world size."""
+    world_size = 1
+    if mpu is not None:
+        if hasattr(mpu, 'get_pipeline_model_parallel_world_size'):
+            # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+            world_size = mpu.get_pipeline_model_parallel_world_size()
+        elif hasattr(mpu, 'get_pipe_parallel_world_size'):
+            # DeepSpeed Topology
+            world_size = mpu.get_pipe_parallel_world_size()
+    return world_size
+
+
+def bwc_pipeline_parallel_group(mpu=None):
+    """Backwards-compatible way of querying the pipeline parallel group."""
+    if mpu is None:
+        return None
+    if hasattr(mpu, 'get_pipeline_model_parallel_group'):
+        # Megatron
+        return mpu.get_pipeline_model_parallel_group()
+    elif hasattr(mpu, 'get_pipe_parallel_group'):
+        # DeepSpeed Topology
+        return mpu.get_pipe_parallel_group()
+    assert False, 'mpu does not support pipeline parallel group'

deepspeed/utils/comms_logging.py

@@ -0,0 +1,378 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import math
+from deepspeed.utils import log_dist
+
+
+def get_caller_func(frame=3):
+    import sys
+    return sys._getframe(frame).f_code.co_name
+
+
+def print_rank_0(message):
+    import deepspeed.comm as dist
+    if dist.get_rank() == 0:
+        print(message)
+
+
+# Helper function to pretty-print message sizes
+def convert_size(size_bytes):
+    if size_bytes == 0:
+        return "0B"
+    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
+    i = int(math.floor(math.log(size_bytes, 1024)))
+    p = math.pow(1024, i)
+    s = round(size_bytes / p, 2)
+    return "%s %s" % (s, size_name[i])
+
+
+# Helper function to calculate algbw and busbw.
+# See https://gist.github.com/jeffra/b5e80466b4c86be00ea3b6f130fb7a36 and https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md
+def calc_bw_log(comm_op, size, duration):
+    import deepspeed.comm as dist
+
+    n = dist.get_world_size()
+    tput = 0
+    busbw = 0
+    if comm_op == "all_to_all_single":
+        tput = (size / duration)
+        busbw = (size / duration) * ((n - 1) / n)
+    elif comm_op == "all_gather" or comm_op == "all_gather_into_tensor" or comm_op == "reduce_scatter" or comm_op == "reduce_scatter_tensor":
+        size *= n
+        tput = (size / duration)
+        busbw = (size / duration) * ((n - 1) / n)
+    elif comm_op == "all_reduce" or comm_op == "all_reduce_coalesced" or comm_op == "inference_all_reduce":
+        tput = (size * 2 / duration)
+        busbw = (size / duration) * (2 * (n - 1) / n)
+    elif comm_op == "send" or comm_op == "recv" or comm_op == "isend" or comm_op == "irecv" or comm_op == "broadcast" or comm_op == "reduce" or comm_op == "gather" or comm_op == "scatter" or comm_op == "barrier":
+        tput = (size / duration)
+        busbw = tput
+    else:
+        print_rank_0("wrong comm_op specified")  # noqa: F821
+        exit(0)
+
+    # convert to Gbps
+    tput *= 8
+    busbw *= 8
+
+    tput /= 1e6
+    busbw /= 1e6
+
+    return tput, busbw
+
+
+class CommsLogger:
+
+    def __init__(self):
+        from deepspeed.comm.constants import COMMS_LOGGER_VERBOSE_DEFAULT, COMMS_LOGGER_DEBUG_DEFAULT, COMMS_LOGGER_PROF_OPS_DEFAULT, COMMS_LOGGER_PROF_ALL_DEFAULT, COMMS_LOGGER_ENABLED_DEFAULT
+        self.comms_dict = {}
+        self.verbose = COMMS_LOGGER_VERBOSE_DEFAULT
+        self.debug = COMMS_LOGGER_DEBUG_DEFAULT
+        self.prof_ops = COMMS_LOGGER_PROF_OPS_DEFAULT
+        self.prof_all = COMMS_LOGGER_PROF_ALL_DEFAULT
+        self.enabled = COMMS_LOGGER_ENABLED_DEFAULT
+
+    def configure(self, comms_config):
+        self.enabled = comms_config.comms_logger_enabled
+        if self.enabled:
+            self.verbose = comms_config.comms_logger.verbose
+            self.debug = comms_config.comms_logger.debug
+            self.prof_ops = comms_config.comms_logger.prof_ops
+            self.prof_all = comms_config.comms_logger.prof_all
+
+    # There are three settings for the op profiler:
+    # - Global profiling (profile all comms)
+    # - Op-type profiling (e.g. profile all all_reduce comms)
+    # - Op profiling (e.g. profile a specific all_reduce op)
+    def start_profiling_comms(self):
+        self.prof_all = True
+
+    def stop_profiling_comms(self):
+        self.prof_all = True
+
+    # E.g. start_profiling_op('all_reduce')
+    def start_profiling_op(self, op_name_list):
+        self.prof_ops = list(set(self.prof_ops) | set(op_name_list))
+
+    def stop_profiling_op(self, op_name_list):
+        self.prof_ops = [op for op in self.prof_ops if op not in op_name_list]
+
+    # Add log entry
+    def append(self, raw_name, record_name, latency, msg_size):
+        algbw, busbw = calc_bw_log(raw_name, msg_size, latency)
+        if record_name in self.comms_dict.keys():
+            # If this comm_op has already been logged with this message size, just add to existing record
+            if msg_size in self.comms_dict[record_name].keys():
+                self.comms_dict[record_name][msg_size][0] += 1
+                self.comms_dict[record_name][msg_size][1].append(latency)
+                self.comms_dict[record_name][msg_size][2].append(algbw)
+                self.comms_dict[record_name][msg_size][3].append(busbw)
+            # If this is a new message size for this comm_op, add new record under existing comm_op
+            else:
+                self.comms_dict[record_name][msg_size] = [1, [latency], [algbw], [busbw]]
+        else:
+            # Create entirely new record
+            self.comms_dict[record_name] = {msg_size: [1, [latency], [algbw], [busbw]]}
+        # If verbose, print every comm op
+        # TODO: Add to tensorboard
+        if self.verbose:
+            log_str = f"comm op: {record_name} | time (ms): {latency:.2f} | msg size: {convert_size(msg_size)} | algbw (Gbps): {algbw:.2f} | busbw (Gbps): {busbw:.2f}"
+            log_dist(log_str, [0])
+
+    def get_raw_data(self):
+        """
+        Get the raw communication data dictionary.
+
+        Returns:
+            dict: Raw communication data in format {record_name: {msg_size: [count, [latencies], [algbws], [busbws]]}}
+        """
+        return self.comms_dict.copy()
+
+    def has_data(self):
+        """
+        Check if any communication data has been logged.
+
+        Returns:
+            bool: True if communication data exists, False otherwise
+        """
+        return len(self.comms_dict) > 0
+
+    def reset_data(self):
+        """
+        Clear all logged communication data.
+        """
+        self.comms_dict.clear()
+
+    def get_operation_names(self):
+        """
+        Get list of all logged communication operation names.
+
+        Returns:
+            list: List of operation names that have been logged
+        """
+        return list(self.comms_dict.keys())
+
+    def get_total_operations(self):
+        """
+        Get total number of communication operations logged across all types.
+
+        Returns:
+            int: Total count of operations
+        """
+        total = 0
+        for record_name in self.comms_dict:
+            for msg_size in self.comms_dict[record_name]:
+                total += self.comms_dict[record_name][msg_size][0]  # count is at index 0
+        return total
+
+    def get_operation_summary(self, operation_name):
+        """
+        Get summary statistics for a specific operation type.
+
+        Args:
+            operation_name (str): Name of the communication operation
+
+        Returns:
+            dict: Summary statistics for the operation, or None if operation not found
+        """
+        if operation_name not in self.comms_dict:
+            return None
+
+        from deepspeed.utils.timer import trim_mean
+
+        # Create a snapshot to avoid concurrent modification issues
+        op_data = self.comms_dict[operation_name].copy()
+        summary = {}
+
+        for msg_size, vals in op_data.items():
+            count = vals[0]
+            total_lat = sum(vals[1])
+            avg_lat = trim_mean(vals[1], 0.1)
+            avg_algbw = trim_mean(vals[2], 0.1)
+            avg_busbw = trim_mean(vals[3], 0.1)
+
+            summary[msg_size] = {
+                "count": count,
+                "total_latency_ms": total_lat,
+                "avg_latency_ms": avg_lat,
+                "tput_avg_gbps": avg_algbw,
+                "busbw_avg_gbps": avg_busbw,
+                "msg_size_bytes": msg_size,
+                "msg_size_str": convert_size(msg_size)
+            }
+
+        return summary
+
+    # Print summary at end of iteration, epoch, or training
+    def log_all(self, print_log=True, show_straggler=False, return_dict=False):
+        """
+        Print and/or return communication operation statistics.
+
+        Args:
+            print_log (bool, optional): Whether to print the summary to console. Defaults to True.
+            show_straggler (bool, optional): Whether to include straggler effect analysis. Defaults to False.
+            return_dict (bool, optional): Whether to return statistics as a dictionary. Defaults to False.
+
+        Returns:
+            dict or None: If return_dict=True, returns a comprehensive dictionary with the following structure:
+            {
+                "summary": {
+                    "operation_name": {
+                        message_size_bytes: {
+                            "count": int,                    # Number of operations with this message size
+                            "total_latency_ms": float,      # Sum of all latencies for this message size
+                            "avg_latency_ms": float,        # Average latency (outliers trimmed)
+                            "tput_avg_gbps": float,         # Average algorithmic bandwidth in Gbps
+                            "busbw_avg_gbps": float,        # Average bus bandwidth in Gbps
+                            "msg_size_bytes": int,          # Message size in bytes
+                            "msg_size_str": str             # Human-readable message size (e.g., "678.86 MB")
+                        }
+                    }
+                },
+                "straggler_analysis": {                     # Only present if show_straggler=True
+                    "operation_name": {
+                        message_size_bytes: {
+                            "count": int,                    # Number of operations
+                            "total_comm_lat_ms": float,     # Total communication latency (min across ranks)
+                            "total_straggler_ms": float,    # Total straggler effect
+                            "avg_comm_lat_ms": float,       # Average communication latency
+                            "avg_straggler_ms": float,      # Average straggler effect
+                            "msg_size_bytes": int,          # Message size in bytes
+                            "msg_size_str": str             # Human-readable message size
+                        }
+                    }
+                } if show_straggler else None,
+                "metadata": {
+                    "world_size": int,                      # Number of processes in distributed setup
+                    "rank": int,                            # Current process rank
+                    "timestamp": str                        # ISO format timestamp when log_all was called
+                }
+            }
+
+            Returns None if return_dict=False.
+
+        Note:
+            - Statistics use trimmed mean (10% trimmed from both ends) to remove outliers
+            - Straggler analysis requires distributed communication and may impact performance
+            - All bandwidth values are in Gbps (Gigabits per second)
+            - Latency values are in milliseconds
+        """
+        import torch
+        from deepspeed.utils.timer import trim_mean
+        import deepspeed.comm as dist
+        from deepspeed.comm.reduce_op import ReduceOp
+        from deepspeed.accelerator import get_accelerator
+        from datetime import datetime
+
+        # Create a snapshot of the dictionary to avoid concurrent modification issues
+        # This prevents "dictionary changed size during iteration" errors when
+        # communication operations are happening in other threads
+        comms_dict_snapshot = self.comms_dict.copy()
+
+        # Initialize return dictionary structure
+        result_dict = {
+            "summary": {},
+            "straggler_analysis": None,
+            "metadata": {
+                "world_size": dist.get_world_size() if dist.is_initialized() else 1,
+                "rank": dist.get_rank() if dist.is_initialized() else 0,
+                "timestamp": datetime.now().isoformat()
+            }
+        } if return_dict else None
+
+        if print_log:
+            print(
+                f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total Latency(ms)': <20}{'Avg Latency(ms)': <20}{'tput_avg (Gbps)': <20}{'busbw_avg (Gbps)': <20}"
+            )
+
+        for record_name in comms_dict_snapshot.keys():
+            if print_log:
+                print(record_name)
+
+            # Initialize operation entry in result dict
+            if return_dict:
+                result_dict["summary"][record_name] = {}
+
+            for msg_size, vals in sorted(comms_dict_snapshot[record_name].items()):
+                # vals[0] is the count for each msg size
+                count = vals[0]
+                # vals[1] is a list of latency records for each msg size
+                total_lat = sum(vals[1])
+                # vals[2] and vals[3] are the lists of algbw and busbw, respectively
+                # Get rid of outliers when we print
+                avg_lat = trim_mean(vals[1], 0.1)
+                avg_algbw = trim_mean(vals[2], 0.1)
+                avg_busbw = trim_mean(vals[3], 0.1)
+
+                # Store data in result dictionary
+                if return_dict:
+                    result_dict["summary"][record_name][msg_size] = {
+                        "count": count,
+                        "total_latency_ms": total_lat,
+                        "avg_latency_ms": avg_lat,
+                        "tput_avg_gbps": avg_algbw,
+                        "busbw_avg_gbps": avg_busbw,
+                        "msg_size_bytes": msg_size,
+                        "msg_size_str": convert_size(msg_size)
+                    }
+
+                if print_log:
+                    print(
+                        f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{avg_lat: <20.2f}{avg_algbw: <20.2f}{avg_busbw: <20.2f}"
+                    )
+
+        if show_straggler:
+            if return_dict:
+                result_dict["straggler_analysis"] = {}
+
+            if print_log:
+                print("_______________________________")
+                print("Breakdown with straggler effect")
+                print("-------------------------------")
+                print(
+                    f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total comm lat(ms)': <20}{'Total straggler(ms)': <20}{'Avg comm lat(ms)': <20}{'Avg straggler(ms)': <20}"
+                )
+
+            device = get_accelerator().current_device_name()
+            for record_name in comms_dict_snapshot.keys():
+                if print_log:
+                    print(record_name)
+
+                # Initialize operation entry in straggler dict
+                if return_dict:
+                    result_dict["straggler_analysis"][record_name] = {}
+
+                for msg_size, vals in sorted(comms_dict_snapshot[record_name].items()):
+                    # vals[0] is the count for each msg size
+                    count = vals[0]
+                    # vals[1] is a list of latency records for each msg size
+                    lats = torch.tensor(vals[1], device=device)
+                    min_lats = torch.tensor(vals[1], device=device)
+                    dist.all_reduce(min_lats, op=ReduceOp.MIN)
+                    total_lat = min_lats.sum().item()
+                    total_straggler = (lats - min_lats).sum().item()
+                    avg_lat = trim_mean(min_lats.tolist(), 0.1)
+                    avg_straggler = trim_mean((lats - min_lats).tolist(), 0.1)
+
+                    # Store straggler data in result dictionary
+                    if return_dict:
+                        result_dict["straggler_analysis"][record_name][msg_size] = {
+                            "count": count,
+                            "total_comm_lat_ms": total_lat,
+                            "total_straggler_ms": total_straggler,
+                            "avg_comm_lat_ms": avg_lat,
+                            "avg_straggler_ms": avg_straggler,
+                            "msg_size_bytes": msg_size,
+                            "msg_size_str": convert_size(msg_size)
+                        }
+
+                    if print_log:
+                        print(
+                            f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{total_straggler: <20.2f}{avg_lat: <20.2f}{avg_straggler: <20.2f}"
+                        )
+
+        # Return the dictionary if requested
+        return result_dict if return_dict else None

deepspeed/utils/config.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+
+#########################################
+# Timers
+#########################################
+# Timers. By default, timers are enabled.
+# Users can configure in ds_config.json as below example:
+TIMERS_FORMAT = '''
+Timers should be enabled as:
+"timers": {
+  "throughput": {
+    "enabled": true,
+    "synchronized": true
+  }
+}
+'''
+
+TIMERS = "timers"
+TIMERS_THROUGHPUT = "throughput"
+
+
+def get_timers_config(param_dict):
+    if param_dict and TIMERS in param_dict and TIMERS_THROUGHPUT in param_dict[TIMERS]:
+        timers_config_dict = param_dict[TIMERS][TIMERS_THROUGHPUT]
+    else:
+        timers_config_dict = {}
+    return DeepSpeedThroughputTimerConfig(**timers_config_dict)
+
+
+class DeepSpeedThroughputTimerConfig(DeepSpeedConfigModel):
+    """ Configure throughput timers """
+
+    enabled: bool = True
+    """ Turn on/off throughput timers """
+
+    synchronized: bool = True
+    """ Whether to synchronize a device when measuring the time.
+        Synchronizing a device is required to produce the most accurate timer measurements.
+        However, this comes at the expense of performance degradation. The CPU timer provides
+        sufficient accuracy in many cases.
+      """

deepspeed/utils/debug.py

@@ -0,0 +1,187 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed.comm as dist
+
+# For lazy import with printflock()
+fcntl = None
+
+# for debug purposes map module and param objects to their fully qualified names
+module_names = {}
+param_names = {}
+
+
+def debug_clear_module_and_param_names():
+    global module_names
+    global param_names
+    module_names = {}
+    param_names = {}
+
+
+def debug_extract_module_and_param_names(model):
+    # extract the fully qualified names as soon as the model is acquired
+    global module_names
+    global param_names
+    # XXX: can probably make a map of param2module and vice-versa
+    module_names = {module: name for name, module in model.named_modules()}
+    param_names = {param: name for name, param in model.named_parameters()}
+
+
+def debug_module2name(module):
+    if module in module_names:
+        return module_names[module]
+    else:
+        return "unknown"
+
+
+def debug_module2name_id(module):
+    return f"name={debug_module2name(module)}"
+
+
+def debug_module2name_class(module):
+    return f"name={debug_module2name(module)} {module.__class__.__name__}"
+
+
+def debug_param2name(param):
+    if param in param_names:
+        return param_names[param]
+    else:
+        return "unknown"
+
+
+def ds_id(param):
+    if hasattr(param, "ds_id"):
+        return param.ds_id
+    else:
+        return "none"
+
+
+def ds_shape(param):
+    if hasattr(param, "ds_shape"):
+        return param.ds_shape
+    else:
+        return param.shape
+
+
+def debug_param2name_id(param):
+    return f"name={debug_param2name(param)} id={ds_id(param)}"
+
+
+def debug_param2name_id_shape(param):
+    return f"name={debug_param2name(param)} id={ds_id(param)} shape={ds_shape(param)}"
+
+
+def debug_param2name_id_shape_device(param):
+    return f"name={debug_param2name(param)} id={ds_id(param)} shape={ds_shape(param)} device={param.device}"
+
+
+def debug_param2name_id_numel(param):
+    return f"name={debug_param2name(param)} id={ds_id(param)} numel={param.numel()}"
+
+
+def debug_param2name_id_shape_status(param):
+    return f"name={debug_param2name(param)} id={ds_id(param)} shape={ds_shape(param)} status={param.ds_status}"
+
+
+def printflock(*msgs):
+    """
+
+    For printing messages for all concurrent gpus w/o getting interleaved text.
+
+    This is useful when debugging issues where multi-gpus don't sync.
+
+    1. Enable the force debug in say partitioning and zero3 files
+    2. Override the usual versions with ::
+
+        def print_rank_0(message, debug=False, force=False):
+            rank = deepspeed.comm.get_rank()
+            printflock(f"[{rank}] {message}")
+    3. run the program and you get both logs non-interleaved
+
+    But this makes it very difficult to make sense of the output, so the ``log_rank_file`` helper
+    function might be more useful, as it's easier to send each log stream into a separate file and
+    then compare those.
+
+    """
+    global fcntl
+    if fcntl is None:
+        import fcntl
+
+    with open(__file__, "r") as fh:
+        fcntl.flock(fh, fcntl.LOCK_EX)
+        try:
+            print(*msgs)
+        finally:
+            fcntl.flock(fh, fcntl.LOCK_UN)
+
+
+fh = None
+
+
+def log_rank_file(rank, *msgs):
+    """
+    Print to a log file of the given rank
+
+    This is useful for debugging hanging in sync processes. Here is a possible workflow:
+
+    1. Enable the force debug in say partitioning and zero3 files
+    2. Override the usual versions of print_rank_0 in those files with ::
+
+        def print_rank_0(message, debug=False, force=False):
+            rank = deepspeed.comm.get_rank()
+            log_rank_file(rank, message)
+
+    3. run the program
+    4. fix up the expected differences, e.g. different cuda numbers ::
+
+        perl -pi -e 's|cuda:1|cuda:0|' log_rank_*
+
+    5. now diff and see where names and ids diverge - you will find where the gpus don't do the same
+    work (e.g. when some layers get conditionally skipped on one gpu but not all)
+
+        diff -u log_rank_0.txt log_rank_1.txt | less
+
+    """
+    global fh
+    if fh is None:
+        fh = open(f"log_rank_{rank}.txt", "w")
+    for m in msgs:
+        fh.write(f"{m}\n")
+    fh.flush()
+
+
+def print_backward_tensors(tensor):
+
+    def _print_bwd_tensors(grad_fn):
+        print(f"Backward tensors in {grad_fn}")
+        for funcs in grad_fn.next_functions:
+            if funcs[0]:
+                try:
+                    tensor = getattr(funcs[0], 'variable')
+                    print(funcs[0])
+                    print(f"Tensor - id: {id(tensor)}, shape: {tensor.shape}, data: {tensor}, grad: {tensor.grad}")
+                except AttributeError as e:
+                    _print_bwd_tensors(funcs[0])
+
+    if hasattr(tensor, 'grad_fn'):
+        _print_bwd_tensors(tensor.grad_fn)
+
+
+def print_rank(*msg, force=False):
+    """print something on all global ranks with [rank] prefix.
+    """
+    if not force:
+        return
+    global_rank = dist.get_rank()
+    print(f"[{global_rank}]", *msg)
+
+
+def print_rank0(*msg, force=False):
+    """print something only on rank 0"""
+    if not force:
+        return
+    global_rank = dist.get_rank()
+    if global_rank == 0:
+        print(f"[{global_rank}]", *msg)

deepspeed/utils/exceptions.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+
+class DeprecatedException(Exception):
+    pass

deepspeed/utils/groups.py

@@ -0,0 +1,916 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# The file has been adapted from https://github.com/NVIDIA/Megatron-LM and retains the following license from the original file
+
+# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+ Support different forms of parallelism in DeepSpeed using multiple process groups.
+ Given that there are multiple scenarios and use-cases, this file is going to be updated
+ frequently. For now, the group creation needed for the training scenario is being implemented.
+ For inference and other new scenarios, the code will be either reused or added to this file.
+"""
+
+import os
+
+from deepspeed import comm as dist
+from deepspeed.utils import log_dist
+from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size, bwc_pipeline_parallel_world_size
+from deepspeed.utils.exceptions import DeprecatedException
+from deepspeed.accelerator import get_accelerator
+
+# Expert parallel group that the current rank belongs to.
+_EXPERT_PARALLEL_GROUP = {}
+# Mapping of expert parallel group to ranks
+_EXPERT_PARALLEL_GROUP_RANKS = {}
+# Expert data parallel group that the current rank belongs to.
+_EXPERT_DATA_PARALLEL_GROUP = {}
+# Mapping of expert data parallel group to ranks
+_EXPERT_DATA_PARALLEL_GROUP_RANKS = {}
+# dist world group needs to be cloned for some cases
+_WORLD_GROUP = None
+# ZeRO parameter  partitioning group that the current rank belongs to.
+_ZERO_PARAM_INTRA_PARALLEL_GROUP = None
+# global object to maintain mpu object if passed by a Megatron client
+mpu = None
+# global object that stores tensor parallel world size for experts
+expert_tensor_parallel_world_size = 1
+# All to All quantized graident communication groups
+_ALL_TO_ALL_GROUP = {}
+
+mesh_device = None
+
+_DEVICE_MESH_SPLIT_UNSUPPORTED = "No backend for the parent process group or its backend does not support splitting"
+
+
+# Deprecated groups initialize function.
+def initialize(ep_size=1, mpu=None):
+    """ Deprecated function. Retained to inform the users."""
+    raise DeprecatedException(
+        "Please do not use the groups.initialize() API as it is deprecated. Instead, pass the desired ep_size to deepspeed.moe.layer.MoE(..,ep_size,..)"
+    )
+
+
+def _ensure_divisibility(numerator, denominator):
+    """Ensure that numerator is divisible by the denominator."""
+    assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator)
+
+
+# ======== Start: Tensor Parallel Group Attributes ========
+
+# Intra-layer model parallel group that the current rank belongs to.
+_TENSOR_MODEL_PARALLEL_GROUP = None
+
+# Model parallel group (both intra- and pipeline) that the current rank belongs to.
+_MODEL_PARALLEL_GROUP = None
+# Data parallel group that the current rank belongs to.
+_DATA_PARALLEL_GROUP = None
+
+# These values enable us to change the mpu sizes on the fly.
+_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
+_MPU_TENSOR_MODEL_PARALLEL_RANK = None
+
+
+def _init_tp_groups_with_new_group(tensor_model_parallel_size=1, data_parallel_size=None):
+    """Initialize TP/DP groups with explicit rank lists.
+
+    This mirrors a 2D DeviceMesh shaped as (data_parallel, tensor_parallel),
+    while avoiding DeviceMesh's optimized split_group path.
+    """
+
+    global _DATA_PARALLEL_GROUP
+    global _MODEL_PARALLEL_GROUP
+    global _TENSOR_MODEL_PARALLEL_GROUP
+
+    world_size = dist.get_world_size()
+    _ensure_divisibility(world_size, tensor_model_parallel_size)
+
+    if data_parallel_size is None:
+        data_parallel_size = world_size // tensor_model_parallel_size
+    else:
+        assert data_parallel_size * tensor_model_parallel_size == world_size, (
+            f"data_parallel_size ({data_parallel_size}) * tensor_model_parallel_size "
+            f"({tensor_model_parallel_size}) must equal world_size ({world_size})")
+
+    rank = dist.get_rank()
+    data_parallel_group = None
+    tensor_model_parallel_group = None
+
+    for tensor_rank in range(tensor_model_parallel_size):
+        ranks = list(range(tensor_rank, world_size, tensor_model_parallel_size))
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            data_parallel_group = group
+
+    for data_rank in range(data_parallel_size):
+        start = data_rank * tensor_model_parallel_size
+        ranks = list(range(start, start + tensor_model_parallel_size))
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            tensor_model_parallel_group = group
+
+    assert data_parallel_group is not None, 'data parallel group is not initialized'
+    assert tensor_model_parallel_group is not None, 'tensor parallel group is not initialized'
+
+    _DATA_PARALLEL_GROUP = data_parallel_group
+    _TENSOR_MODEL_PARALLEL_GROUP = tensor_model_parallel_group
+    _MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP
+
+    return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
+
+
+def _init_tp_mesh_device(tensor_model_parallel_size=1, data_parallel_size=None):
+    """Initialize model data parallel groups."""
+
+    global _DATA_PARALLEL_GROUP
+    global _MODEL_PARALLEL_GROUP
+    global _TENSOR_MODEL_PARALLEL_GROUP
+
+    if _TENSOR_MODEL_PARALLEL_GROUP is not None:
+        return
+
+    if data_parallel_size is None:
+        data_parallel_size = dist.get_world_size() // tensor_model_parallel_size
+
+    if os.environ.get("TORCH_DISTRIBUTED_DEBUG", "").upper() == "DETAIL":
+        log_dist("TORCH_DISTRIBUTED_DEBUG=DETAIL detected; initializing TP mesh groups with new_group", ranks=[0])
+        return _init_tp_groups_with_new_group(tensor_model_parallel_size, data_parallel_size)
+
+    try:
+        mesh_device = dist.initialize_mesh_device((data_parallel_size, tensor_model_parallel_size),
+                                                  ("data_parallel", "tensor_parallel"))
+    except RuntimeError as exc:
+        if _DEVICE_MESH_SPLIT_UNSUPPORTED not in str(exc):
+            raise
+        log_dist("DeviceMesh process-group splitting is unsupported; falling back to new_group TP mesh groups",
+                 ranks=[0])
+        return _init_tp_groups_with_new_group(tensor_model_parallel_size, data_parallel_size)
+
+    _TENSOR_MODEL_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="tensor_parallel")
+    _DATA_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="data_parallel")
+
+    # They are always equal only in 2D (DP + TP) parallelism.
+    # _MODEL_PARALLEL_GROUP is assigned the same value as _TENSOR_MODEL_PARALLEL_GROUP
+    # to allow for future potential changes.
+    _MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP
+
+    return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
+
+
+def get_tensor_model_parallel_group():
+    """Get the tensor model parallel group the caller rank belongs to."""
+
+    assert _TENSOR_MODEL_PARALLEL_GROUP is not None, \
+        'intra_layer_model parallel group is not initialized'
+    return _TENSOR_MODEL_PARALLEL_GROUP
+
+
+def get_model_parallel_group():
+    """Get the model parallel group the caller rank belongs to."""
+
+    assert _MODEL_PARALLEL_GROUP is not None, \
+        'model parallel group is not initialized'
+    return _MODEL_PARALLEL_GROUP
+
+
+def get_data_parallel_group():
+    """Get the data parallel group the caller rank belongs to."""
+    assert _DATA_PARALLEL_GROUP is not None, \
+        'data parallel group is not initialized'
+    return _DATA_PARALLEL_GROUP
+
+
+def set_tensor_model_parallel_world_size(world_size):
+    """Set the tensor model parallel size"""
+    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size
+
+
+def get_tensor_model_parallel_world_size():
+    """Return world size for the tensor model parallel group."""
+    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
+        return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    return dist.get_world_size(group=get_tensor_model_parallel_group())
+
+
+def get_model_parallel_world_size():
+    return get_tensor_model_parallel_world_size()
+
+
+def set_tensor_model_parallel_rank(rank):
+    """Set tensor model parallel rank."""
+    global _MPU_TENSOR_MODEL_PARALLEL_RANK
+    _MPU_TENSOR_MODEL_PARALLEL_RANK = rank
+
+
+def get_tensor_model_parallel_rank():
+    """Return my rank for the tensor model parallel group."""
+    global _MPU_TENSOR_MODEL_PARALLEL_RANK
+    if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None:
+        return _MPU_TENSOR_MODEL_PARALLEL_RANK
+    return dist.get_rank(group=get_tensor_model_parallel_group())
+
+
+def get_model_parallel_rank():
+    return get_tensor_model_parallel_rank()
+
+
+def get_tensor_model_parallel_src_rank():
+    """Calculate the global rank corresponding to the first local rank
+    in the tensor model parallel group."""
+    global_rank = dist.get_rank()
+    local_world_size = get_tensor_model_parallel_world_size()
+    return (global_rank // local_world_size) * local_world_size
+
+
+def get_data_parallel_world_size():
+    """Return world size for the data parallel group."""
+    return dist.get_world_size(group=get_data_parallel_group())
+
+
+def get_data_parallel_rank():
+    """Return my rank for the data parallel group."""
+    return dist.get_rank(group=get_data_parallel_group())
+
+
+# ======== End: Tensor Parallel Group Attributes ========
+
+
+# Not currently used. Helper function to create a model (tensor) parallel group.
+def _create_model_parallel(model_parallel_size_):
+    """
+    Initialize model data parallel groups.
+
+    Arguments:
+        model_parallel_size: number of GPUs used to parallelize model.
+
+    Returns:
+        Tuple of data parallel group and model parallel group
+
+    Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
+    use 2 GPUs to parallelize the model. The present function will
+    create 4 model parallel groups and 2 data parallel groups as:
+        4 model parallel groups:
+            [g0, g1], [g2, g3], [g4, g5], [g6, g7]
+        2 data parallel groups:
+            [g0, g2, g4, g6], [g1, g3, g5, g7]
+    Note that for efficiency, the caller should make sure adjacent ranks
+    are on the same DGX box. For example if we are using 2 DGX-1 boxes
+    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
+    ranks 8 to 15 belong to the second box.
+    """
+    log_dist(f'Creating model parallel group with size {model_parallel_size_}', ranks=[0])
+    # Get world size and rank. Ensure some consistencies.
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    model_parallel_size = min(model_parallel_size_, world_size)
+    _ensure_divisibility(world_size, model_parallel_size)
+    rank = dist.get_rank()
+
+    _DATA_PARALLEL_GROUP = None
+    _MODEL_PARALLEL_GROUP = None
+    # Build the data parallel groups.
+    for i in range(model_parallel_size):
+        ranks = range(i, world_size, model_parallel_size)
+        group = dist.new_group(ranks)
+        if i == (rank % model_parallel_size):
+            _DATA_PARALLEL_GROUP = group
+
+    # Build the model parallel groups.
+    for i in range(world_size // model_parallel_size):
+        ranks = range(i * model_parallel_size, (i + 1) * model_parallel_size)
+        group = dist.new_group(ranks)
+        if i == (rank // model_parallel_size):
+            _MODEL_PARALLEL_GROUP = group
+
+    return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
+
+
+def _create_expert_and_data_parallel(expert_parallel_size_,
+                                     mp_size=None,
+                                     pp_size=None,
+                                     mp_mode="tp",
+                                     use_data_before_expert_parallel_=False,
+                                     folding_spec=None):
+    """Create expert and data parallel groups.
+
+    When mp_size is None or 1: legacy consecutive ordering (backward compatible).
+    When mp_size > 1 and folding_spec is not None: AutoEP+AutoTP folding tables.
+    When mp_size > 1 and mp_mode=="tp": TP-strided rank ordering.
+    When mp_size > 1 and mp_mode=="sp": consecutive rank ordering.
+
+    Note: Caller of this function is responsible to check if the groups already exist.
+
+    Example - E + D parallel (legacy path)
+    world_size = 16
+    expert_parallel_size = 2 # number of experts in same group
+    expert_data_parallel_group = [0,2,4,6,8,10,12,14], [1,3,5,7,9,11,13,15] - all reduce is only on MoE params
+    expert_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - no all reduce, but all to all
+    data_parallel_group = [0,1,...,15] - all reduce is only on non-MoE
+
+    Args:
+        expert_parallel_size_ (int): Expert parallel group size.
+        mp_size (int, optional): Model parallel size (TP or SP). None treated as 1.
+        pp_size (int, optional): Pipeline parallel size. None falls back to mpu.
+        mp_mode (str): "tp" for TP-strided ordering, "sp" for consecutive ordering.
+        use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology.
+        folding_spec: Optional AutoEP+AutoTP folding topology spec.
+    """
+    assert dist.is_initialized()
+
+    # Resolve parameters for backward compat
+    effective_mp_size = 1 if mp_size is None else mp_size
+
+    log_dist(f'Creating expert and data parallel groups with size {expert_parallel_size_}', ranks=[0])
+    world_size = dist.get_world_size()
+
+    # Resolve pp_size
+    if pp_size is not None:
+        pp_world_size = pp_size
+    else:
+        pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu)
+
+    rank = dist.get_rank()
+
+    pp_stride = world_size // pp_world_size
+    _ensure_divisibility(pp_stride, expert_parallel_size_)
+
+    group_name = f"ep_size_{expert_parallel_size_}"
+
+    global _EXPERT_DATA_PARALLEL_GROUP
+    global _EXPERT_DATA_PARALLEL_GROUP_RANKS
+    global _EXPERT_PARALLEL_GROUP
+    global _EXPERT_PARALLEL_GROUP_RANKS
+
+    # Legacy path: mp_size <= 1 (preserves exact original behavior)
+    if effective_mp_size <= 1:
+        ep_stride = pp_stride // expert_parallel_size_
+
+        # Build the expert data parallel groups.
+        # Only create group if it does not already exist
+        if group_name not in _EXPERT_DATA_PARALLEL_GROUP:
+            for pp_stage_start in range(0, world_size, pp_stride):
+                for i in range(expert_parallel_size_):
+                    if use_data_before_expert_parallel_:
+                        ranks = range(pp_stage_start + i * ep_stride, pp_stage_start + (i + 1) * ep_stride)
+                    else:
+                        ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, expert_parallel_size_)
+                    group = dist.new_group(ranks)
+                    log_dist(
+                        f'Creating expert data parallel process group named {group_name} with ranks: {list(ranks)}',
+                        [0])
+                    if rank in ranks:
+                        _EXPERT_DATA_PARALLEL_GROUP[group_name] = group
+                        _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks
+
+        # Build the expert parallel groups.
+        # Only create group if it does not already exist
+        if group_name not in _EXPERT_PARALLEL_GROUP:
+            if use_data_before_expert_parallel_:
+                for pp_stage_start in range(0, world_size, pp_stride):
+                    for i in range(ep_stride):
+                        ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, ep_stride)
+                        group = dist.new_group(ranks)
+                        log_dist(
+                            f'creating expert parallel process group named {group_name} '
+                            f'with ranks: {list(ranks)}', [0])
+                        if rank in ranks:
+                            _EXPERT_PARALLEL_GROUP[group_name] = group
+                            _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks
+            else:
+                for i in range(world_size // expert_parallel_size_):
+                    ranks = range(i * expert_parallel_size_, (i + 1) * expert_parallel_size_)
+                    group = dist.new_group(ranks)
+                    log_dist(
+                        f'creating expert parallel process group named {group_name} '
+                        f'with ranks: {list(ranks)}', [0])
+                    if rank in ranks:
+                        _EXPERT_PARALLEL_GROUP[group_name] = group
+                        _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks
+        return
+
+    # New path: mp_size > 1
+    if use_data_before_expert_parallel_:
+        raise NotImplementedError("use_data_before_expert_parallel_ is not supported with mp_size > 1")
+
+    if group_name in _EXPERT_PARALLEL_GROUP:
+        if folding_spec is not None:
+            from deepspeed.module_inject.auto_ep_folding import assert_group_matches_spec
+            assert_group_matches_spec(
+                {
+                    "ep": [_EXPERT_PARALLEL_GROUP_RANKS[group_name]],
+                    "edp": [_EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name]],
+                },
+                folding_spec,
+            )
+        return  # Already created
+
+    folding_tables = None
+    if folding_spec is not None:
+        from deepspeed.module_inject.auto_ep_folding import expected_folding_group_tables
+        folding_tables = expected_folding_group_tables(folding_spec)
+
+    for pp_stage_start in range(0, world_size, pp_stride):
+        stage_ranks = list(range(pp_stage_start, pp_stage_start + pp_stride))
+        stage_rank_set = set(stage_ranks)
+
+        if folding_tables is not None:
+            ep_groups_list = [list(group) for group in folding_tables.ep_groups if set(group).issubset(stage_rank_set)]
+            edp_groups_list = [
+                list(group) for group in folding_tables.edp_groups if set(group).issubset(stage_rank_set)
+            ]
+        else:
+            # Preserve the existing TP-strided native MoE topology when no
+            # folding spec was provided by the AutoEP+AutoTP path.
+            if mp_mode == "tp" and effective_mp_size > 1:
+                num_tp_groups = len(stage_ranks) // effective_mp_size
+                ordered_stage_ranks = []
+                for dp_lane in range(effective_mp_size):
+                    for tp_group_idx in range(num_tp_groups):
+                        ordered_stage_ranks.append(stage_ranks[tp_group_idx * effective_mp_size + dp_lane])
+            else:
+                ordered_stage_ranks = stage_ranks
+
+            num_ep_groups = len(ordered_stage_ranks) // expert_parallel_size_
+            ep_groups_list = [
+                ordered_stage_ranks[g * expert_parallel_size_:(g + 1) * expert_parallel_size_]
+                for g in range(num_ep_groups)
+            ]
+            edp_groups_list = [[ep_groups_list[g][pos] for g in range(num_ep_groups)]
+                               for pos in range(expert_parallel_size_)]
+
+        for ep_ranks in ep_groups_list:
+            group = dist.new_group(ep_ranks)
+            log_dist(f'creating expert parallel process group named {group_name} with ranks: {ep_ranks}', [0])
+            if rank in ep_ranks:
+                _EXPERT_PARALLEL_GROUP[group_name] = group
+                _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ep_ranks
+
+        for edp_ranks in edp_groups_list:
+            group = dist.new_group(edp_ranks)
+            log_dist(f'Creating expert data parallel process group named {group_name} with ranks: {edp_ranks}', [0])
+            if rank in edp_ranks:
+                _EXPERT_DATA_PARALLEL_GROUP[group_name] = group
+                _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = edp_ranks
+
+
+def _get_expert_parallel_ranks(world_size,
+                               tensor_parallel_size_,
+                               expert_parallel_size_,
+                               pipeline_parallel_size_=1,
+                               use_data_before_expert_parallel_=False):
+    """Generate expert parallel and expert data parallel group ranks list.
+
+        Example - E + M + D parallel
+        world_size = 16
+        model_degree = 2
+        expert_degree = 4 # number of experts in same group
+        mp_group = [0, 1], [2,3], [4,5] ...
+        data_parallel_group =[0,2,4,6,8,10, 12,14],                 [1,3,5,7,9,11,13,15]
+        expert_parallel_group = [0,2,4,6], [8,10,12,14]             [1,3,5,7], [9,11,13,15]
+        expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14],    [1,9],[3,11],[5,13],[7,15]
+
+    Args:
+        world_size (int): Distributed world size.
+        tensor_parallel_size_ (int): Tensor parallel group size.
+        expert_parallel_size_ (int): Expert parallel group size.
+        pipeline_parallel_size_ (int): Pipeline parallel group size
+        use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology
+    Returns:
+        Expert parallel group ranks and Expert data parallel group ranks list.
+    """
+    _ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_)
+    dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_)
+    _ensure_divisibility(dp_world_size, expert_parallel_size_)
+
+    # Generate data parallel groups
+    data_parallel_groups = []
+    dp_group_size = tensor_parallel_size_
+    pp_stride = world_size // pipeline_parallel_size_
+
+    if use_data_before_expert_parallel_:
+        dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_
+        for pp_stage_start in range(0, world_size, pp_stride):
+            pp_stage_next = pp_stage_start + pp_stride
+            for i in range(dp_group_size):
+                data_parallel_groups.append(list())
+                for ds in range(dp_stride):
+                    # [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30]
+                    # [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31]
+                    data_parallel_groups[-1].extend(
+                        list(
+                            range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next,
+                                  dp_stride * tensor_parallel_size_)))
+    else:
+        for pp_stage_start in range(0, world_size, pp_stride):
+            pp_stage_next = pp_stage_start + pp_stride
+            for i in range(dp_group_size):
+                data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size)))
+
+    expert_parallel_groups = []
+    expert_data_parallel_groups = []
+    for dp_ranks in data_parallel_groups:
+        # partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14]
+        part_ep_groups = []
+        for i in range(0, dp_world_size, expert_parallel_size_):
+            part_ep_groups.append(dp_ranks[i:i + expert_parallel_size_])
+        expert_parallel_groups.extend(part_ep_groups)
+
+        # zip part_ep_groups get expert data parallel ranks, e.g [0,8],[2,10],[4,12],[6,14]
+        for expert_dp_ranks in zip(*part_ep_groups):
+            expert_data_parallel_groups.append(list(expert_dp_ranks))
+
+    return expert_parallel_groups, expert_data_parallel_groups
+
+
+def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False):
+    """
+        Create expert and data parallel groups based on MPU (model parallel) group.
+
+        Note: Caller of this function is responsible to check if the groups already exist.
+
+        Example - E + M + D parallel
+        world_size = 16
+        model_degree = 2
+        expert_degree = 4 # number of experts in same group
+        mp_group = [0, 1], [2,3], [4,5] ...
+        data_parallel_group =[0,2,4,6,8,10, 12,14],                 [1,3,5,7,9,11,13,15]
+        expert_parallel_group = [0,2,4,6], [8,10,12,14]             [1,3,5,7], [9,11,13,15]
+        expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14],    [1,9],[3,11],[5,13],[7,15]
+    """
+    assert dist.is_initialized(), "dist is not initialized"
+    tensor_parallel_size_ = bwc_tensor_model_parallel_world_size(mpu)
+
+    global expert_tensor_parallel_world_size
+    expert_tensor_parallel_world_size = tensor_parallel_size_
+
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    dp_world_size = _get_data_parallel_world_size()
+    pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu)
+
+    _ensure_divisibility(world_size, tensor_parallel_size_)
+    _ensure_divisibility(dp_world_size, expert_parallel_size_)
+
+    log_dist(
+        f"Creating deepspeed groups with model parallel size {tensor_parallel_size_}, "
+        f"pipeline parallel size {pp_world_size}, expert parallel size {expert_parallel_size_}, "
+        f"world size {world_size}, dp world size {dp_world_size}", [0])
+
+    global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP
+    global _EXPERT_PARALLEL_GROUP_RANKS, _EXPERT_DATA_PARALLEL_GROUP_RANKS
+
+    group_name = f"ep_size_{expert_parallel_size_}"
+
+    # Only create groups if they don't already exist
+    # Need to check conditions outside the group creation loop because of the way torch.dist group creation works
+    if group_name not in _EXPERT_DATA_PARALLEL_GROUP and group_name not in _EXPERT_PARALLEL_GROUP:
+        expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks(
+            world_size, tensor_parallel_size_, expert_parallel_size_, pp_world_size, use_data_before_expert_parallel_)
+        for ranks in expert_parallel_groups:
+            group = dist.new_group(ranks)
+            if rank in list(ranks):
+                _EXPERT_PARALLEL_GROUP[group_name] = group
+                _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks
+
+        for ranks in expert_data_parallel_groups:
+            group = dist.new_group(ranks)
+            if rank in list(ranks):
+                _EXPERT_DATA_PARALLEL_GROUP[group_name] = group
+                _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks
+
+
+def _get_max_expert_size():
+    """Get the maximum ep_size from all the created groups."""
+    assert _EXPERT_PARALLEL_GROUP is not None, "Warning! Process group not initialized"
+    keylist = []
+    for key in _EXPERT_PARALLEL_GROUP.keys():
+        # index 2 is ep_size in the group name: ep_size_<ep_size>
+        index = 2
+        keylist.append(int(key.split('_')[index]))
+    return max(keylist) if len(keylist) > 0 else None
+
+
+def _get_max_expert_size_name():
+    """Get the name of the group with max. ep_size"""
+    return f'ep_size_{_get_max_expert_size()}'
+
+
+def _get_max_expert_parallel_group():
+    """Get the max expert parallel size."""
+    return _get_expert_parallel_group(_get_max_expert_size_name())
+
+
+def _get_expert_parallel_group(group_name):
+    """Get the expert parallel group the caller rank belongs to."""
+    assert group_name in _EXPERT_PARALLEL_GROUP, \
+        'expert parallel group is not initialized'
+    return _EXPERT_PARALLEL_GROUP[group_name]
+
+
+def _get_expert_parallel_group_ranks(group_name):
+    """Get the ranks of the expert parallel group the caller rank belongs to."""
+    assert group_name in _EXPERT_PARALLEL_GROUP_RANKS, \
+        'expert parallel group is not initialized'
+    return _EXPERT_PARALLEL_GROUP_RANKS[group_name]
+
+
+def _get_expert_parallel_group_dict():
+    """Get the expert parallel group dict."""
+    return _EXPERT_PARALLEL_GROUP
+
+
+def _get_expert_data_parallel_group(group_name):
+    """Get the expert data parallel group the caller rank belongs to."""
+    assert group_name in _EXPERT_DATA_PARALLEL_GROUP, \
+        'expert data parallel group is not initialized'
+    return _EXPERT_DATA_PARALLEL_GROUP[group_name]
+
+
+def _get_expert_data_parallel_group_ranks(group_name):
+    """Get the ranks of the expert data parallel group the caller rank belongs to."""
+    assert group_name in _EXPERT_DATA_PARALLEL_GROUP_RANKS, \
+        'expert data parallel group is not initialized'
+    return _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name]
+
+
+def _get_expert_data_parallel_group_dict():
+    """Get the expert data parallel group dict."""
+    return _EXPERT_DATA_PARALLEL_GROUP
+
+
+def _clone_world_group():
+    """Create a clone of the world group
+        Note: We need to clone the dist world group because we
+        use dist.get_global_rank() utility function in DeepSpeed at many places.
+        As that function does not work on dist.group.WORLD, we
+        need to keep a clone of it.
+    """
+    assert dist.is_initialized(), "dist is not initialized"
+    global _WORLD_GROUP
+    if _WORLD_GROUP is None:
+        # If not cloned already, clone the world group
+        _WORLD_GROUP = dist.new_group(ranks=range(dist.get_world_size()))
+    return _WORLD_GROUP
+
+
+def _get_local_all_to_all_group():
+    assert dist.is_initialized(), 'dist is not initialized'
+    global _ALL_TO_ALL_GROUP
+    device_per_node = get_accelerator().device_count()
+    num_local = dist.get_world_size() // device_per_node
+    if num_local == 0 and dist.get_world_size() > 0:
+        assert dist.get_world_size() >= 1, 'num_gpus must >=1, cannot initialize All-To-All'
+        cur_rank = []
+        for i in range(dist.get_world_size()):
+            cur_rank.append(i)
+        _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=cur_rank)
+    elif num_local == 1:
+        assert dist.get_world_size(
+        ) == device_per_node, 'num_gpus not equal to device per node, cannot initialize All-To-All'
+        _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=[i for i in range(device_per_node)])
+    else:
+        assert dist.get_world_size() > device_per_node, 'num_nodes<2 cannot initialize All-To-All'
+        for i in range(num_local):
+            local_rank = [j + device_per_node * i for j in range(device_per_node)]
+            _ALL_TO_ALL_GROUP[f"local_{i}"] = dist.new_group(ranks=local_rank)
+
+        for i in range(device_per_node):
+            cur_rank = []
+            for j in range(num_local):
+                cur_rank.append(i + j * device_per_node)
+            _ALL_TO_ALL_GROUP[f"global_{i}"] = dist.new_group(ranks=cur_rank)
+    return _ALL_TO_ALL_GROUP
+
+
+def _get_data_parallel_group():
+    """Get the data parallel group the caller rank belongs to."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    global mpu
+    if mesh_device is not None:
+        return mesh_device.get_group(mesh_dim="data_parallel")
+    if mpu is not None:
+        if hasattr(mpu, 'initialize_sequence_parallel'):
+            return None
+        else:
+            return mpu.get_data_parallel_group()
+
+    # Return the clone of dist world group
+    return _clone_world_group()
+
+
+def _get_data_parallel_group_ranks():
+    """Get the ranks of data parallel group the caller rank belongs to."""
+    assert dist.is_initialized(), \
+        'dist is not initialized'
+    global mpu
+    if mpu is not None:
+        return mpu.get_data_parallel_group_ranks()
+    # Return all ranks
+    return range(dist.get_world_size())
+
+
+def _get_broadcast_src_rank():
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_global_rank(_get_sequence_data_parallel_group(), 0)
+
+
+def _get_expert_broadcast_src_rank(group_name):
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_global_rank(_get_expert_data_parallel_group(group_name), 0)
+
+
+def _get_expert_parallel_world_size(group_name):
+    """Return world size for the expert parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_world_size(group=_get_expert_parallel_group(group_name))
+
+
+def _get_expert_data_parallel_world_size(group_name):
+    """Return world size for the expert data parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_world_size(group=_get_expert_data_parallel_group(group_name))
+
+
+def _get_expert_parallel_rank(group_name):
+    """Return my rank for the expert parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_rank(group=_get_expert_parallel_group(group_name))
+
+
+def _get_expert_parallel_src_rank(group_name):
+    """Calculate the global rank corresponding to a local rank zero
+    in the expert parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    global_rank = dist.get_rank()
+    local_world_size = _get_expert_parallel_world_size(group_name)
+    return (global_rank // local_world_size) * local_world_size
+
+
+def _get_expert_data_parallel_rank(group_name):
+    """Return my rank for the expert data parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_rank(group=_get_expert_data_parallel_group(group_name))
+
+
+def _get_data_parallel_world_size():
+    """Return world size for the data parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    if mesh_device is not None:
+        return dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel"))
+    global mpu
+    if mpu is not None:
+        if hasattr(mpu, 'initialize_sequence_parallel'):
+            return None
+        else:
+            return mpu.get_data_parallel_world_size()
+    return dist.get_world_size(group=_get_data_parallel_group())
+
+
+def _get_model_parallel_world_size():
+    """Return world size for the model parallel group."""
+    global mpu
+    if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'):
+        return 1
+    return mpu.get_model_parallel_world_size()
+
+
+def _get_data_parallel_rank():
+    """Return my rank for the data parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    return dist.get_rank(group=_get_data_parallel_group())
+
+
+def _get_sequence_parallel_world_size():
+    """Return world size for the model parallel group."""
+    assert dist.is_initialized(), 'dist is not initialized'
+    global mpu
+    if mesh_device is not None:
+        return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel"))
+    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'):
+        return mpu.get_sequence_parallel_world_size()
+    return 1
+
+
+def _get_sequence_parallel_rank():
+    """Return my rank for the sequence parallel group."""
+    global mpu
+    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'):
+        return mpu.get_sequence_parallel_rank()
+    if mesh_device is not None:
+        return dist.get_rank(mesh_device.get_group(mesh_dim="sequence_parallel"))
+    return 0
+
+
+def _get_sequence_parallel_group():
+    global mpu
+    if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'):
+        if mesh_device is None:
+            raise KeyError("No sequence parallel group found")
+        return mesh_device.get_group(mesh_dim="sequence_parallel")
+    return mpu.get_sequence_parallel_group()
+
+
+def _get_sequence_data_parallel_world_size():
+    """Return world size for the model parallel group."""
+    global mpu
+    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'):
+        return mpu.get_sequence_data_parallel_world_size()
+    return _get_data_parallel_world_size()
+
+
+def _get_sequence_data_parallel_rank():
+    """Return my rank for the data parallel group."""
+    global mpu
+    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'):
+        return mpu.get_sequence_data_parallel_rank()
+    return _get_data_parallel_rank()
+
+
+def _get_sequence_data_parallel_group():
+    global mpu
+    # When sequence parallelism is enabled, the process group for zero sharding and
+    # gradient allreduce must be across both dimensions of data and sequence parallelism.
+    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'):
+        return mpu.get_sequence_data_parallel_group()
+    return _get_data_parallel_group()
+
+
+def _get_expert_model_parallel_world_size():
+    global expert_tensor_parallel_world_size
+    return expert_tensor_parallel_world_size
+
+
+def _create_zero_param_parallel_group(group_size):
+    """
+        Create parameter partitioning group within ZeRO data parallel groups.
+
+        Example - ZP + D parallel
+        world_size = 16
+        zero_hpz_partition_size = 2 # number of ranks with replicated params (dual partitioning)
+        zero_param_intra_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - segmented (subgroup) with rep partition
+        data_parallel_group = [0,1,...,15] - all reduce is on ZeRO model
+    """
+    assert dist.is_initialized()
+    global _ZERO_PARAM_INTRA_PARALLEL_GROUP
+    # Only create group if it does not already exist
+    assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is None, \
+        'ZeRO parameter intra parallel group is already initialized'
+
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+
+    zero_param_parallel_size_ = min(group_size, world_size)
+    _ensure_divisibility(world_size, zero_param_parallel_size_)
+
+    # Build the ZeRO param intra parallel groups.
+    for i in range(world_size // zero_param_parallel_size_):
+        ranks = range(i * zero_param_parallel_size_, (i + 1) * zero_param_parallel_size_)
+        group = dist.new_group(ranks)
+        if i == (rank // zero_param_parallel_size_):
+            _ZERO_PARAM_INTRA_PARALLEL_GROUP = group
+
+
+def _get_zero_param_intra_parallel_group():
+    """Get the ZeRO parameter partitioning intra parallel group the caller rank belongs to."""
+    #assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is not None, \
+    #    'ZeRO parameter partitioning group is not initialized'
+    #TODO: Add warning
+    return _ZERO_PARAM_INTRA_PARALLEL_GROUP
+
+
+def _zero_param_parallel_is_initialized():
+    """Check if ZeRO data parallel with parameter partititioning groups are initialized."""
+    ###TODO: assert that MPU is not set
+    if _ZERO_PARAM_INTRA_PARALLEL_GROUP is None and _DATA_PARALLEL_GROUP is None:
+        return False
+
+
+def _get_zero_param_intra_parallel_rank_in_mygroup():
+    """Return my rank for the ZeRO parameter inter parallel group."""
+    return dist.get_rank(group=_get_zero_param_intra_parallel_group())
+
+
+def _get_zero_param_intra_parallel_group_world_size():
+    """Return world size for the ZeRO parameter parallel group."""
+    return dist.get_world_size(group=_get_zero_param_intra_parallel_group())
+
+
+def _get_zero_param_intra_parallel_group_ranks():
+    """Return all ranks for the ZeRO parameter intra parallel group."""
+    return dist.get_all_ranks_from_group(group=_get_zero_param_intra_parallel_group())

deepspeed/utils/init_on_device.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from typing import Callable
+from torch import Tensor
+from packaging import version as pkg_version
+
+
+class OnDevice(object):
+    """
+    Create modules/tensors w. specific devices and dtypes. Examples:
+
+    Create MyModule which consists of many different sub-modules and parameters. In this case we can create
+    MyModule as a collection of 'meta' tensors by passing `device='meta'` or we can create the module _directly_
+    on a CUDA device by passing `device=f'cuda:{local_rank}'` (where `local_rank` is the local GPU id.
+
+    with OnDevice(dtype=torch.float16, device='meta'):
+        model = MyModel()
+
+    with OnDevice(dtype=torch.float16, device=f'cuda:{local_rank}'):
+        model = MyModel()
+
+    """
+
+    _orig_torch_empty = torch.empty
+    _orig_torch_zeros = torch.zeros
+    _orig_torch_ones = torch.ones
+    _orig_torch_full = torch.full
+
+    def __init__(self, dtype, device="meta", enabled=True):
+        self.dtype = dtype
+        self.enabled = enabled
+        self.device = device
+
+        if device == "meta":
+            if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__):
+                raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+")
+
+    def fp_tensor_constructor(self, fn: Callable, target_fp_dtype: torch.dtype) -> Callable:
+
+        def wrapped_fn(*args, **kwargs) -> Tensor:
+            if kwargs.get("device", None) is None:
+                kwargs['device'] = self.device
+            tensor: Tensor = fn(*args, **kwargs)
+            if tensor.is_floating_point():
+                tensor = tensor.to(target_fp_dtype)
+            return tensor
+
+        return wrapped_fn
+
+    def get_new_tensor_fn_for_dtype(self, dtype: torch.dtype) -> Callable:
+
+        def new_tensor(cls, *args) -> Tensor:
+            tensor = OnDevice._orig_torch_empty(0, device=self.device).new_empty(*args)
+            if tensor.is_floating_point():
+                tensor = tensor.to(dtype)
+            return tensor
+
+        return new_tensor
+
+    def __enter__(self):
+        if not self.enabled:
+            return
+        torch.Tensor.__old_new__ = torch.Tensor.__new__
+        torch.Tensor.__new__ = self.get_new_tensor_fn_for_dtype(self.dtype)
+        torch.empty = self.fp_tensor_constructor(self._orig_torch_empty, self.dtype)
+        torch.zeros = self.fp_tensor_constructor(self._orig_torch_zeros, self.dtype)
+        torch.ones = self.fp_tensor_constructor(self._orig_torch_ones, self.dtype)
+        torch.full = self.fp_tensor_constructor(self._orig_torch_full, self.dtype)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if not self.enabled:
+            return
+        torch.Tensor.__new__ = torch.Tensor.__old_new__
+        torch.empty = self._orig_torch_empty
+        torch.zeros = self._orig_torch_zeros
+        torch.ones = self._orig_torch_ones
+        torch.full = self._orig_torch_full

deepspeed/utils/logging.py

@@ -0,0 +1,231 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import functools
+import logging
+import sys
+import os
+import torch
+from deepspeed.utils.torch import required_torch_version
+
+log_levels = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
+
+
+class LoggerFactory:
+
+    @staticmethod
+    def create_logger(name=None, level=logging.WARNING):
+        """create a logger
+
+        Args:
+            name (str): name of the logger
+            level: level of logger
+
+        Raises:
+            ValueError is name is None
+        """
+
+        if name is None:
+            raise ValueError("name for logger cannot be None")
+
+        formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] "
+                                      "[%(filename)s:%(lineno)d:%(funcName)s] %(message)s")
+
+        logger_ = logging.getLogger(name)
+        logger_.setLevel(level)
+        logger_.propagate = False
+        ch = logging.StreamHandler(stream=sys.stdout)
+        ch.setLevel(level)
+        ch.setFormatter(formatter)
+        logger_.addHandler(ch)
+        if required_torch_version(min_version=2.6) and os.getenv("DISABLE_LOGS_WHILE_COMPILING", "0") == "1":
+            excluded_set = {
+                item.strip()
+                for item in os.getenv("LOGGER_METHODS_TO_EXCLUDE_FROM_DISABLE", "").split(",")
+            }
+            ignore_set = {'info', 'debug', 'error', 'warning', 'critical', 'exception', 'isEnabledFor'} - excluded_set
+            for method in ignore_set:
+                original_logger = getattr(logger_, method)
+                torch._dynamo.config.ignore_logger_methods.add(original_logger)
+        return logger_
+
+
+logger = LoggerFactory.create_logger(name="DeepSpeed", level=logging.WARNING)
+
+
+@functools.lru_cache(None)
+def warning_once(*args, **kwargs):
+    """
+    This method is identical to `logger.warning()`, but will emit the warning with the same message only once
+
+    Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache.
+    The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to
+    another type of cache that includes the caller frame information in the hashing function.
+    """
+    logger.warning(*args, **kwargs)
+
+
+logger.warning_once = warning_once
+
+
+def print_configuration(args, name):
+    logger.info("{}:".format(name))
+    for arg in sorted(vars(args)):
+        dots = "." * (29 - len(arg))
+        logger.info("  {} {} {}".format(arg, dots, getattr(args, arg)))
+
+
+def get_dist_msg(message, ranks=None):
+    from deepspeed import comm as dist
+    """Return a message with rank prefix when one of following conditions is met:
+
+      + not dist.is_initialized()
+      + dist.get_rank() in ranks if ranks is not None or ranks = [-1]
+
+    If neither is met, `None` is returned.
+
+    Example: "hello" => "[Rank 0] hello"
+
+    Args:
+        message (str)
+        ranks (list)
+    """
+    should_log = not dist.is_initialized()
+    ranks = ranks or []
+    my_rank = dist.get_rank() if dist.is_initialized() else -1
+    if ranks and not should_log:
+        should_log = ranks[0] == -1
+        should_log = should_log or (my_rank in set(ranks))
+    if should_log:
+        return "[Rank {}] {}".format(my_rank, message)
+    else:
+        return None
+
+
+def log_dist(message, ranks=None, level=logging.INFO):
+    """Log message when get_dist_msg() deems it should be logged, see its docstring for details.
+
+    Args:
+        message (str)
+        ranks (list)
+        level (int)
+    """
+    final_message = get_dist_msg(message, ranks)
+    if final_message is not None:
+        logger.log(level, final_message)
+
+
+def print_dist(message, ranks=None):
+    """print message when get_dist_msg() deems it should be logged, see its docstring for details.
+
+    Use this function instead of `log_dist` when the log level shouldn't impact whether the message should be printed or not.
+
+    Args:
+        message (str)
+        ranks (list)
+    """
+    final_message = get_dist_msg(message, ranks)
+    if final_message is not None:
+        print(final_message)
+
+
+@functools.lru_cache(None)
+def _log_dist_once_cached(message, ranks_key, level):
+    ranks_arg = list(ranks_key) if ranks_key is not None else None
+    log_dist(message, ranks=ranks_arg, level=level)
+
+
+def log_dist_once(message, ranks=None, level=logging.INFO):
+    # Identical to `log_dist`, but will emit each unique message only once per process.
+    # ranks is a list which is unhashable, so convert to tuple for caching
+    ranks_key = tuple(ranks) if ranks is not None else None
+    _log_dist_once_cached(message, ranks_key, level)
+
+
+logger.log_dist_once = log_dist_once
+
+
+def print_json_dist(message, ranks=None, path=None):
+    from deepspeed import comm as dist
+    """Print message when one of following condition meets
+
+    + not dist.is_initialized()
+    + dist.get_rank() in ranks if ranks is not None or ranks = [-1]
+
+    Args:
+        message (str)
+        ranks (list)
+        path (str)
+
+    """
+    should_log = not dist.is_initialized()
+    ranks = ranks or []
+    my_rank = dist.get_rank() if dist.is_initialized() else -1
+    if ranks and not should_log:
+        should_log = ranks[0] == -1
+        should_log = should_log or (my_rank in set(ranks))
+    if should_log:
+        message['rank'] = my_rank
+        import json
+        with open(path, 'w') as outfile:
+            json.dump(message, outfile)
+            os.fsync(outfile)
+
+
+def get_log_level_from_string(log_level_str):
+    """converts a log level string into its numerical equivalent. e.g. "info" => `logging.INFO`
+    """
+    log_level_str = log_level_str.lower()
+    if log_level_str not in log_levels:
+        raise ValueError(
+            f"{log_level_str} is not one of the valid logging levels. Valid log levels are {log_levels.keys()}.")
+    return log_levels[log_level_str]
+
+
+def set_log_level_from_string(log_level_str, custom_logger=None):
+    """Sets a log level in the passed `logger` and its handlers from string. e.g. "info" => `logging.INFO`
+
+    Args:
+        log_level_str: one of 'debug', 'info', 'warning', 'error', 'critical'
+        custom_logger: if `None` will use the default `logger` object
+    """
+    log_level = get_log_level_from_string(log_level_str)
+    if custom_logger is None:
+        custom_logger = logger
+    custom_logger.setLevel(log_level)
+    for handler in custom_logger.handlers:
+        handler.setLevel(log_level)
+
+
+def get_current_level():
+    """
+    Return logger's current log level
+    """
+    return logger.getEffectiveLevel()
+
+
+def should_log_le(max_log_level_str):
+    """
+    Args:
+        max_log_level_str: maximum log level as a string
+
+    Returns ``True`` if the current log_level is less or equal to the specified log level. Otherwise ``False``.
+
+    Example:
+
+        ``should_log_le("info")`` will return ``True`` if the current log level is either ``logging.INFO`` or ``logging.DEBUG``
+    """
+
+    if not isinstance(max_log_level_str, str):
+        raise ValueError(f"{max_log_level_str} is not a string")
+
+    max_log_level = get_log_level_from_string(max_log_level_str)
+    return get_current_level() <= max_log_level

deepspeed/utils/mixed_precision_linkage.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import types
+from deepspeed.utils import get_full_hp_param, get_full_hp_grad, get_hp_fragment_mapping
+from deepspeed.utils import set_full_hp_param, set_full_hp_grad
+
+
+def link_hp_params(lp_param_list, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload,
+                   param_group_index, partition_start, partition_size, dp_group):
+    local_lp_param_and_offset = _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group)
+
+    for lp_param, lp_start in local_lp_param_and_offset:
+        lp_param._hp_mapping = get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict,
+                                                       offload_gradient_dict, use_offload, param_group_index,
+                                                       partition_start, partition_size)
+
+
+def lazy_init_hp_params_optimizer_state(lp_param_list, flat_hp_partition, optimizer_state):
+    for lp in lp_param_list:
+        if lp._hp_mapping is not None:
+            lp._hp_mapping.set_optim_state_fragment(flat_hp_partition, optimizer_state[flat_hp_partition])
+
+
+def _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group):
+    current_offset = 0
+    param_and_offset_list = []
+    partition_end = partition_start + partition_size
+    index_in_param_group = 0
+    for i, lp_param in enumerate(lp_param_list):
+        lp_param._hp_mapping = None
+        lp_param._dp_group = dp_group
+        lp_param.get_full_hp_param = types.MethodType(get_full_hp_param, lp_param)
+        lp_param.get_full_hp_grad = types.MethodType(get_full_hp_grad, lp_param)
+        lp_param.set_full_hp_param = types.MethodType(set_full_hp_param, lp_param)
+        lp_param.set_full_hp_grad = types.MethodType(set_full_hp_grad, lp_param)
+
+        # lp_param overlaps with partition if both are true
+        # 1) current_offset < partition_end,
+        # 2) current_offset + lp_param.numel() >= partition_start
+        lp_param_end = current_offset + lp_param.numel()
+        if current_offset < partition_end and lp_param_end > partition_start:
+            param_and_offset_list.append((lp_param, current_offset))
+            lp_param._index_in_param_group = index_in_param_group
+            # Indices for params in this partition/GPU
+            index_in_param_group += 1
+        current_offset += lp_param.numel()
+
+    return param_and_offset_list

deepspeed/utils/numa.py

@@ -0,0 +1,205 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+# return a list of list for cores to numa mapping
+# [
+#     [ cores for numa 0 ]
+#     [ cores belong to numa 1 ]
+#     ...
+# ]
+
+import os
+import psutil
+import shutil
+import subprocess
+
+
+# return a list of list for cores to numa mapping
+# [
+#     [ cores for numa 0 ]
+#     [ cores belong to numa 1 ]
+#     ...
+# ]
+def get_numa_cores():
+    ret = []
+    try:
+        output = subprocess.check_output(['numactl', '--hardware']).decode("utf-8")
+    except Exception:
+        return []
+    lines = output.split('\n')
+    for line in lines:
+        if line.startswith('available:'):
+            num_numas = int(line.split(' ')[1])
+            break
+    for numa in range(num_numas):
+        for line in lines:
+            if line.startswith(f'node {numa} cpus:'):
+                cores = line.split(' ')[3:]
+                ret.append([int(core) for core in cores])
+    return ret
+
+
+def check_for_numactl_pkg():
+    libs = dict(
+        dpkg=["-l", "numactl", "apt"],
+        pacman=["-Q", "numactl", "pacman"],
+        rpm=["-q", "numactl", "yum"],
+    )
+
+    found = False
+    for pkgmgr, data in libs.items():
+        flag, lib, tool = data
+        path = shutil.which(pkgmgr)
+        if path is not None:
+            cmd = [pkgmgr, flag, lib]
+            result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            if result.wait() == 0:
+                found = True
+            else:
+                print(f"please install the {lib} package with {tool}")
+            break
+    return found
+
+
+def parse_range(rng):
+    try:
+        value = int(rng)
+        return range(value, value + 1)
+    except ValueError:
+        # value is not a single number
+        parts = rng.split('-')
+        if len(parts) != 2:
+            raise ValueError("Bad range: '%s', range must be either a number or two number separated by dash" %
+                             (rng, ))
+        start = int(parts[0])
+        end = int(parts[1])
+        if start > end:
+            raise ValueError("Bad range: '%s', range end must larger than or equal to start" % (rng, ))
+        return range(start, end + 1)
+
+
+# parse comma and dash separated range list into list
+# i.e. "0,2-4,6" --> [0, 2, 3, 4, 6]
+# rules:
+# 1. Range list number be comma separated, each item are either a single number,
+#    or a range marked by two numbers (both number are included in the range)
+# 2. Sub ranges must be in ascend order and not overlap with each other
+# 3. No space in the range expression
+def parse_range_list(range_str):
+    number_list = []
+    last = -1
+    range_list = range_str.split(',')
+    for sub_range in range_list:
+        sub_number_list = parse_range(sub_range)
+        if sub_number_list[0] <= last:
+            raise ValueError(
+                "Bad range: '%s', sub ranges must not overlap with each other and should be in ascend order" %
+                (range_str, ))
+        last = sub_number_list[-1]
+        number_list.extend(sub_number_list)
+    return number_list
+
+
+def get_numactl_cmd(bind_core_list, num_local_procs, local_rank):
+    numactl_cmd = []
+    check_for_numactl_pkg()
+    if 'KMP_AFFINITY' in os.environ.keys():
+        raise ValueError("Environment variable KMP_AFFINITY conflicts with numactl "
+                         "because it interfere with how many CPU cores numactl can set. "
+                         "Unset KMP_AFFINITY before launching deepspeed.\n\n"
+                         "\t$ unset KMP_AFFINITY\n"
+                         "\t$ deepspeed <deepspeed command parameters>")
+    if bind_core_list is not None:
+        core_list = parse_range_list(bind_core_list)
+        total_cores = len(core_list)
+    else:
+        total_cores = psutil.cpu_count(logical=False)
+        core_list = range(total_cores)
+    cores_per_rank = total_cores // num_local_procs
+    assert cores_per_rank >= 1, "At least one core needs to be assigned to each rank"
+    core_list_for_rank = core_list[cores_per_rank * local_rank:cores_per_rank * (local_rank + 1)]
+    numactl_cmd.append("numactl")
+
+    # check if all cores belong to same numa, if true, bind process to that numa domain with -m parameter
+    numa_cores = get_numa_cores()
+    num_numas = len(numa_cores)
+
+    numa_mode = "normal"
+
+    non_empty_numa_list = []
+    empty_numa_list = []
+    previous_numa_cores = []
+    numa_node_list = []
+    numa_node_list_list = []
+    for i in range(num_numas):
+        # look for empty numa which is HBM numa
+        if numa_cores[i] == []:
+            empty_numa_list.append(i)
+        else:
+            non_empty_numa_list.append(i)
+
+            # check for fakenuma
+            if numa_cores[i] == previous_numa_cores:
+                if numa_node_list == []:
+                    #first duplication, add previous node into list
+                    numa_node_list.append(i - 1)
+                numa_node_list.append(i)
+            else:
+                if numa_node_list != []:
+                    numa_node_list_list.append(numa_node_list)
+                    numa_node_list = []
+        previous_numa_cores = numa_cores[i]
+    if numa_node_list != []:
+        numa_node_list_list.append(numa_node_list)
+
+    if empty_numa_list != [] and len(empty_numa_list) == len(non_empty_numa_list):
+        numa_mode = "flat_hbm"
+        numa_dict = dict(zip(non_empty_numa_list, empty_numa_list))
+    elif numa_node_list_list != []:
+        numa_mode = "fake"
+
+    if numa_mode == "normal":
+        for i in range(num_numas):
+            if set(core_list_for_rank) <= set(numa_cores[i]):
+                numactl_cmd.append("-m")
+                numactl_cmd.append(f"{i}")
+                break
+    elif numa_mode == "flat_hbm":
+        for i in range(num_numas):
+            if set(core_list_for_rank) <= set(numa_cores[i]):
+                numactl_cmd.append("-p")
+                numactl_cmd.append(f"{numa_dict[i]}")
+                break
+    elif numa_mode == "fake":
+        for i in range(num_numas):
+            if set(core_list_for_rank) <= set(numa_cores[i]):
+                for nodes in numa_node_list_list:
+                    if i in nodes:
+                        numactl_cmd.append("-m")
+                        numactl_cmd.append(f"{','.join(map(str, nodes))}")
+                        break
+                # the following construct break the outer loop if inner loop breaks
+                else:
+                    continue
+                break
+
+    numactl_cmd.append("-C")
+    last_core = core_list_for_rank[0]
+    first_core = last_core
+    core_list_str = f"{last_core}"
+    for core_id in core_list_for_rank[1:]:
+        if core_id == last_core + 1:
+            last_core = core_id
+            continue
+        else:
+            if first_core == last_core:
+                core_list_str = f"{core_list_str},{core_id}"
+            else:
+                core_list_str = f"{core_list_str}-{last_core},{core_id}"
+            first_core = core_id
+            last_core = core_id
+    if first_core != last_core:
+        core_list_str = f"{core_list_str}-{last_core}"
+    numactl_cmd.append(f"{core_list_str}")
+    return cores_per_rank, numactl_cmd

deepspeed/utils/nvtx.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.compiler import is_compiling
+
+enable_nvtx = True
+DEEPSPEED_NVTX_DOMAIN = "DeepSpeed"
+
+
+def _range_push(accelerator, msg):
+    if getattr(accelerator, "supports_nvtx_domain", False):
+        return accelerator.range_push(msg, domain=DEEPSPEED_NVTX_DOMAIN)
+    return accelerator.range_push(msg)
+
+
+def _range_pop(accelerator):
+    if getattr(accelerator, "supports_nvtx_domain", False):
+        return accelerator.range_pop(domain=DEEPSPEED_NVTX_DOMAIN)
+    return accelerator.range_pop()
+
+
+def instrument_w_nvtx(func):
+    """Decorator that records an NVTX range for the duration of the function call.
+       Skips NVTX instrumentation when torch.compile is active to avoid graph breaks.
+    """
+
+    def wrapped_fn(*args, **kwargs):
+        if enable_nvtx and not is_compiling():
+            _range_push(get_accelerator(), func.__qualname__)
+        ret_val = func(*args, **kwargs)
+        if enable_nvtx and not is_compiling():
+            _range_pop(get_accelerator())
+        return ret_val
+
+    return wrapped_fn

deepspeed/utils/static_cache.py

@@ -0,0 +1,230 @@
+# Copyright (c) DeepSpeed Team
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""CUDA-graph-compatible static KV cache for hybrid engine rollout.
+
+Derived from HuggingFace transformers ``StaticCache`` / ``StaticLayer``, but
+with a critical difference: the write position is supplied externally via a
+shared tensor instead of an internal ``cumulative_length`` counter.
+
+Why this matters
+----------------
+Transformers' ``StaticLayer.update()`` maintains its own ``cumulative_length``
+tensor that advances on every call.  During CUDA graph capture the captured
+forward "freezes" this counter at whatever value it had at capture time.
+On replay the counter does *not* advance, so subsequent KV writes go to the
+wrong positions and the model silently produces incorrect logits.
+
+Our ``DeepSpeedStaticCache`` instead reads the write position from a shared
+tensor (``write_position``) that the caller updates in-place before each graph
+replay.  Because ``write_position`` is a real tensor at a fixed address, CUDA
+graph replays read the current value each time.
+
+The caller (HybridEngineRollout) must call ``cache.set_write_position(pos)``
+before each replay, where ``pos`` is a scalar ``torch.long`` tensor on the
+correct device.
+"""
+
+import torch
+
+
+class DeepSpeedStaticLayer:
+    """A single layer's static KV cache whose write position is externally set.
+
+    Parameters
+    ----------
+    max_cache_len : int
+        Maximum number of tokens the cache can hold (last dim size).
+    """
+
+    is_compileable = True
+    is_sliding = False
+
+    def __init__(self, max_cache_len: int):
+        self.max_cache_len = max_cache_len
+        self.keys: torch.Tensor | None = None
+        self.values: torch.Tensor | None = None
+        self.is_initialized = False
+        self._write_position: torch.Tensor | None = None
+
+    def set_write_position(self, pos: torch.Tensor):
+        self._write_position = pos
+
+    def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None:
+        self.dtype = key_states.dtype
+        self.device = key_states.device
+        max_batch_size, num_heads = key_states.shape[:2]
+        self.max_batch_size = max_batch_size
+        self.num_heads = num_heads
+        self.k_head_dim = key_states.shape[-1]
+        self.v_head_dim = value_states.shape[-1]
+
+        self.keys = torch.zeros(
+            (max_batch_size, num_heads, self.max_cache_len, self.k_head_dim),
+            dtype=self.dtype,
+            device=self.device,
+        )
+        self.values = torch.zeros(
+            (max_batch_size, num_heads, self.max_cache_len, self.v_head_dim),
+            dtype=self.dtype,
+            device=self.device,
+        )
+        torch._dynamo.mark_static_address(self.keys)
+        torch._dynamo.mark_static_address(self.values)
+        self.is_initialized = True
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if not self.is_initialized:
+            self.lazy_initialization(key_states, value_states)
+
+        kv_length = key_states.shape[-2]
+
+        if self._write_position is not None:
+            cache_position = torch.arange(kv_length, device=self.device) + self._write_position
+        else:
+            cache_position = torch.arange(kv_length, device=self.device)
+
+        try:
+            self.keys.index_copy_(2, cache_position, key_states)
+            self.values.index_copy_(2, cache_position, value_states)
+        except NotImplementedError:
+            self.keys[:, :, cache_position] = key_states
+            self.values[:, :, cache_position] = value_states
+
+        return self.keys, self.values
+
+    def get_mask_sizes(self, query_length: int) -> tuple[int, int]:
+        return self.max_cache_len, 0
+
+    def get_seq_length(self) -> int:
+        if not self.is_initialized:
+            return 0
+        if self._write_position is not None:
+            return self._write_position + 1
+        return 0
+
+    def get_max_cache_shape(self) -> int:
+        return self.max_cache_len
+
+    def reset(self) -> None:
+        if self.is_initialized:
+            self.keys.zero_()
+            self.values.zero_()
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        if self.is_initialized:
+            self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))
+            self.values = self.values.index_select(0, beam_idx.to(self.values.device))
+
+
+class DeepSpeedStaticCache:
+    """CUDA-graph-compatible static KV cache.
+
+    Drop-in replacement for ``transformers.StaticCache`` in the graph-capture
+    decode path of ``HybridEngineRollout``.  All layers share a single
+    ``write_position`` tensor that the caller updates before each graph replay.
+
+    Parameters
+    ----------
+    config : PreTrainedConfig
+        HuggingFace model config (used to determine number of layers and head
+        dimensions).
+    batch_size : int
+        Batch size for eager initialization.
+    max_cache_len : int
+        Maximum sequence length (prompt + generated tokens).
+    device : torch.device | int | str | None
+        Device for eager initialization.
+    dtype : torch.dtype | None
+        Dtype for eager initialization.
+    """
+
+    def __init__(
+        self,
+        config,
+        batch_size: int = 1,
+        max_cache_len: int = 4096,
+        device=None,
+        dtype=None,
+    ):
+        self.config = config
+        text_config = getattr(config, "text_config", config)
+        num_layers = getattr(text_config, "num_hidden_layers", 1)
+        self._layers = [DeepSpeedStaticLayer(max_cache_len) for _ in range(num_layers)]
+        self._max_cache_len = max_cache_len
+        self._write_position: torch.Tensor | None = None
+
+        if dtype is not None and device is not None and batch_size > 0:
+            num_heads = getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))
+            head_dim = getattr(text_config, "hidden_size", 1) // getattr(text_config, "num_attention_heads", 1)
+            self.early_initialization(batch_size, num_heads, head_dim, dtype, device)
+
+    @property
+    def layers(self):
+        return self._layers
+
+    def set_write_position(self, pos: torch.Tensor):
+        """Set the write position shared by all layers.
+
+        Must be called before each graph replay with the decode step position
+        as a scalar ``torch.long`` tensor on the correct device.  The tensor is
+        stored by reference so subsequent in-place updates (e.g.
+        ``pos.fill_(new_val)``) are immediately visible to all layers.
+        """
+        self._write_position = pos
+        for layer in self._layers:
+            layer.set_write_position(pos)
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        *args,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if layer_idx >= len(self._layers):
+            raise IndexError(f"layer_idx {layer_idx} out of range (cache has {len(self._layers)} layers)")
+        return self._layers[layer_idx].update(key_states, value_states, *args, **kwargs)
+
+    def early_initialization(
+        self,
+        batch_size: int,
+        num_heads: int,
+        head_dim: int,
+        dtype: torch.dtype,
+        device,
+    ):
+        for layer in self._layers:
+            fake_k = torch.zeros((batch_size, num_heads, 0, head_dim), dtype=dtype, device=device)
+            fake_v = torch.zeros((batch_size, num_heads, 0, head_dim), dtype=dtype, device=device)
+            layer.lazy_initialization(fake_k, fake_v)
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        if layer_idx >= len(self._layers):
+            return 0
+        return self._layers[layer_idx].get_seq_length()
+
+    def get_max_cache_shape(self, layer_idx: int = 0) -> int:
+        if layer_idx >= len(self._layers):
+            return self._max_cache_len
+        return self._layers[layer_idx].get_max_cache_shape()
+
+    def get_mask_sizes(self, query_length: int, layer_idx: int = 0) -> tuple[int, int]:
+        if layer_idx >= len(self._layers):
+            return self._max_cache_len, 0
+        return self._layers[layer_idx].get_mask_sizes(query_length)
+
+    def reset(self):
+        for layer in self._layers:
+            layer.reset()
+
+    def __len__(self):
+        return len(self._layers)

deepspeed/utils/tensor_fragment.py

@@ -0,0 +1,481 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from dataclasses import dataclass
+from deepspeed import comm as dist
+from typing import Dict, List, Callable
+
+
+@dataclass
+class fragment_address:
+    numel: int
+    start: int
+
+
+@dataclass
+class tensor_fragment:
+    lp_fragment: torch.Tensor
+    lp_fragment_address: fragment_address
+    hp_fragment: torch.Tensor
+    hp_fragment_address: fragment_address
+    gradient_dict: Dict
+    offload_gradient_dict: Dict
+    use_offload: bool
+    param_group_index: int
+    optim_fragment: Dict = None
+
+    def update_hp(self):
+        self.hp_fragment.data.copy_(self.lp_fragment.data)
+
+    def update_lp(self):
+        self.lp_fragment.data.copy_(self.hp_fragment.data)
+
+    def get_optim_state_fragment(self, key):
+        if key in self.optim_fragment:
+            return self.optim_fragment[key]
+        else:
+            raise ValueError(f'{key} not found in optimizer state fragment')
+
+    def set_optim_state_fragment(self, flat_hp_partition, optim_fragment):
+        self.optim_fragment = {
+            key: value.narrow(0, self.hp_fragment_address.start, self.hp_fragment_address.numel)
+            for key, value in optim_fragment.items()
+            if torch.is_tensor(value) and value.shape == flat_hp_partition.shape
+        }
+
+    def get_hp_fragment_address(self):
+        return self.hp_fragment_address
+
+    def get_optim_state_keys(self):
+        return list(self.optim_fragment.keys())
+
+    def get_hp_fragment(self, optim_state_key=None):
+        if optim_state_key is None:
+            return self.hp_fragment
+        return self.get_optim_state_fragment(optim_state_key)
+
+    def get_lp_grad_fragment(self, index_in_param_group):
+        if self.use_offload:
+            gradient_dict = self.offload_gradient_dict
+        else:
+            gradient_dict = self.gradient_dict
+
+        if self.param_group_index not in gradient_dict or gradient_dict[self.param_group_index] is None:
+            raise ValueError("Gradients are only available immediately after backward and before engine step")
+
+        return gradient_dict[self.param_group_index][index_in_param_group]
+
+
+def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys):
+    for key in opt_keys:
+        hp_param = flat_hp_tensor
+        buffer = torch.zeros_like(hp_param)
+
+        for lp in lp_tensors:
+            if lp._hp_mapping is not None:
+                hp_fragment_address = lp._hp_mapping.get_hp_fragment_address()
+                hp_fragment = buffer.narrow(0, hp_fragment_address.start, hp_fragment_address.numel)
+                hp_fragment.data.copy_(lp._hp_mapping.get_hp_fragment(optim_state_key=key).data)
+                lp._hp_mapping.hp_fragment = hp_fragment
+
+        optim_state[hp_param][key] = buffer
+
+
+def get_full_hp_param(self, optim_state_key=None):
+    reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten()
+    if self._hp_mapping is not None:
+        lp_frag_address = self._hp_mapping.lp_fragment_address
+        reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel)
+        hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key)
+        reduce_fragment.data.copy_(hp_fragment.data)
+    dist.all_reduce(reduce_buffer, group=self._dp_group)
+    return reduce_buffer.reshape_as(self)
+
+
+def set_full_hp_param(self, value, optim_state_key=None):
+    if self._hp_mapping is not None:
+        lp_frag_address = self._hp_mapping.lp_fragment_address
+        value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel)
+        hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key)
+        hp_fragment.data.copy_(value_fragment.data)
+
+
+def get_full_hp_grad(self):
+    reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten()
+    if self._hp_mapping is not None:
+        lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group)
+        hp_grad_fragment = lp_grad_fragment.to(torch.float32).flatten()
+
+        lp_frag_address = self._hp_mapping.lp_fragment_address
+        reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel)
+
+        if self.view(-1).shape == hp_grad_fragment.shape:
+            reduce_buffer.data.copy_(hp_grad_fragment.data)
+        else:
+            reduce_fragment.data.copy_(hp_grad_fragment.data)
+
+    dist.all_reduce(reduce_buffer, group=self._dp_group)
+    return reduce_buffer.reshape_as(self)
+
+
+def set_full_hp_grad(self, value):
+    if self._hp_mapping is not None:
+        lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group)
+        lp_frag_address = self._hp_mapping.lp_fragment_address
+        value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel)
+        lp_grad_fragment.data.copy_(value_fragment.data.reshape_as(lp_grad_fragment.data))
+        if hasattr(self, '_zero_optimizer'):
+            self._zero_optimizer.update_offload_overflow_tracker(value)
+
+
+def safe_get_full_fp32_param(param):
+    """Assemble and return the fp32 parameter of a low-precision (e.g., fp16) parameter.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+
+        Returns:
+            Union[torch.Tensor, None]: A tensor on accelerator device
+    """
+    # ZeRO stage 3 param
+    if hasattr(param, 'ds_id'):
+        return param._z3_optimizer.get_full_hp_param(param)
+
+    # ZeRO stage 1, 2, and bf16_optimizer params
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_param()
+    return None
+
+
+def safe_set_full_fp32_param(param, value):
+    """Update the partitioned fp32 parameter of a low-precision (e.g., fp16) parameter.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+            value (``torch.Tensor``): New value
+    """
+    # ZeRO stage 3 param
+    if hasattr(param, 'ds_id'):
+        param._z3_optimizer.set_full_hp_param(value, param)
+
+    # ZeRO stage 1, 2, and bf16_optimizer params
+    if hasattr(param, '_hp_mapping'):
+        param.set_full_hp_param(value)
+
+
+def safe_get_full_optimizer_state(param, optim_state_key):
+    """Assemble and return the fp32 optimizer state of a low-precision (e.g., fp16) parameter.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+            optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer)
+
+        Returns:
+            Union[torch.Tensor, None]: A tensor on accelerator device
+"""
+    # ZeRO stage 3 param
+    if hasattr(param, 'ds_id'):
+        return param._z3_optimizer.get_full_hp_param(param, optim_state_key)
+
+    # ZeRO stage 1, 2, and bf16_optimizer params
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_param(optim_state_key)
+    return None
+
+
+def safe_set_full_optimizer_state(param, value, optim_state_key):
+    """Update the partitioned fp32 optimizer state of a low-precision (e.g., fp16) parameter.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+            value (``torch.Tensor``): New value
+            optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer)
+    """
+    # ZeRO stage 3 param
+    if hasattr(param, 'ds_id'):
+        param._z3_optimizer.set_full_hp_param(value, param, optim_state_key)
+
+    # ZeRO stage 1, 2, and bf16_optimizer params
+    if hasattr(param, '_hp_mapping'):
+        param.set_full_hp_param(value, optim_state_key)
+
+
+# TODO: Figure out the correct return dtype
+def safe_get_full_grad(param):
+    """
+        Assemble and return the fp32 gradient of a low-precision (e.g., fp16) parameter.
+        The return data type is that used for gradient accumulation. This is usually the param data type,
+        but could also be different (e.g., bf16 param training with fp32 gradient accumulation).
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+
+        Returns:
+            Union[torch.Tensor, None]: A tensor on accelerator device
+    """
+    if param.grad is not None:
+        return param.grad
+
+    # ZeRO stage 3 param
+    if hasattr(param, 'ds_id'):
+        return param._z3_optimizer.get_fp32_grad_for_param(param)
+
+    # ZeRO stage 1, 2, and bf16_optimizer params
+    if hasattr(param, '_hp_mapping'):
+        return param.get_full_hp_grad()
+
+    return None
+
+
+def safe_set_full_grad(param, value):
+    """
+        Update the partitioned gradient of a low-precision (e.g., fp16) parameter.
+        To avoid precision issues, the update value should have the data type of
+        gradient accumulation.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+            value (``torch.Tensor``): The un-partitioned new gradient value.
+    """
+    if param.grad is not None:
+        param.grad.copy_(value)
+    elif hasattr(param, 'ds_id'):
+        # ZeRO stage 3 param
+        param._z3_optimizer.set_fp32_grad_for_param(value, param)
+    elif hasattr(param, '_hp_mapping'):
+        # ZeRO stage 1, 2, and bf16_optimizer params
+        param.set_full_hp_grad(value)
+
+
+### Local API  START ###
+def safe_get_local_grad(param):
+    """
+        Get the local gradient partition of a ZeRO-3 partitioned parameter.
+        The return data type is that used for gradient accumulation. This is usually the param data type,
+        but could also be different (e.g., bf16 param training with fp32 gradient accumulation).
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+
+        Returns:
+            Union[torch.Tensor, None]: A tensor on accelerator device
+    """
+    assert hasattr(param, 'ds_id'), 'This API is only defined for ZeRO-3 partitioned parameters'
+    return param._z3_optimizer.get_local_fp32_grad_for_param(param)
+
+
+def safe_set_local_grad(param, value):
+    """
+        Update the local gradient partition of a ZeRO-3 partitioned parameter.
+        To avoid precision issues, the update value should have the data type of
+        gradient accumulation.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter.
+            value (``torch.Tensor``): New value of local gradient partition.
+    """
+    assert hasattr(param, 'ds_id'), 'This API is only defined for ZeRO-3 partitioned parameters'
+    param._z3_optimizer.set_local_grad_for_param(value, param)
+
+
+def safe_get_local_fp32_param(param):
+    """Get the local partition of a ZeRO-3 partitioned parameter in fp32 precision.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter.
+
+        Returns:
+            Union[torch.Tensor, None]: A tensor on accelerator device
+    """
+    assert hasattr(param, 'ds_id'), 'This API is only defined for ZeRO-3 partitioned parameters'
+    return param._z3_optimizer.get_local_fp32_param(param)
+
+
+def safe_get_local_optimizer_state(param, optim_state_key):
+    """Get the local optimizer state partition of ZeRO-3 partitioned parameter in fp32 precision.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter
+            optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer)
+
+        Returns:
+            Union[torch.Tensor, None]: A tensor on accelerator device
+    """
+    assert hasattr(param, 'ds_id'), 'This API is only defined for ZeRO-3 partitioned parameters'
+    return param._z3_optimizer.get_local_fp32_param(param, optim_state_key)
+
+
+def safe_set_local_optimizer_state(param, value, optim_state_key):
+    """Update the local optimizer state partition of a ZeRO-3 partitioned parameter.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter.
+            value (``torch.Tensor``): New value of local optimizer state partition.
+            optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer).
+    """
+    assert hasattr(param, 'ds_id'), 'This API is only defined for ZeRO-3 partitioned parameters'
+    param._z3_optimizer.set_local_hp_param(value, param, optim_state_key)
+
+
+def safe_set_local_fp32_param(param, value):
+    """Update the local partition of ZeRO-3 partitioned parameter.
+
+        Args:
+            param (``torch.nn.Parameter``): A model parameter.
+            value (``torch.Tensor``): New value of local parameter partition.
+    """
+    assert hasattr(param, 'ds_id'), 'This API is only defined for ZeRO-3 partitioned parameters'
+    param._z3_optimizer.set_local_hp_param(value, param)
+
+
+### Local API  END ###
+
+
+### VECTORIZED API  BEGIN ###
+def safe_update_full_grad_vectorized(param_list: List[torch.nn.Parameter], update_func: Callable):
+    """
+        Vectorized update of the partitioned gradients of a list of low-precision (e.g., fp16) parameters.
+        To avoid precision issues, the update value should have the data type of
+        gradient accumulation.
+
+        Args:
+            param_list (``List[torch.nn.Parameter]``): List of model parameters
+            update_func (``torch.Tensor``): A function that takes current full gradient value and returns new one.
+    """
+    partitioned_grad_params = []
+    for p in param_list:
+        if p.grad is not None:
+            p.grad.copy_(update_func(p.grad, p))
+        elif p.requires_grad:
+            partitioned_grad_params.append(p)
+
+    if not partitioned_grad_params:
+        return
+
+    if hasattr(partitioned_grad_params[0], 'ds_id'):
+        # ZeRO stage 3 param
+        partitioned_grad_params[0]._z3_optimizer.update_fp32_grad_for_param_vectorized(
+            update_func, partitioned_grad_params)
+    elif hasattr(partitioned_grad_params[0], '_hp_mapping'):
+        # ZeRO stage 1, 2, and bf16_optimizer params
+        for p in partitioned_grad_params:
+            old_grad = safe_get_full_grad(p)
+            new_grad = update_func(old_grad, p)
+            p.set_full_hp_grad(new_grad)
+
+
+### VECTORIZED API  END ###
+
+
+def get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload,
+                            param_group_index, partition_start, partition_size):
+    lp_end = lp_param.numel() + lp_start
+    hp_start = partition_start
+    hp_end = partition_start + partition_size
+
+    fragment_start = max(lp_start, hp_start)
+    fragment_end = min(lp_end, hp_end)
+    assert fragment_start < fragment_end, \
+        f'fragment start {fragment_start} should be < fragment_end {fragment_end}'
+
+    fragment_numel = fragment_end - fragment_start
+    hp_frag_address = fragment_address(start=fragment_start - hp_start, numel=fragment_numel)
+    hp_fragment_tensor = flat_hp_partition.narrow(0, hp_frag_address.start, hp_frag_address.numel)
+
+    lp_frag_address = fragment_address(start=fragment_start - lp_start, numel=fragment_numel)
+    lp_fragment_tensor = lp_param.flatten().narrow(0, lp_frag_address.start, lp_frag_address.numel)
+
+    return tensor_fragment(lp_fragment=lp_fragment_tensor,
+                           lp_fragment_address=lp_frag_address,
+                           hp_fragment=hp_fragment_tensor,
+                           hp_fragment_address=hp_frag_address,
+                           gradient_dict=gradient_dict,
+                           offload_gradient_dict=offload_gradient_dict,
+                           use_offload=use_offload,
+                           param_group_index=param_group_index)
+
+
+'''
+Logic for lp_param to hp_param mapping
+
+lp      lp0 lp1 lp2         lp3  lp4            <-------  indices/names
+lp      [  ][  ][          ][   ][         ]    <-------- tensors
+flat_lp [                                  ]     <-------- flat lp params
+flat_hp            [                 ]   <------------------ flat hp partition on current rank
+full_hp [                                        ] <------- full flat hp params
+
+
+lp2
+ full numel = 16
+ lp_frag
+   numel = 12
+   frag_start = 3
+   frag_end  = 15
+ hp_frag
+    numel = 12
+    frag_start = 0
+    frag_end = 11
+
+ hp_frag.copy_(lp_frag)
+
+
+lp3:
+  full numel = 4
+  lp_frag
+     numel = 4
+     start = 0
+     end = 3
+  hp_frag
+     numel = 4
+     start = 12
+     end = 15
+
+
+lp4:
+   full numel = 12
+   lp_frag
+     numel = 4
+     start = 0
+     end = 3
+  hp_frag
+     numel = 4
+     start = 16
+     end = 19
+
+
+
+Visual depiction of above
+lp              {         }
+flat_lp [                                ]
+flat_hp            (                 )
+
+
+flat_lp [       {  (      }          )   ]
+                lx  hx   ly          hy
+                    ly-hx
+
+
+lp                             {       }
+flat_lp [                                ]
+flat_hp            (                 )
+
+
+flat_lp [          (            {     ) }  ]
+                   hx           lx   hy ly
+                                   hy-lx
+
+lp                        {   }
+flat_lp [                                ]
+flat_hp            (                 )
+
+
+flat_lp [          (       {   }      )   ]
+                   hx      lx  ly    hy
+                             ly-lx
+
+lp -> (lx, hy)
+flat_hp -> (hx, hy)
+'''

deepspeed/utils/timer.py

@@ -0,0 +1,313 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import time
+from numpy import mean
+from deepspeed.utils.logging import print_dist
+from deepspeed.accelerator import get_accelerator
+
+FORWARD_MICRO_TIMER = 'fwd_microstep'
+FORWARD_GLOBAL_TIMER = 'fwd'
+BACKWARD_MICRO_TIMER = 'bwd_microstep'
+BACKWARD_GLOBAL_TIMER = 'bwd'
+BACKWARD_INNER_MICRO_TIMER = 'bwd_inner_microstep'
+BACKWARD_INNER_GLOBAL_TIMER = 'bwd_inner'
+BACKWARD_REDUCE_MICRO_TIMER = 'bwd_allreduce_microstep'
+BACKWARD_REDUCE_GLOBAL_TIMER = 'bwd_allreduce'
+STEP_MICRO_TIMER = 'step_microstep'
+STEP_GLOBAL_TIMER = 'step'
+TIME_EPSILON = 1e-6
+
+try:
+    import psutil
+
+    PSUTILS_INSTALLED = True
+except ImportError:
+    PSUTILS_INSTALLED = False
+    pass
+
+
+class CudaEventTimer(object):
+
+    def __init__(self, start_event: get_accelerator().Event, end_event: get_accelerator().Event):
+        self.start_event = start_event
+        self.end_event = end_event
+
+    def get_elapsed_msec(self):
+        get_accelerator().current_stream().wait_event(self.end_event)
+        self.end_event.synchronize()
+        return self.start_event.elapsed_time(self.end_event)
+
+
+class SynchronizedWallClockTimer:
+    """Group of timers. Borrowed from Nvidia Megatron code"""
+
+    class Timer:
+        """Timer."""
+
+        def __init__(self, name):
+            self.name_ = name
+            self.started_ = False
+            self.event_timers = []
+            self.use_host_timer = get_accelerator().use_host_timers()
+            self.start_event = None
+            self.elapsed_records = None
+            self.start_time = 0.0
+            self.end_time = 0.0
+
+        def start(self):
+            """Start the timer."""
+            assert not self.started_, f"{self.name_} timer has already been started"
+            if self.use_host_timer:
+                self.start_time = time.time()
+            else:
+                event_class = get_accelerator().Event
+                self.start_event = event_class(enable_timing=True)
+                self.start_event.record()
+            self.started_ = True
+
+        def stop(self, reset=False, record=False):
+            """Stop the timer."""
+            assert self.started_, "timer is not started"
+            event_class = get_accelerator().Event
+            if self.use_host_timer:
+                self.end_time = time.time()
+                self.event_timers.append(self.end_time - self.start_time)
+            else:
+                event_class = get_accelerator().Event
+                end_event = event_class(enable_timing=True)
+                end_event.record()
+                self.event_timers.append(CudaEventTimer(self.start_event, end_event))
+                self.start_event = None
+            self.started_ = False
+
+        def _get_elapsed_msec(self):
+            if self.use_host_timer:
+                self.elapsed_records = [et * 1000.0 for et in self.event_timers]
+            else:
+                self.elapsed_records = [et.get_elapsed_msec() for et in self.event_timers]
+            return sum(self.elapsed_records)
+
+        def reset(self):
+            """Reset timer."""
+            self.started_ = False
+            self.start_event = None
+            self.elapsed_records = None
+            self.event_timers.clear()
+
+        def elapsed(self, reset=True):
+            """Calculate the elapsed time."""
+            started_ = self.started_
+            # If the timing in progress, end it first.
+            if self.started_:
+                self.stop()
+            # Get the elapsed time.
+            elapsed_ = self._get_elapsed_msec()
+            # Reset the elapsed time
+            if reset:
+                self.reset()
+            # If timing was in progress, set it back.
+            if started_:
+                self.start()
+            return elapsed_
+
+        def mean(self):
+            self.elapsed(reset=False)
+            return trim_mean(self.elapsed_records, 0.1)
+
+    def __init__(self):
+        self.timers = {}
+
+    def get_timers(self):
+        return self.timers
+
+    def __call__(self, name):
+        if name not in self.timers:
+            self.timers[name] = self.Timer(name)
+        return self.timers[name]
+
+    @staticmethod
+    def memory_usage():
+        alloc = "mem_allocated: {:.4f} GB".format(get_accelerator().memory_allocated() / (1024 * 1024 * 1024))
+        max_alloc = "max_mem_allocated: {:.4f} GB".format(get_accelerator().max_memory_allocated() /
+                                                          (1024 * 1024 * 1024))
+        cache = "cache_allocated: {:.4f} GB".format(get_accelerator().memory_cached() / (1024 * 1024 * 1024))
+        max_cache = "max_cache_allocated: {:.4f} GB".format(get_accelerator().max_memory_cached() /
+                                                            (1024 * 1024 * 1024))
+        return " | {} | {} | {} | {}".format(alloc, max_alloc, cache, max_cache)
+
+    def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None):
+        """Log a group of timers."""
+        assert normalizer > 0.0
+        string = "time (ms)"
+        for name in names:
+            if name in self.timers:
+                elapsed_time = (self.timers[name].elapsed(reset=reset) / normalizer)
+                string += " | {}: {:.2f}".format(name, elapsed_time)
+
+        # timers logging should be independent of the global log level it's already conditional on wall_clock_breakdown being True, so using use_logger=False will always print the stats
+        print_dist(string, ranks=ranks or [0])
+
+    def get_mean(self, names, normalizer=1.0, reset=True):
+        """Get the mean of a group of timers."""
+        assert normalizer > 0.0
+        means = {}
+        for name in names:
+            if name in self.timers:
+                elapsed_time = (self.timers[name].mean() * 1000.0 / normalizer)
+                means[name] = elapsed_time
+        return means
+
+
+class NoopTimer:
+
+    class Timer:
+
+        def start(self):
+            ...
+
+        def reset(self):
+            ...
+
+        def stop(self, **kwargs):
+            ...
+
+        def elapsed(self, **kwargs):
+            return 0
+
+        def mean(self):
+            return 0
+
+    def __init__(self):
+        self.timer = self.Timer()
+
+    def __call__(self, name):
+        return self.timer
+
+    def get_timers(self):
+        return {}
+
+    def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None):
+        ...
+
+    def get_mean(self, names, normalizer=1.0, reset=True):
+        ...
+
+
+class ThroughputTimer:
+
+    def __init__(self, config, batch_size, start_step=2, steps_per_output=None, monitor_memory=False, logging_fn=None):
+        from deepspeed.utils import logger
+        self.config = config
+        self.start_time = 0
+        self.end_time = 0
+        self.started = False
+        self.batch_size = 1 if batch_size is None else batch_size
+        self.start_step = start_step
+        self.epoch_count = 0
+        self.micro_step_count = 0
+        self.global_step_count = 0
+        self.total_elapsed_time = 0
+        self.step_elapsed_time = 0
+        self.steps_per_output = steps_per_output
+        self.monitor_memory = monitor_memory
+        self.logging = logging_fn
+        if self.logging is None:
+            self.logging = logger.info
+        self.initialized = False
+
+        if self.monitor_memory and not PSUTILS_INSTALLED:
+            raise ImportError("Unable to import 'psutils', please install package")
+
+    def update_epoch_count(self):
+        self.epoch_count += 1
+        self.micro_step_count = 0
+
+    def _init_timer(self):
+        self.initialized = True
+
+    def start(self):
+        if not self.config.enabled:
+            return
+        self._init_timer()
+        self.started = True
+        if self.global_step_count >= self.start_step:
+            if self.config.synchronized:
+                get_accelerator().synchronize()
+            self.start_time = time.time()
+
+    def _is_report_boundary(self):
+        if self.steps_per_output is None:
+            return False
+        return self.global_step_count % self.steps_per_output == 0
+
+    def stop(self, global_step=False, report_speed=True):
+        if not self.config.enabled or not self.started:
+            return
+        self.started = False
+        self.micro_step_count += 1
+        if global_step:
+            self.global_step_count += 1
+
+        if self.start_time > 0:
+            if self.config.synchronized:
+                get_accelerator().synchronize()
+            self.end_time = time.time()
+            duration = self.end_time - self.start_time
+            self.total_elapsed_time += duration
+            self.step_elapsed_time += duration
+
+            if global_step:
+                if report_speed and self._is_report_boundary():
+                    self.logging(
+                        "epoch={}/micro_step={}/global_step={}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, "
+                        "MemAllocated={}GB, MaxMemAllocated={}GB".format(
+                            self.epoch_count,
+                            self.micro_step_count,
+                            self.global_step_count,
+                            self.avg_samples_per_sec(),
+                            self.batch_size / (self.step_elapsed_time + TIME_EPSILON),
+                            round(get_accelerator().memory_allocated() / 1024**3, 2),
+                            round(get_accelerator().max_memory_allocated() / 1024**3, 2),
+                        ))
+                    if self.monitor_memory:
+                        virt_mem = psutil.virtual_memory()
+                        swap = psutil.swap_memory()
+                        self.logging("epoch={}/micro_step={}/global_step={}, vm %: {}, swap %: {}".format(
+                            self.epoch_count,
+                            self.micro_step_count,
+                            self.global_step_count,
+                            virt_mem.percent,
+                            swap.percent,
+                        ))
+                self.step_elapsed_time = 0
+
+    def avg_samples_per_sec(self):
+        if self.global_step_count > 0:
+            total_step_offset = self.global_step_count - self.start_step
+            avg_time_per_step = self.total_elapsed_time / total_step_offset
+            # training samples per second
+            return self.batch_size / avg_time_per_step
+        return float("-inf")
+
+
+def trim_mean(data, trim_percent):
+    """Compute the trimmed mean of a list of numbers.
+
+    Args:
+        data (list): List of numbers.
+        trim_percent (float): Percentage of data to trim.
+
+    Returns:
+        float: Trimmed mean.
+    """
+    assert 0.0 <= trim_percent <= 1.0
+    n = len(data)
+    # Account for edge case of empty list
+    if len(data) == 0:
+        return 0
+    data.sort()
+    k = int(round(n * (trim_percent)))
+    return mean(data[k:n - k])

deepspeed/utils/torch.py

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import logging
+import sys
+
+from packaging import version as pkg_version
+
+import torch
+
+_logger = logging.getLogger(__name__)
+
+
+def required_torch_version(min_version=None, max_version=None):
+    assert min_version or max_version, "Must provide a min_version or max_version argument"
+
+    torch_version = pkg_version.parse(torch.__version__)
+
+    if min_version and pkg_version.parse(str(min_version)) > torch_version:
+        return False
+
+    if max_version and pkg_version.parse(str(max_version)) < torch_version:
+        return False
+
+    return True
+
+
+def is_functorch_transforming():
+    """True when called under torch.func.grad / vmap / jacrev / etc."""
+    # peek_interpreter_stack is private; getattr keeps this working if it moves.
+    _functorch = getattr(torch._C, "_functorch", None)
+    if _functorch is None:
+        return False
+    peek = getattr(_functorch, "peek_interpreter_stack", None)
+    return peek is not None and peek() is not None
+
+
+def register_grad_hook(param, hook):
+    if required_torch_version(min_version=2.1):
+        return param.register_post_accumulate_grad_hook(hook)
+    else:
+        param_tmp = param.expand_as(param)
+        grad_acc = param_tmp.grad_fn.next_functions[0][0]
+        return grad_acc.register_hook(hook)
+
+
+def jit_script_compat(fn):
+    fn_name = getattr(fn, "__qualname__", getattr(fn, "__name__", repr(fn)))
+
+    can_try_compile = (required_torch_version(min_version=2.0) and hasattr(torch, "compile")
+                       and not (sys.version_info >= (3, 12) and not required_torch_version(min_version=2.4)))
+
+    if can_try_compile:
+        try:
+            return torch.compile(fn)
+        except Exception:
+            _logger.debug(
+                "torch.compile failed for %s, falling back to torch.jit.script",
+                fn_name,
+                exc_info=True,
+            )
+
+    try:
+        return torch.jit.script(fn)
+    except Exception:
+        _logger.debug(
+            "torch.jit.script failed for %s, returning unmodified function",
+            fn_name,
+            exc_info=True,
+        )
+        return fn

deepspeed/utils/types.py

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from enum import IntEnum
+
+
+class ActivationFuncType(IntEnum):
+    UNKNOWN = 0
+    GELU = 1
+    ReLU = 2
+    GATED_GELU = 3
+    GATED_SILU = 4
+
+
+GATED_ACTIVATION_TYPES = [
+    ActivationFuncType.GATED_GELU,
+    ActivationFuncType.GATED_SILU,
+]
+
+
+class NormType(IntEnum):
+    UNKNOWN = 0
+    LayerNorm = 1
+    GroupNorm = 2
+    RMSNorm = 3

deepspeed/utils/z3_leaf_module.py

@@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from typing import List, Tuple, Type, Union, Optional, TYPE_CHECKING
+
+from .logging import logger
+
+if TYPE_CHECKING:
+    from deepspeed.runtime.zero.leaf_module_config import DeepSpeedZeroLeafModuleConfig
+
+
+def z3_leaf_module(model: torch.nn.Module) -> bool:
+    """Returns whether a module in `model` has been flagged as a 'leaf' module.
+        See `set_z3_leaf_modules` for more details.
+        Args:
+            model (torch.nn.Module): The model to which the leaf module flag will be applied.
+        Returns:
+            bool: Whether the module has been flagged as a 'leaf' module.
+    """
+    return hasattr(model, '_z3_leaf') and model._z3_leaf
+
+
+def z3_leaf_parameter(model: torch.nn.Parameter) -> bool:
+    """Returns whether a parameter belongs to a leaf module.
+        See `set_z3_leaf_modules` for more details.
+        Args:
+            model (torch.nn.Parameter): The parameter to which the leaf module flag will be applied.
+        Returns:
+            bool: Whether the parameter belongs to a leaf module.
+    """
+    return hasattr(model, 'ds_z3_leaf_module')
+
+
+def get_z3_leaf_modules(model: torch.nn.Module) -> List[torch.nn.Module]:
+    """Returns a list of modules in `model` that have been flagged as 'leaf' modules.
+        See `set_z3_leaf_modules` for more details.
+        Args:
+            model (torch.nn.Module): The model to which the leaf module flag will be applied.
+        Returns:
+            List[torch.nn.Module]: A list of modules that have been flagged as 'leaf' modules.
+    """
+    return [module for module in model.modules() if z3_leaf_module(module)]
+
+
+def set_z3_leaf_module(model: torch.nn.Module, flag: bool):
+    model._z3_leaf = flag
+
+
+def _fully_qualified_class_name(module: torch.nn.Module) -> str:
+    cls = module.__class__
+    return f"{cls.__module__}.{cls.__qualname__}"
+
+
+def _do_set_z3_leaf_modules(model: torch.nn.Module,
+                            leaf_module_classes: Union[List[Type], List[str]],
+                            flag: bool,
+                            raise_if_not_found: bool = True) -> List[torch.nn.Module]:
+    assert all(isinstance(module_class, (type, str)) for module_class in leaf_module_classes), \
+        f'leaf_module_classes must be a list of types or names, got {leaf_module_classes}'
+
+    leaf_modules: List[torch.nn.Module] = []
+
+    def _set_z3_leaf_flag(module_instance: torch.nn.Module):
+        nonlocal leaf_modules
+        for module in leaf_module_classes:
+            if isinstance(module, type) and isinstance(module_instance, module):
+                module_instance._z3_leaf = flag
+                leaf_modules.append(module_instance)
+                break
+
+            if isinstance(module, str):
+                if (module_instance.__class__.__name__ == module
+                        or _fully_qualified_class_name(module_instance) == module):
+                    module_instance._z3_leaf = flag
+                    leaf_modules.append(module_instance)
+                    break
+
+    model.apply(_set_z3_leaf_flag)
+
+    if len(leaf_modules) == 0 and raise_if_not_found:
+        raise ValueError(f'No modules of type {leaf_module_classes} found in model {model}')
+
+    return leaf_modules
+
+
+def set_z3_leaf_modules_by_name(model: torch.nn.Module,
+                                module_names: List[str],
+                                flag: bool = True,
+                                raise_if_not_found: bool = True) -> Tuple[List[torch.nn.Module], List[str]]:
+    """Sets a leaf flag for modules referenced by their names in ``model.named_modules()``.
+        Args:
+            model (torch.nn.Module): The model containing the modules to update.
+            module_names (List[str]): Module names as returned by ``named_modules()``.
+            flag (bool): Desired flag state.
+            raise_if_not_found (bool): Whether to raise when no module matches a provided name.
+        Returns:
+            Tuple[List[torch.nn.Module], List[str]]: Matched modules and missing module names.
+    """
+    modules_by_name = dict(model.named_modules())
+    leaf_modules: List[torch.nn.Module] = []
+    missing: List[str] = []
+
+    for name in module_names:
+        module = modules_by_name.get(name)
+        if module is None:
+            missing.append(name)
+            continue
+        module._z3_leaf = flag
+        leaf_modules.append(module)
+
+    if missing and raise_if_not_found:
+        raise ValueError(f'No modules named {missing} found in model {model}')
+
+    return leaf_modules, missing
+
+
+def set_z3_leaf_modules_by_suffix(model: torch.nn.Module,
+                                  module_name_suffixes: List[str],
+                                  flag: bool = True,
+                                  raise_if_not_found: bool = True) -> Tuple[List[torch.nn.Module], List[str]]:
+    """Sets a leaf flag for modules referenced by suffixes of ``model.named_modules()`` names."""
+    modules_by_name = dict(model.named_modules())
+    leaf_modules: List[torch.nn.Module] = []
+    missing: List[str] = []
+    seen_ids = set()
+
+    for suffix in module_name_suffixes:
+        matched = False
+        for name, module in modules_by_name.items():
+            if name.endswith(suffix):
+                module._z3_leaf = flag
+                module_id = id(module)
+                if module_id not in seen_ids:
+                    seen_ids.add(module_id)
+                    leaf_modules.append(module)
+                matched = True
+        if not matched:
+            missing.append(suffix)
+
+    if missing and raise_if_not_found:
+        raise ValueError(f'No modules matching suffixes {missing} found in model {model}')
+
+    return leaf_modules, missing
+
+
+def set_z3_leaf_modules(model: torch.nn.Module,
+                        leaf_module_classes: Union[List[Type], List[str]],
+                        raise_if_not_found: bool = True) -> List[torch.nn.Module]:
+    """Sets a flag within a module in `model` to instruct ZeRO3 to stop setting hooks recursively when it encounters a module class listed in `leaf_module_classes`.
+       This is particularly useful in the context of Mixture of Experts (MoE) models. In MoE models, the computation order of experts varies across forward passes. This variability can disrupt ZeRO3's functionality, as ZeRO3 relies on tracking the computation order of modules to prefetch parameters efficiently. By designating a module as a 'leaf' node, ZeRO3 will prefetch parameters for all child modules upon entering the module.
+       Another scenario where this functionality is beneficial is in models with excessively fine-grained nested modules, where it helps to avoid the overhead associated with hooks.
+        Args:
+            model (torch.nn.Module): The model to which the leaf module flag will be applied.
+            leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules.
+            raise_if_not_found (bool): Whether to raise a ``ValueError`` when none of the provided classes
+                match a module inside ``model``.
+        Returns:
+            List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`.
+    """
+    return _do_set_z3_leaf_modules(model, leaf_module_classes, True, raise_if_not_found)
+
+
+def unset_z3_leaf_modules(model: torch.nn.Module,
+                          leaf_module_classes: List[Type],
+                          raise_if_not_found: bool = True) -> List[torch.nn.Module]:
+    """Unsets a flag within a module in `model` to instruct ZeRO3 to resume setting hooks recursively when it encounters a module class listed in `leaf_module_classes`.
+        See `set_z3_leaf_modules` for more details.
+        Args:
+            model (torch.nn.Module): The model to which the leaf module flag will be applied.
+            leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules.
+            raise_if_not_found (bool): Whether to raise a ``ValueError`` when none of the provided classes
+                match a module inside ``model``.
+        Returns:
+            List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`.
+    """
+    return _do_set_z3_leaf_modules(model, leaf_module_classes, False, raise_if_not_found)
+
+
+def apply_zero_leaf_module_config(model: torch.nn.Module,
+                                  leaf_cfg: Optional["DeepSpeedZeroLeafModuleConfig"]) -> List[torch.nn.Module]:
+    """Apply ZeRO leaf module configuration to ``model``.
+
+    Args:
+        model (torch.nn.Module): Root module to update.
+        leaf_cfg (DeepSpeedZeroLeafModuleConfig | None): Parsed configuration. If ``None``
+            no changes are applied.
+
+    Returns:
+        List[torch.nn.Module]: Modules flagged as leaves.
+    """
+    if leaf_cfg is None:
+        return []
+
+    from deepspeed.runtime.zero.leaf_module_config import (
+        DEFAULT_LEAF_MODULE_CLASSES,
+        DEFAULT_LEAF_MODULE_NAMES,
+        DEFAULT_LEAF_MODULE_NAME_SUFFIXES,
+    )
+
+    matched_modules: List[torch.nn.Module] = []
+    matched_ids = set()
+
+    customized_classes = leaf_cfg.classes != DEFAULT_LEAF_MODULE_CLASSES
+    customized_names = leaf_cfg.names != DEFAULT_LEAF_MODULE_NAMES
+    customized_suffixes = leaf_cfg.name_suffixes != DEFAULT_LEAF_MODULE_NAME_SUFFIXES
+
+    if leaf_cfg.classes:
+        class_matched = set_z3_leaf_modules(model, leaf_cfg.classes, raise_if_not_found=False)
+        for module in class_matched:
+            module_id = id(module)
+            if module_id not in matched_ids:
+                matched_ids.add(module_id)
+                matched_modules.append(module)
+
+    if leaf_cfg.names:
+        name_matched, missing_names = set_z3_leaf_modules_by_name(model,
+                                                                  leaf_cfg.names,
+                                                                  flag=True,
+                                                                  raise_if_not_found=False)
+        for module in name_matched:
+            module_id = id(module)
+            if module_id not in matched_ids:
+                matched_ids.add(module_id)
+                matched_modules.append(module)
+
+        if missing_names and customized_names:
+            logger.warning(f"ZeRO leaf module configuration contains unknown module names: {missing_names}")
+
+    if leaf_cfg.name_suffixes:
+        suffix_matched, missing_suffixes = set_z3_leaf_modules_by_suffix(model,
+                                                                         leaf_cfg.name_suffixes,
+                                                                         flag=True,
+                                                                         raise_if_not_found=False)
+        for module in suffix_matched:
+            module_id = id(module)
+            if module_id not in matched_ids:
+                matched_ids.add(module_id)
+                matched_modules.append(module)
+
+        if missing_suffixes and customized_suffixes:
+            logger.warning(f"ZeRO leaf module configuration contains unmatched module suffixes: {missing_suffixes}")
+
+    if not matched_modules and (customized_classes or customized_names or customized_suffixes):
+        logger.warning("ZeRO leaf module configuration did not match any modules; hooks will be applied as usual")
+
+    return matched_modules

deepspeed/utils/zero_to_fp32.py

@@ -0,0 +1,790 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+#   python zero_to_fp32.py . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS, AUTOEP_LAYERS_KEY,
+                                            AUTOEP_LAYERS_KEY_LEGACY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
+                                            AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def _has_autoep_zero3_partitioned_metadata(state_dict):
+    autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY)
+    if autoep_layers is None:
+        autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY_LEGACY)
+    if not isinstance(autoep_layers, list):
+        return False
+    return any(
+        isinstance(entry, dict)
+        and entry.get(AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY) == AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT
+        for entry in autoep_layers)
+
+
+def _raise_if_autoep_zero3_partitioned_state(state_dict):
+    if _has_autoep_zero3_partitioned_metadata(state_dict):
+        raise NotImplementedError("zero_to_fp32 does not support AutoEP ZeRO-3 partition-native checkpoints. "
+                                  "AutoEP expert parameters are partitioned over expert replica groups, so "
+                                  "global data-parallel consolidation would produce incomplete expert tensors. "
+                                  "Use ds_to_universal.py for expert-aware conversion.")
+
+
+def _raise_if_autoep_zero3_partitioned_checkpoint(model_files):
+    for file in model_files:
+        state_dict = torch.load(file, map_location=device, weights_only=False)
+        _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device, weights_only=False)
+        _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+    total_files = len(files)
+    state_dicts = []
+    for f in tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    # parse_model_states rejects AutoEP ZeRO-3 partition-native checkpoints
+    # before the expensive optimizer-shard load below.
+    model_files = get_model_state_files(ds_checkpoint_dir)
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        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):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 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
+        - ``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
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    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)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    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)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``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``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info("Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info("Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    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 to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. 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")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+                                               args.output_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)