chore: import upstream snapshot with attribution

deepspeed/runtime/zero/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .partition_parameters import ZeroParamType
+from .partition_parameters import ZeroParamStatus
+from .partition_parameters import Init
+from .partition_parameters import GatheredParameters
+from .partition_parameters import register_external_parameter
+from .parameter_offload import DeepSpeedZeRoOffload
+from .partition_parameters import DeepSpeedTensorOverride
+
+from .tiling import TiledLinear
+from .tiling import TiledLinearReturnBias
+
+from .mics import MiCS_Init
+
+from .stage3 import unwrap_model_for_generation

deepspeed/runtime/zero/config.py

@@ -0,0 +1,401 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import sys
+from typing import Optional, Dict, Any
+from enum import Enum
+from pydantic import Field, model_validator
+from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel
+from deepspeed.utils import logger
+from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum
+from deepspeed.runtime.zenflow.zenflow_config import ZenFlowConfig
+from .leaf_module_config import DeepSpeedZeroLeafModuleConfig
+
+# ZeRO optimization. By default, this optimization is not enabled.
+# Users have to configure the desired optimization (0 means disabled) in params.json as below example:
+ZERO_FORMAT = """
+ZeRO optimization should be enabled as:
+"session_params": {
+  "zero_optimization": {
+    "stage": [0|1|2],
+    "stage3_max_live_parameters" : 1000000000,
+    "stage3_max_reuse_distance" : 1000000000,
+    "stage3_use_all_reduce_for_fetch_params": [true|false],
+    "stage3_module_granularity_threshold": 0,
+    "allgather_partitions": [true|false],
+    "use_multi_rank_bucket_allreduce": [true|false],
+    "stage3_allgather_sequential": [true|false],
+    "allgather_bucket_size": 500000000,
+    "reduce_scatter": [true|false],
+    "contiguous_gradients" : [true|false]
+    "overlap_comm": [true|false],
+    "reduce_bucket_size": 500000000,
+    "load_from_fp32_weights": [true|false],
+    "cpu_offload": [true|false] (deprecated),
+    "cpu_offload_param" : [true|false] (deprecated),
+    "cpu_offload_use_pin_memory": [true|false] (deprecated),
+    "sub_group_size" : 1000000000000,
+    "offload_param": {...},
+    "offload_optimizer": {...},
+    "ignore_unused_parameters": [true|false],
+    "round_robin_gradients": [true|false],
+    "zero_hpz_partition_size": 1,
+    "zero_quantized_weights": [true|false],
+    "zero_quantized_nontrainable_weights": [true|false],
+    "zero_quantized_gradients": [true|false],
+    "memory_efficient_linear": [true|false],
+    "override_module_apply": [true|false],
+    "zeropp_loco_param": {...},
+    "log_trace_cache_warnings" : [true|false],
+    "enable_sanity_checks": [true|false],
+    }
+}
+"""
+
+ZERO_OPTIMIZATION = "zero_optimization"
+
+
+def read_zero_config_deprecated(param_dict):
+    zero_config_dict = {}
+    zero_config_dict["stage"] = 1 if param_dict[ZERO_OPTIMIZATION] else 0
+    if zero_config_dict["stage"] > 0:
+        zero_config_dict["allgather_bucket_size"] = get_scalar_param(param_dict, "allgather_size", 5e8)
+    logger.warning(
+        "DeepSpeedConfig: this format of ZeRO optimization setup is deprecated. Please use the following format: {}".
+        format(ZERO_FORMAT))
+    return zero_config_dict
+
+
+def get_zero_config(param_dict):
+    if ZERO_OPTIMIZATION in param_dict:
+        zero_config_dict = param_dict[ZERO_OPTIMIZATION]
+        if isinstance(zero_config_dict, bool):
+            zero_config_dict = read_zero_config_deprecated(param_dict)
+    else:
+        zero_config_dict = {}
+    return DeepSpeedZeroConfig(**zero_config_dict)
+
+
+class ZeroStageEnum(int, Enum):
+    """ Enum class for possible zero stages """
+    disabled = 0
+    optimizer_states = 1
+    gradients = 2
+    weights = 3
+    max_stage = 3
+
+
+class DeepSpeedZeroConfig(DeepSpeedConfigModel):
+    """
+    Sets parameters for ZeRO optimizations.
+    """
+
+    stage: ZeroStageEnum = 0
+    """
+    Chooses different stages of ZeRO Optimizer. Stage 0, 1, 2, and 3 refer
+    to disabled, optimizer state partitioning, and optimizer+gradient state
+    partitioning, and optimizer+gradient+parameter partitioning, respectively.
+    """
+
+    contiguous_gradients: bool = True
+    """
+    Copies the gradients to a contiguous buffer as they are produced. Avoids
+    memory fragmentation during backward pass.
+    """
+
+    reduce_scatter: bool = True
+    """
+    Uses reduce or reduce scatter instead of allreduce to average gradients
+    """
+
+    reduce_bucket_size: int = Field(pp_int(5e8), ge=0)
+    """
+    Number of elements reduced/allreduced at a time. Limits the memory required
+    for the allgather for large model sizes
+    """
+
+    use_multi_rank_bucket_allreduce: bool = True
+    """
+    Combine the reduce buckets of the different ranks and do an All-Reduce instead of multiple Reduce ops.
+    This feature is useful when the model is small and we want to scale it on too many GPUs which therefore
+    reduces the message sizes of each packet.
+    """
+
+    allgather_partitions: bool = True
+    """
+    Chooses between allgather collective or a series of broadcast collectives
+    to gather updated parameters from all the GPUs at the end of each step
+    """
+
+    allgather_bucket_size: int = Field(pp_int(5e8), ge=0)
+    """
+    Number of elements allgathered at a time. Limits the memory required for
+    the allgather for large model sizes
+    """
+
+    overlap_comm: Optional[bool] = None  # None for dynamic default value (see validator `overlap_comm_valid` below)
+    """
+    Attempts to overlap the reduction of the gradients with backward computation
+    """
+
+    load_from_fp32_weights: bool = True
+    """
+    Boolean indicating whether to initialize fp32 master weights from fp32
+    copies in checkpoint (no precision loss) or from model's fp16 copies (with
+    precision loss). This can be used to initialize optimizer state even when
+    checkpoint is missing optimizer state.
+    """
+
+    elastic_checkpoint: bool = False
+    """
+    Legacy elastic checkpoint support. ZeRO-3 elastic checkpointing is no
+    longer supported; use Universal Checkpointing instead.
+    """
+
+    offload_param: Optional[DeepSpeedZeroOffloadParamConfig] = None
+    """
+    Enable offloading of model parameters to CPU or NVMe. This frees up GPU
+    memory for larger models or batch sizes. Valid only with stage 3. Expects a
+    dictionary containing values for :any:`DeepSpeedZeroOffloadParamConfig`.
+    """
+
+    offload_optimizer: Optional[DeepSpeedZeroOffloadOptimizerConfig] = None
+    """
+    Enable offloading of optimizer state to CPU or NVMe, and optimizer
+    computation to CPU. This frees up GPU memory for larger models or batch
+    sizes. Valid for ZeRO stage 1, 2, 3. Expects a dictionary containing values
+    for :any:`DeepSpeedZeroOffloadOptimizerConfig`.
+    """
+
+    zenflow: Optional[ZenFlowConfig] = None
+    """Enable ZenFlow"""
+
+    sub_group_size: int = Field(pp_int(1e9), ge=0)
+    """
+    Tile size for parameter processing to fit massive models (with trillions of
+    parameters). Used by ZeRO3-Offload and ZeRO-Infinity
+    """
+
+    cpu_offload_param: Optional[bool] = Field(
+        None,
+        json_schema_extra={
+            "deprecated": True,
+            "new_param": "offload_param",
+            "new_param_fn": (lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu)
+                             if val else None)
+        },
+    )
+    """ Deprecated, please use ``offload_param`` """
+
+    cpu_offload_use_pin_memory: Optional[bool] = Field(
+        None,
+        json_schema_extra={
+            "deprecated": True,
+            "new_param": "offload_param or offload_optimizer",
+            "set_new_param": False
+        },
+    )
+    """ Deprecated, please use ``offload_param`` or ``offload_optimizer`` """
+
+    cpu_offload: Optional[bool] = Field(
+        None,
+        json_schema_extra={
+            "deprecated":
+            True,
+            "new_param":
+            "offload_optimizer",
+            "new_param_fn": (lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu)
+                             if val else None)
+        },
+    )
+    """ Deprecated, please use ``offload_optimizer`` """
+
+    prefetch_bucket_size: int = Field(pp_int(5e7), ge=0, alias="stage3_prefetch_bucket_size")
+    """
+    Maximum number of parameter elements to fetch ahead of use. Used by ZeRO3,
+    ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference.
+    """
+
+    param_persistence_threshold: int = Field(pp_int(1e5), ge=0, alias="stage3_param_persistence_threshold")
+    """
+    Do not partition parameters smaller than this threshold. Smaller values use
+    less memory, but can greatly increase communication (especially
+    latency-bound messages).
+    """
+
+    model_persistence_threshold: int = Field(pp_int(sys.maxsize, "sys.maxsize"),
+                                             ge=0,
+                                             alias="stage3_model_persistence_threshold")
+    """
+    Maximum number of parameter elements that can be persisted in GPU and not
+    partitioned. This imposes an upper bound on the number of unpartitioned
+    parameters resulting from param_persistence_threshold setting. Used by
+    ZeRO3-Offload, ZeRO-Infinity and ZeRO-Inference.
+    """
+
+    max_live_parameters: int = Field(pp_int(1e9), ge=0, alias="stage3_max_live_parameters")
+    """
+    The maximum number of parameters resident per GPU before releasing. Smaller
+    values use less memory, but perform more communication.
+    """
+
+    max_reuse_distance: int = Field(pp_int(1e9), ge=0, alias="stage3_max_reuse_distance")
+    """
+    Do not release a parameter if it will be reused within this threshold of
+    parameters. Smaller values use less memory, but perform more communication.
+    """
+
+    gather_16bit_weights_on_model_save: bool = Field(False, alias="stage3_gather_16bit_weights_on_model_save")
+    """
+    Consolidate the weights before saving the model by ``save_16bit_model()``.
+    Since the weights are partitioned across GPUs, they aren’t part of
+    ``state_dict``, so this function automatically gathers the weights when
+    this option is enabled and then saves the fp16 model weights.
+    """
+
+    module_granularity_threshold: int = Field(pp_int(0), alias="stage3_module_granularity_threshold")
+    """
+    The granularity of a module is determined by the ratio of "parameter_count / (1 + descendant count)".
+    ZeRO3 classifies modules with a granularity below the threshold as fine-grained,
+    which are treated as integral units during parameter fetching. This reduces host overhead
+    and the separate allgather overhead introduced by hooks for fine-grained layers when fetching parameters.
+    """
+
+    use_all_reduce_for_fetch_params: bool = Field(False, alias="stage3_use_all_reduce_for_fetch_params")
+    """
+    Use all_reduce op when fetching module parameters at stage3. This improves performance by reducing
+    the overhead of concatenation and slicing on the host.
+    """
+
+    allgather_sequential: bool = Field(default=False, alias="stage3_allgather_sequential")
+    """
+    Performs allgather on individual parameters sequentially, bypassing the standard parameter bucketing
+    mechanism in stage3. This significantly reduces data copy overhead (eliminating copy-to-bucket operations)
+    and lowers peak memory usage by avoiding the allocation of large temporary flattening buffers.
+    Recommended for scenarios with high memory pressure.
+    """
+
+    stage3_gather_fp16_weights_on_model_save: bool = Field(False,
+                                                           json_schema_extra={
+                                                               "deprecated": True,
+                                                               "new_param": "gather_16bit_weights_on_model_save"
+                                                           })
+    """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """
+
+    ignore_unused_parameters: bool = True
+    """
+    Unused parameters in modules may be unexpected in static networks, but
+    could be normal in dynamic networks. This controls whether or not training
+    should terminate with an error message when unused parameters are detected.
+    This is set to ``True`` by default, which means unused parameters are
+    ignored and training continues. Now is just used in stage 2.
+    """
+
+    legacy_stage1: bool = False
+    """
+    For backward-compatibility enable old ZeRO stage 1 implementation. Use at
+    your own risk, will be deprecated soon.
+    """
+
+    round_robin_gradients: bool = False
+    """
+    Stage 1 and 2 optimization for CPU offloading that parallelizes gradient
+    copying to CPU memory among ranks by fine-grained gradient partitioning.
+    Performance benefit grows with gradient accumulation steps (more copying
+    between optimizer steps) or GPU count (increased parallelism).
+    """
+    zero_hpz_partition_size: int = Field(1, ge=0)
+    """
+    Number of ranks in zero parameters partitioning secondary group
+    """
+    zero_quantized_weights: bool = False
+    """
+    Boolean indicating whether to quantize zero parameters (weights)
+    for efficient all_gather comm
+    """
+    zero_quantized_nontrainable_weights: bool = False
+    """
+    Boolean indicating whether to quantize non-trainable zero parameters (weights)
+    for efficient memory usage and communication. Different from zero_quantized_weights
+    that stores the weights in original precision and only perform quantization during communication,
+    this flag will store the weights in quantized precision. This is useful for LoRA training.
+    """
+    zero_quantized_gradients: bool = False
+    """
+    Boolean indicating whether to use quantized zero gradients
+    for efficient all_2_all_reduce comm
+    """
+    zeropp_loco_param: Optional[Dict[str, Any]] = None
+    """
+    This dictionary contains parameters for using LoCo-Zero++, with two key parameters:
+    - `err_beta`: A coefficient for the moving average of quantization errors before and after gradient computation.
+    It ranges between 0 and 1, with a default value of 0.8.
+    - `reset_T`: The number of steps after which the moving-average error buffer is cleared. The default value is 1024.
+    These parameters can be adjusted based on performance needs. Example configuration in ds config:
+    "zeropp_loco_param": { "err_beta": 0.8, "reset_T": 1024 }.
+    See LoCo paper for more details: (https://arxiv.org/abs/2407.04480).
+    """
+
+    mics_shard_size: int = Field(-1, json_schema_extra={"new_param": "mics_shard_size"})
+
+    mics_hierarchical_params_gather: bool = False
+
+    memory_efficient_linear: bool = True
+    """
+    Use memory efficient linear implementation, for Stage 3.
+    """
+    """
+    Whether force load checkpoint in pipeline mode, current only for Stage 3.
+    """
+    pipeline_loading_checkpoint: bool = False
+
+    override_module_apply: bool = True
+    """
+    Override nn.Module apply function, for Stage 3.
+    """
+
+    log_trace_cache_warnings: bool = False
+    """
+    Whether to log warnings from trace cache, such as invalidation events.
+    """
+
+    enable_sanity_checks: bool = False
+    """
+    Enable internal sanity checks, which could be useful for debugging
+    """
+
+    save_muon_momentum_buffer_in_memory: bool = False
+    """
+    When using the Muon optimizer with ZeRO Stage 3, keeps the Muon momentum
+    buffer in GPU/CPU memory instead of swapping to NVMe with other optimizer
+    states. Only relevant when using NVMe offloading.
+    """
+
+    leaf_module: DeepSpeedZeroLeafModuleConfig = Field(default_factory=DeepSpeedZeroLeafModuleConfig)
+    """
+    Configuration for modules that should be treated as ZeRO3 leaf modules.
+    """
+
+    # Validators
+    @model_validator(mode="after")
+    def overlap_comm_valid(self):
+        if self.overlap_comm is None:
+            self.overlap_comm = self.stage == ZeroStageEnum.weights
+        return self
+
+    @model_validator(mode="after")
+    def offload_ratio_check(self):
+        offload_config = self.offload_optimizer
+        if offload_config and offload_config.ratio < 1.0:
+            assert self.stage == ZeroStageEnum.weights, "Partial offloading only supported for ZeRO Stage 3."
+        return self
+
+    @model_validator(mode="after")
+    def elastic_checkpoint_deprecated(self):
+        if self.stage == ZeroStageEnum.weights and self.elastic_checkpoint:
+            logger.warning(
+                "ZeRO-3 elastic checkpointing is deprecated and no longer supported. Use Universal Checkpointing instead."
+            )
+        return self

deepspeed/runtime/zero/contiguous_memory_allocator.py

@@ -0,0 +1,287 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed import comm as dist
+
+
+def print_rank_0(message):
+    if dist.get_rank() == 0:
+        print(message)
+
+
+class ContiguousMemoryAllocator(object):
+
+    def __init__(self, size, dtype, device):
+        self.buffer = torch.zeros(size, dtype=dtype, device=device)
+
+        #address to contiguous size available
+        self.contiguous_sizes = {}
+
+        self.contiguous_sizes[0] = size
+
+        #tensor id to its address
+        self.tensor_addresses = {}
+
+        #tensor address to its size
+        self.tensor_sizes = {}
+
+        #tensor address to ids
+        self.tensor_ids = {}
+
+        #id to tensors
+        self.tensor_map = {}
+
+        #id to params. Maps each tensor buffer to list of parameters that uses it
+        self.id_to_params = {}
+
+        self.total_size = size
+        self.total_free = size
+        self.largest_contiguous = size
+        self.max_allocated = 0
+
+        self.count = 0
+
+    #create a tensor of size from the pre-allocated buffer
+    #if not enough free space will fail
+    #if not enough contiguous space, will defragment and allocate
+    def allocate_tensor(self, size):
+        free_before = self.total_free
+
+        assert size <= self.total_free, "Not enough memory in buffer. Allocation failed"
+        if self.largest_contiguous < size:
+            print_rank_0("Needs defragmentation to allocate. Before Defragmentation:")
+            self.print_allocation(resolution=100)
+            self._defragment_memory()
+            #set the param data to the new tensor buffer locations
+            self._reset_param_data()
+            print_rank_0("After defragmentation:")
+            self.print_allocation(resolution=100)
+
+        self.total_free = self.total_free - size
+
+        allocated = self.total_size - self.total_free
+        if allocated > self.max_allocated:
+            self.max_allocated = allocated
+
+        tensor_address = self._get_new_tensor_address(size)
+
+        ret_tensor = self._get_new_tensor(tensor_address, size)
+        print_rank_0(
+            f"Free before allocation {free_before}. Allocating {size}. Free after allocation {self.total_free}. Max allocated {self.max_allocated}"
+        )
+        assert self.total_free + size == free_before, "Allocation bookkeeping error"
+
+        return ret_tensor
+
+    #assigns the tensor data to the param data and keeps track of the assignment
+    #any change the underlying buffer from defragmentation will cause a
+    #reassignment of the param data
+    def assign_to_param(self, tensor, param, numel, shape):
+        tensor_id = id(tensor)
+
+        assert tensor_id in self.tensor_map.keys(), "No such tensor allocated by the allocator."
+        assert tensor.numel() >= numel, "Assert tensor buffer does is not large enough"
+        assert tensor_id not in self.id_to_params.keys(), "This tensor has already been assigned to a param"
+
+        self.id_to_params[tensor_id] = [param]
+
+        replicated_tensor = tensor.narrow(0, 0, numel).view(shape)
+        param.data = replicated_tensor.data
+        param.contiguous_tensor_id = tensor_id
+
+    #deletes the tensor and frees up the underlying buffer
+    def release_tensor(self, tensor):
+        free_before = self.total_free
+        tensor_id = id(tensor)
+        tensor_size = tensor.numel()
+        self._release_tensor(tensor_id)
+        self._unassign_params(tensor_id)
+        self.total_free += tensor_size
+        print_rank_0(
+            f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.")
+        assert self.total_free - tensor_size == free_before, "Release bookkeeping error"
+
+    def release_tensor_with_id(self, tensor_id):
+        free_before = self.total_free
+        assert tensor_id in self.tensor_map.keys(), "Invalid tensor id"
+        tensor = self.tensor_map[tensor_id]
+        tensor_size = tensor.numel()
+        self._release_tensor(tensor_id)
+        self._unassign_params(tensor_id)
+        self.total_free += tensor_size
+        print_rank_0(
+            f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.")
+        assert self.total_free - tensor_size == free_before, "Release bookkeeping error"
+
+    #shows the current memory allocation at specified resolution
+    def print_allocation(self, resolution=200):
+        total_size = self.buffer.numel() * 1.0
+        empty = []
+        for addr, size in self.contiguous_sizes.items():
+            start = int(addr * resolution / total_size)
+            end = int((addr + size) * resolution / total_size)
+            empty.extend(range(start, end))
+        s = ''
+        for i in range(resolution):
+            s += '.' if i in empty else '|'
+        print_rank_0(s)
+
+    def max_allocated(self):
+        return self.max_allocated
+
+    #to be called after defragmentation that moves the tensor buffers
+    #this call reassigns the data of all the parameters using the tensor buffers
+    def _reset_param_data(self):
+        for id, tensor in self.tensor_map.items():
+            for param in self.id_to_params[id]:
+                param.data = tensor.narrow(0, 0, param.numel()).view(param.data.shape).data
+
+    def _unassign_params(self, tensor_id):
+        if tensor_id in self.id_to_params.keys():
+            del self.id_to_params[tensor_id]
+
+    def _release_tensor(self, tensor_id):
+        assert tensor_id in self.tensor_addresses, f"Tensor id {tensor_id} not found"
+
+        address = self.tensor_addresses[tensor_id]
+        contiguous_size = self.tensor_map[tensor_id].numel()
+
+        del self.tensor_addresses[tensor_id]
+        del self.tensor_ids[address]
+        del self.tensor_map[tensor_id]
+        del self.tensor_sizes[address]
+
+        self._consolidate_address(address, contiguous_size)
+        self.largest_contiguous = self._largest_contiguous()
+
+    def _consolidate_address(self, address, contiguous_size):
+
+        #consolidate next buffer
+        end_address = address + contiguous_size
+        if end_address in self.contiguous_sizes:
+            contiguous_size += self.contiguous_sizes[end_address]
+            del self.contiguous_sizes[end_address]
+
+        #consolidate previous buffer
+        for addr, size in self.contiguous_sizes.items():
+            if addr + size == address:
+                del self.contiguous_sizes[addr]
+                contiguous_size += size
+                address = addr
+                break
+
+        self.contiguous_sizes[address] = contiguous_size
+
+    def _defragment_memory(self):
+        empty_addresses = sorted(self.contiguous_sizes.keys())
+        tensor_addresses = sorted(self.tensor_addresses.values())
+
+        tensor_index = 0
+
+        while tensor_index < len(tensor_addresses):
+
+            empty_addr = empty_addresses[0]
+            empty_size = self.contiguous_sizes[empty_addr]
+
+            tensor_addr = tensor_addresses[tensor_index]
+            tensor_size = self.tensor_sizes[tensor_addr]
+            tensor_id = self.tensor_ids[tensor_addr]
+            tensor = self.tensor_map[self.tensor_ids[tensor_addr]]
+
+            assert tensor_size == tensor.numel(), \
+                f"Size mismatch. {tensor_size} is allocated at addr {tensor_addr} but tensor size is {tensor.numel()} "
+
+            assert empty_addr != tensor_addr, \
+                f"Cannot have same empty address {empty_addr} and tensor address {tensor_addr}"
+
+            if empty_addr < tensor_addr:
+
+                if empty_size >= tensor_size:
+                    dest_buffer = self.buffer.narrow(0, empty_addr, tensor_size)
+                    src_buffer = self.buffer.narrow(0, tensor_addr, tensor_size)
+                    dest_buffer.data.copy_(src_buffer.data)
+                else:
+
+                    #print_rank_0(f'empty addr : {empty_addr}, empty size {empty_size} tensor addr {tensor_addr} tensor size {tensor_size}')
+                    src_addr = tensor_addr
+                    dest_addr = empty_addr
+                    while src_addr < (tensor_addr + tensor_size):
+                        copy_size = min(empty_size, tensor_addr + tensor_size - src_addr)
+
+                        dest_buffer = self.buffer.narrow(0, dest_addr, copy_size)
+                        src_buffer = self.buffer.narrow(0, src_addr, copy_size)
+
+                        dest_buffer.data.copy_(src_buffer.data)
+
+                        src_addr += copy_size
+                        dest_addr += copy_size
+
+                self._replace_old_address_with_new(tensor_id, empty_addr)
+
+                tensor_index += 1
+
+            else:
+                tensor_index += 1
+
+            empty_addresses = sorted(self.contiguous_sizes.keys())
+
+    def _replace_old_address_with_new(self, tensor_id, new_address):
+
+        tensor = self.tensor_map[tensor_id]
+        tensor_size = tensor.numel()
+        tensor.data = self.buffer.narrow(0, new_address, tensor_size).data
+
+        self._release_tensor(tensor_id)
+        self._mark_as_occupied(new_address, tensor_size)
+
+        self.tensor_ids[new_address] = tensor_id
+        self.tensor_map[tensor_id] = tensor
+        self.tensor_addresses[tensor_id] = new_address
+        self.tensor_sizes[new_address] = tensor_size
+
+    def _get_new_tensor_address(self, size):
+        tensor_address = None
+        for address, contiguous_size in self.contiguous_sizes.items():
+            if contiguous_size >= size and \
+                    (tensor_address is None or \
+                    contiguous_size < self.contiguous_sizes[tensor_address]):
+                tensor_address = address
+        assert tensor_address is not None, "address cannot be None"
+        return tensor_address
+
+    def _get_new_tensor(self, address, size):
+        available_contiguous_size = self.contiguous_sizes[address]
+
+        assert size <= available_contiguous_size, \
+            f"Tensor numel {size} is large than available contiguous size {available_contiguous_size}"
+        self.count += 1
+        new_tensor = self.buffer.narrow(0, address, size)
+        tensor_id = id(new_tensor)
+        self.tensor_addresses[tensor_id] = address
+        self.tensor_sizes[address] = size
+
+        self.tensor_ids[address] = tensor_id
+        self.tensor_map[tensor_id] = new_tensor
+
+        self._mark_as_occupied(address, size)
+
+        return new_tensor
+
+    def _largest_contiguous(self):
+        if len(self.contiguous_sizes) > 0:
+            return max([size for _, size in self.contiguous_sizes.items()])
+        else:
+            return 0
+
+    def _mark_as_occupied(self, address, size):
+        available_contiguous_size = self.contiguous_sizes[address]
+        del self.contiguous_sizes[address]
+
+        if available_contiguous_size != size:
+            self.contiguous_sizes[address + size] = available_contiguous_size - size
+
+        self.largest_contiguous = self._largest_contiguous()

deepspeed/runtime/zero/leaf_module_config.py

@@ -0,0 +1,52 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+from pydantic import Field, model_validator
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+
+DEFAULT_LEAF_MODULE_CLASSES: List[str] = [
+    "transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock",
+    "transformers.models.qwen2_moe.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock",
+    "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock",
+]
+DEFAULT_LEAF_MODULE_NAMES: List[str] = []
+DEFAULT_LEAF_MODULE_NAME_SUFFIXES: List[str] = []
+
+
+class DeepSpeedZeroLeafModuleConfig(DeepSpeedConfigModel):
+    """Configuration for ZeRO leaf modules that should bypass hook installation."""
+
+    classes: List[str] = Field(default_factory=lambda: list(DEFAULT_LEAF_MODULE_CLASSES))
+    names: List[str] = Field(default_factory=lambda: list(DEFAULT_LEAF_MODULE_NAMES))
+    name_suffixes: List[str] = Field(default_factory=lambda: list(DEFAULT_LEAF_MODULE_NAME_SUFFIXES))
+
+    @model_validator(mode="before")
+    def _coerce_container_types(cls, values):
+        if values is None:
+            return {}
+        if isinstance(values, dict):
+            coerced = dict(values)
+            for key in ("classes", "names", "name_suffixes"):
+                if key in coerced and isinstance(coerced[key], str):
+                    coerced[key] = [coerced[key]]
+            return coerced
+        raise TypeError("leaf_module configuration must be a mapping of fields to values")
+
+    @model_validator(mode="after")
+    def _validate_entries(self):
+        normalized_classes = [str(cls) for cls in self.classes]
+        normalized_names = [str(name) for name in self.names]
+        normalized_suffixes = [str(suffix) for suffix in self.name_suffixes]
+
+        deduped_classes = list(dict.fromkeys(normalized_classes))
+        deduped_names = list(dict.fromkeys(normalized_names))
+        deduped_suffixes = list(dict.fromkeys(normalized_suffixes))
+
+        object.__setattr__(self, "classes", deduped_classes)
+        object.__setattr__(self, "names", deduped_names)
+        object.__setattr__(self, "name_suffixes", deduped_suffixes)
+        return self

deepspeed/runtime/zero/linear.py

@@ -0,0 +1,203 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#Linear Module to use with ZeRO Stage 3 to allow for parameter memory release
+#after the module execution during forward
+#Instead of saving variables using save_for_backward, we save variable ids
+#Allowing us to retrieve the variable without creating pointer to it
+#Which allows for underlying tensor to be garbage collected
+#When partitioned as needed by the Zero Stage 3 optimizer
+#TODO instead of patching Linear module, we could patch the ctx.save_for_backward
+#ctx.saved_tensors so that this approach works for all nn modules that are built upon
+#torch.nn.function. However the issue is that many modules uses C++ implementations
+#which does not have pytorch implementation. Eg torch.addmm which acts as a functional
+#when implemented outside of torch.autograd.Function
+
+import math
+import functools
+
+import torch
+from torch import Tensor
+from torch.nn.parameter import Parameter
+from torch.nn import init
+from torch.nn.modules.module import Module
+from deepspeed.runtime.utils import noop_decorator
+from deepspeed import comm as dist
+from deepspeed.accelerator import get_accelerator
+
+
+def print_rank_0(message, debug=False, force=False):
+    if dist.get_rank() == 0 and (debug or force):
+        print(message)
+
+
+def _get_legacy_autocast_decorators(device_type):
+    legacy_amp = getattr(getattr(torch, device_type, None), 'amp', None)
+    custom_fwd = getattr(legacy_amp, 'custom_fwd', None)
+    custom_bwd = getattr(legacy_amp, 'custom_bwd', None)
+    if custom_fwd is not None and custom_bwd is not None:
+        return custom_fwd, custom_bwd
+    return noop_decorator, noop_decorator
+
+
+def _get_autocast_decorators():
+    amp = getattr(torch, 'amp', None)
+    custom_fwd = getattr(amp, 'custom_fwd', None)
+    custom_bwd = getattr(amp, 'custom_bwd', None)
+    if custom_fwd is not None and custom_bwd is not None:
+        device_type = get_accelerator().device_name()
+        return functools.partial(custom_fwd, device_type=device_type), functools.partial(custom_bwd,
+                                                                                         device_type=device_type)
+    return _get_legacy_autocast_decorators(get_accelerator().device_name())
+
+
+autocast_custom_fwd, autocast_custom_bwd = _get_autocast_decorators()
+
+
+def _is_autocast_enabled(device_type):
+    try:
+        return torch.is_autocast_enabled(device_type)
+    except TypeError:
+        legacy_getter = getattr(torch, f'is_autocast_{device_type}_enabled', None)
+        if legacy_getter is not None:
+            return legacy_getter()
+        return torch.is_autocast_enabled()
+
+
+def _get_autocast_dtype(device_type):
+    try:
+        return torch.get_autocast_dtype(device_type)
+    except TypeError:
+        legacy_getter = getattr(torch, f'get_autocast_{device_type}_dtype', None)
+        if legacy_getter is not None:
+            return legacy_getter()
+        return None
+
+
+class LinearFunctionForZeroStage3(torch.autograd.Function):
+
+    generate_vmap_rule = True
+
+    @staticmethod
+    # bias is an optional argument
+    def forward(input, weight, bias=None):
+
+        if input.dim() == 2 and bias is not None:
+            # fused op is marginally faster
+            ret = torch.addmm(bias, input, weight.t())
+        else:
+            output = input.matmul(weight.t())
+            if bias is not None:
+                output += bias
+            ret = output
+
+        return ret
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        device_type = get_accelerator().device_name()
+        ctx._dtype = _get_autocast_dtype(device_type)
+        ctx._fwd_used_autocast = _is_autocast_enabled(device_type)
+        input, weight, bias = inputs[0], inputs[1], inputs[2] if len(inputs) > 2 else None
+        ctx.save_for_backward(input, weight, bias)
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, grad_output):
+        # Match @custom_bwd semantics: always run backward under the same
+        # autocast state as forward — including explicitly disabling autocast
+        # when forward did not use it, to guard against outer autocast regions.
+        device_type = get_accelerator().device_name()
+        with torch.autocast(device_type=device_type, enabled=ctx._fwd_used_autocast, dtype=ctx._dtype):
+            input, weight, bias = ctx.saved_tensors
+
+            grad_input = grad_weight = grad_bias = None
+
+            dim = grad_output.dim()
+            if ctx.needs_input_grad[0]:
+                grad_input = grad_output.matmul(weight)
+            if ctx.needs_input_grad[1]:
+                if dim > 2:
+                    grad_weight = grad_output.reshape(-1, grad_output.shape[-1]).t().matmul(
+                        input.reshape(-1, input.shape[-1]))
+                else:
+                    grad_weight = grad_output.t().matmul(input)
+            if bias is not None and ctx.needs_input_grad[2]:
+                if dim > 2:
+                    grad_bias = grad_output.sum([i for i in range(dim - 1)])
+                else:
+                    grad_bias = grad_output.sum(0)
+            return grad_input, grad_weight, grad_bias
+
+
+def zero3_linear_wrap(input, weight, bias=None):
+    return LinearFunctionForZeroStage3.apply(input, weight, bias)
+
+
+class LinearModuleForZeroStage3(Module):
+    r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.
+    The weights are pre-transposed and stored as A^T instead of transposing during each
+    forward. Memory savings proportional to the parameter size.
+
+    Args:
+        in_features: size of each input sample
+        out_features: size of each output sample
+        bias: If set to ``False``, the layer will not learn an additive bias.
+            Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, *, H_{in})` where :math:`*` means any number of
+          additional dimensions and :math:`H_{in} = \text{in\_features}`
+        - Output: :math:`(N, *, H_{out})` where all but the last dimension
+          are the same shape as the input and :math:`H_{out} = \text{out\_features}`.
+
+    Attributes:
+        weight: the learnable weights of the module of shape
+            :math:`(\text{out\_features}, \text{in\_features})`. The values are
+            initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
+            :math:`k = \frac{1}{\text{in\_features}}`
+        bias:   the learnable bias of the module of shape :math:`(\text{out\_features})`.
+                If :attr:`bias` is ``True``, the values are initialized from
+                :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                :math:`k = \frac{1}{\text{in\_features}}`
+
+    Examples::
+
+        >>> m = nn.Linear(20, 30)
+        >>> input = torch.randn(128, 20)
+        >>> output = m(input)
+        >>> print(output.size())
+        torch.Size([128, 30])
+    """
+    __constants__ = ['in_features', 'out_features']
+    in_features: int
+    out_features: int
+    weight: Tensor
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super(LinearModuleForZeroStage3, self).__init__()
+        print("Building ZeRO module")
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.Tensor(out_features, in_features))
+        if bias:
+            self.bias = Parameter(torch.Tensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in)
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: Tensor) -> Tensor:
+        return LinearFunctionForZeroStage3.apply(input, self.weight, self.bias)
+
+    def extra_repr(self) -> str:
+        return 'in_features={}, out_features={}, bias={}'.format(self.in_features, self.out_features, self.bias
+                                                                 is not None)

deepspeed/runtime/zero/mics.py

@@ -0,0 +1,447 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List
+
+import deepspeed
+import torch
+from deepspeed import comm as dist
+from deepspeed.runtime.zero.utils import is_zero_param
+from deepspeed.runtime.zero.mics_utils import (MiCS_CommGroups, create_mics_comm_groups, scale_tensors)
+from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload
+from deepspeed.runtime.zero.partition_parameters import Init, AllGatherCoalescedHandle, ZeroParamStatus
+from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+from deepspeed.utils import instrument_w_nvtx, log_dist, logger
+from deepspeed.accelerator import get_accelerator
+from torch import Tensor
+from torch.nn import Parameter
+
+
+def has_hierarchical_all_gather_groups(comm_groups: MiCS_CommGroups):
+    result = False
+    if comm_groups.param_intra_node_group is not None and comm_groups.param_inter_node_shard_group is not None:
+        result = True
+    return result
+
+
+class MiCS_AllGatherCoalescedHandle(AllGatherCoalescedHandle):
+    """ This handle assumes that no need to
+    copy data out from a contiguous tensor
+    """
+
+    def __init__(self, allgather_handle, params: List[Parameter], partitions: List[Tensor], world_size: int) -> None:
+        super().__init__(allgather_handle, params, partitions, world_size)
+
+    def wait(self, **kwargs) -> None:
+        """
+        """
+        # let the current stream to op
+        try:
+            # print("HANDLE", self.allgather_handle)
+            instrument_w_nvtx(self.allgather_handle.wait)()
+        except (ValueError, RuntimeError) as e:
+            log_dist(
+                "WARNING: Runtime Error while waiting the collective all-gather, possibly due to the _IllegalWork",
+                ranks=[0])
+            log_dist(f"Error message: {e}", ranks=[0])
+
+        if self.complete:
+            return
+
+        for _, param in enumerate(self.params):
+            assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight"
+            param.ds_status = ZeroParamStatus.AVAILABLE
+
+        self.complete = True
+
+
+class MiCS_Init(Init):
+
+    def __init__(self,
+                 module=None,
+                 data_parallel_group=None,
+                 sequence_data_parallel_group=None,
+                 mem_efficient_linear=True,
+                 remote_device=None,
+                 pin_memory=False,
+                 config_dict_or_path=None,
+                 config=None,
+                 enabled=True,
+                 dtype=None,
+                 mpu=None):
+        """A context manager to partition the model parameters during the model
+        construction with MiCS partition strategy. Model states are partitioned
+        to the number of devices specified via ``mics_shard_size`` field in the
+        deepspeed config json file. The context manager also introduces
+        hierarchical communication method to reduce the cost of inter-node
+        communications, which can be enabled with
+        ``mics_hierarchical_params_gather`` field in deepspeed config.
+
+        Args:
+            module (``torch.nn.Module``, optional): If provided, partition the model as
+                if it was constructed in the context.
+            data_parallel_group (``deepspeed.comm`` process group, optional):
+                The group of processes to partition among. Defaults to all processes.
+                Synonymous with sequence data parallel group for param partitioning
+                across both sequence and data parallel groups.
+            mem_efficient_linear (bool, optional): Replace
+                torch.nn.functional.linear with an implementation that allows
+                DeepSpeed to partition parameters. Defaults to ``True``.
+            remote_device (string, optional): The initial device to store model
+                weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU
+                memory. The model may still be moved to GPU based on the
+                offload settings for training. Defaults to param offload device if a config is
+                defined, otherwise GPU.
+            pin_memory (bool, optional): Potentially increase performance by
+                using pinned memory for model weights. ``remote_device`` must be
+                ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``.
+            config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration
+                for swapping fp16 params to NVMe.
+            config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead.
+            enabled (bool, optional): If ``False``, this context has no
+                effect. Defaults to ``True``.
+            dtype (``dtype``, optional): Can be used to change the data type of the parameters.
+                Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None``
+            mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}.
+
+        This context follows the same logic as ``deepspeed.zero.Init()``, but
+        with the modification for partition size of each parameter.
+
+        Examples
+        --------
+
+        #. Allocate a model and partition it among all processes:
+
+            .. code-block:: python
+                # the config_dict_or_path is required to let the context manager know
+                # how partition the parameters.
+                # The configuration has to include the field ``mics_shard_size``
+                with deepspeed.zero.MiCS_Init(config_dict_or_path=ds_config):
+                    model = MyLargeModel()
+
+
+        #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes:
+
+            .. code-block:: python
+
+                with deepspeed.zero.MiCS_Init(data_parallel_group=mpu.get_data_parallel_group(),
+                                              remote_device="cpu",
+                                              pin_memory=True
+                                              config_dict_or_path=ds_config):
+                    model = MyLargeModel()
+
+
+        #. Partition an already-allocated model in CPU memory:
+
+            .. code-block:: python
+
+                model = deepspeed.zero.MiCS_Init(module=model,
+                                                 config_dict_or_path=ds_config)
+        """
+
+        assert config_dict_or_path is not None, "Must provide configuration for MiCS Initialization"
+        _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, mpu)
+        if not dist.is_initialized():
+            dist.init_distributed()
+            assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm"
+
+        if data_parallel_group is None:
+            ds_process_group = dist.get_world_group()
+        else:
+            ds_process_group = data_parallel_group
+
+        if sequence_data_parallel_group is not None:
+            logger.warning(
+                "sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.")
+            if data_parallel_group is not None:
+                raise ValueError(
+                    "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments."
+                )
+            self.ds_process_group = sequence_data_parallel_group
+
+        self.mics_comm_groups = create_mics_comm_groups(
+            _ds_config.mics_shard_size,
+            ds_process_group,
+            hierarchical_allgather=_ds_config.mics_hierarchial_params_gather,
+            mpu=mpu)
+
+        super().__init__(module, data_parallel_group, mem_efficient_linear, remote_device, pin_memory,
+                         config_dict_or_path, config, enabled, dtype, mpu)
+
+    def _convert_to_deepspeed_param(self, param):
+        super()._convert_to_deepspeed_param(param)
+        # attach communication groups to every param
+        param.comm = self.mics_comm_groups
+
+        # record existing all_gather_coalesced implementation
+        # so that we can fallback later
+        old_all_gather_coalesced = param.all_gather_coalesced
+
+        def _param_all_gather_coalesced(params, param_buffers=None, **kwargs):
+            """"""
+            mics_comm_groups: MiCS_CommGroups = params[0].comm
+            hierarchical_all_gather = has_hierarchical_all_gather_groups(mics_comm_groups)
+            if dist.has_coalescing_manager() and hierarchical_all_gather:
+                return self._hierarchical_all_gather_params(params, param_buffers)
+            elif dist.has_coalescing_manager():
+                return self._flat_all_gather_with_coalescing_manager(params, param_buffers)
+            else:
+                return old_all_gather_coalesced(params, **kwargs)
+
+        # change the all_gather_coalesced method
+        param.all_gather_coalesced = _param_all_gather_coalesced
+
+    def _pre_all_gather(self, params, params_buffers=None):
+        # fetches from nvme if the partition is not available and in nvme
+        self._ensure_availability_of_partitioned_params(params)
+
+        for param in params:
+            if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
+                raise RuntimeError(param.ds_summary())
+            param.ds_status = ZeroParamStatus.INFLIGHT
+
+        # ensure that each rank has params in same order. the allgather
+        # is done by flattening the parameter list into a single tensor that
+        # can be allgathered in a single call - this means that if each rank
+        # gives a list of the same parameters in a different order we will
+        # silently get incorrect parameter values, and have very difficult
+        # to debug correctness issues.
+        params = sorted(params, key=lambda p: p.ds_id)
+        return params, params_buffers
+
+    def _flat_all_gather_with_coalescing_manager(self, params, params_buffers=None):
+        """"""
+        # must have to change the status of the param
+        # and ensure they are on the device
+        params, params_buffers = self._pre_all_gather(params, params_buffers)
+
+        mics_comm_groups: MiCS_CommGroups = params[0].comm
+        param_shard_size = mics_comm_groups.param_shard_size
+
+        output_tensors = []
+        input_tensors = []
+        for i, p in enumerate(params):
+            t_size = p.ds_tensor.ds_numel * param_shard_size
+            if params_buffers is not None and params_buffers[i] is not None:
+                assert params_buffers[i].numel(
+                ) == t_size, f'params_to_gather_buffers[{i}] size {params_buffers[i].numel()} does not match with t_size {t_size}'
+                flat_out = params_buffers[i]
+            else:
+                flat_out = torch.empty(t_size, dtype=p.dtype, device=self.local_device, requires_grad=False).view(-1)
+            output_tensors.append(flat_out)
+            _flat_input = p.ds_tensor.data.view(-1)
+            input_tensors.append(_flat_input)
+
+        all_gather_handle = dist.all_gather_coalesced(output_tensors,
+                                                      input_tensors,
+                                                      group=mics_comm_groups.param_shard_group,
+                                                      async_op=True)
+
+        for idx, param in enumerate(params):
+            param.data = output_tensors[idx].narrow(0, 0, param.ds_numel).view(param.ds_shape).data
+
+        return MiCS_AllGatherCoalescedHandle(allgather_handle=all_gather_handle,
+                                             params=params,
+                                             partitions=[],
+                                             world_size=param_shard_size)
+
+    def _hierarchical_all_gather_params(self, params, params_buffers=None):
+        """"""
+        params, params_buffers = self._pre_all_gather(params, params_buffers)
+
+        mics_comm_groups: MiCS_CommGroups = params[0].comm
+        local_rank = dist.get_rank(group=mics_comm_groups.param_intra_node_group)
+        inter_node_comm_group = mics_comm_groups.param_inter_node_shard_group
+        intra_node_comm_group = mics_comm_groups.param_intra_node_group
+        param_shard_size = mics_comm_groups.param_shard_size
+
+        inter_node_size = dist.get_world_size(group=inter_node_comm_group)
+        intra_node_size = dist.get_world_size(group=intra_node_comm_group)
+        param_tensors = []
+        for i, p in enumerate(params):
+            param_size = p.ds_tensor.ds_numel * param_shard_size
+            if params_buffers is not None and params_buffers[i] is not None:
+                assert params_buffers[i].numel(
+                ) == param_size, f'param_buffers[{i}] size {params_buffers[i].numel()} does not match with param_size {param_size}'
+                param_tensor = params_buffers[i]
+            else:
+                param_tensor = torch.empty(param_size, dtype=p.dtype, device=self.local_device,
+                                           requires_grad=False).view(-1)
+            param_tensors.append(param_tensor)
+
+        # inter node all-gather
+        inter_outputs = []
+        inter_inputs = []
+        for i, p in enumerate(params):
+            inter_size = p.ds_tensor.ds_numel * inter_node_size
+            _out = param_tensors[i].narrow(0, local_rank * inter_size, inter_size)
+            inter_outputs.append(_out)
+            inter_inputs.append(p.ds_tensor.data.view(-1).to(self.local_device))
+        # sync enqueue
+        dist.all_gather_coalesced(inter_outputs, inter_inputs, group=inter_node_comm_group, async_op=False)
+
+        # intra node all-gather
+        intra_outputs = []
+        intra_inputs = []
+        for i, p in enumerate(params):
+            # partition param into multiple chunks for allgather
+            # because inter-node all-gather outputs are in a continues memory
+            # while in param memory, those inter-node data are placed in different
+            # location.
+            # each chunk is an intra-node output
+            param_chunk = param_tensors[i].view(
+                (inter_node_size, intra_node_size, p.ds_tensor.ds_numel)).narrow(1, local_rank, 1)
+            param_chunk.copy_(inter_outputs[i].detach().clone().view(param_chunk.size()))
+            output_chunks = torch.chunk(param_tensors[i], inter_node_size)
+            for j, _out in enumerate(output_chunks):
+                intra_chunk_size = intra_node_size * p.ds_tensor.ds_numel
+                local_offset = local_rank * p.ds_tensor.ds_numel
+                _in = param_tensors[i].narrow(0, j * intra_chunk_size + local_offset, p.ds_tensor.ds_numel)
+                intra_outputs.append(_out)
+                intra_inputs.append(_in)
+
+        all_gather_handle = dist.all_gather_coalesced(intra_outputs,
+                                                      intra_inputs,
+                                                      group=intra_node_comm_group,
+                                                      async_op=True)
+        for i, param in enumerate(params):
+            param.data = param_tensors[i].narrow(0, 0, param.ds_numel).view(param.ds_shape).data
+
+        return MiCS_AllGatherCoalescedHandle(
+            allgather_handle=all_gather_handle,
+            params=params,
+            partitions=[],
+            world_size=param_shard_size,
+        )
+
+    def get_partition_dp_group(self, param):
+        return param.comm.param_shard_group
+
+    def get_partition_rank(self):
+        return self.mics_comm_groups.param_shard_rank
+
+    @property
+    def num_partitions(self):
+        return self.mics_comm_groups.param_shard_size
+
+
+class MiCS_Offload(DeepSpeedZeRoOffload):
+    """ Wrapper to change the behavior for parameter sharding
+    """
+
+    def _convert_to_zero_parameters(self, ds_config, module, mpu):
+        """ overload the parent class function for convert the parameters
+
+        """
+        log_dist('Convert to zero parameters from MiCS Offload manager', ranks=[0])
+        non_zero_params = [p for p in module.parameters() if not is_zero_param(p)]
+        if non_zero_params:
+            zero_params = [p for p in module.parameters() if is_zero_param(p)]
+            if zero_params:
+                zero_params[0].convert_to_zero_parameters(param_list=non_zero_params)
+            else:
+                group = None
+                if mpu:
+                    group = mpu.get_data_parallel_group()
+
+                MiCS_Init(module=module,
+                          data_parallel_group=group,
+                          dtype=self.dtype,
+                          config_dict_or_path=ds_config,
+                          remote_device=self.offload_device,
+                          pin_memory=self.offload_param_pin_memory,
+                          mpu=mpu)
+
+
+class MiCS_Optimizer(DeepSpeedZeroOptimizer_Stage3):
+    """
+    MiCS Optimizer
+    """
+
+    def __init__(self,
+                 module,
+                 init_optimizer,
+                 param_names,
+                 timers,
+                 ds_config,
+                 gradient_accumulation_dtype=torch.float16,
+                 **kwargs):
+
+        log_dist("Init MiCS optimizer", ranks=[0])
+        super().__init__(module,
+                         init_optimizer,
+                         param_names,
+                         timers,
+                         ds_config,
+                         gradient_accumulation_dtype=gradient_accumulation_dtype,
+                         **kwargs)
+        first_param = next(module.parameters())
+        # overload the dp_process_group and partition_count
+        assert hasattr(first_param, "comm"), " ".join([
+            "Sharded parameters don't have the MiCS_CommGroups attached.",
+            "Might due to the use of deepspeed.zero.Init context for initializing the weights.",
+            "To use MiCS sharding, please use deepspeed.zero.MiCS_Init instead for initializing parameter."
+        ])
+        self.dp_process_group = first_param.comm.param_shard_group
+        self.partition_count = first_param.comm.param_shard_size
+
+    def initialize_ds_offload(
+        self,
+        *args,
+        **kwargs,
+    ):
+        return MiCS_Offload(*args, **kwargs)
+
+    def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None:
+        grad_buffers = super().partition_grads(params_to_release, grad_partitions)
+        # perform all-reduce among replication groups
+        # the function will perform accumulation boundary check
+        self.allreduce_mics_shard_grads(params_to_release, grad_buffers)
+
+    @instrument_w_nvtx
+    def allreduce_mics_shard_grads(self, params, partitioned_grads_buffers: List[Tensor]):
+        """
+        """
+        # TODO: improve the condition check
+        if not self.is_gradient_accumulation_boundary or \
+            len(partitioned_grads_buffers) == 0:
+            return
+
+        mics_comm_groups: MiCS_CommGroups = params[0].comm
+        param_repli_group = mics_comm_groups.param_repli_group
+        param_repli_size = mics_comm_groups.param_repli_size
+
+        if param_repli_size is None or param_repli_size <= 1:
+            return
+        if not get_accelerator().on_accelerator(partitioned_grads_buffers[0]):
+            raise RuntimeError("Local sharding has no support for CPU offloading")
+
+        if dist.has_all_reduce_coalesced():
+            scale_tensors(partitioned_grads_buffers, param_repli_size)
+            dist.all_reduce_coalesced(tensors=partitioned_grads_buffers, group=param_repli_group)
+        else:
+            # manually coalescing all-reduce
+            aggregated_buffer: Tensor = torch.cat(partitioned_grads_buffers)
+            aggregated_buffer.div_(param_repli_size)
+            dist.all_reduce(aggregated_buffer, group=param_repli_group)
+            offset = 0
+            for grad_buff in partitioned_grads_buffers:
+                grad_buff.view(-1).copy_(aggregated_buffer.narrow(0, offset, grad_buff.numel()))
+                offset += grad_buff.numel()
+
+    def load_state_dict(self,
+                        state_dict_list,
+                        load_optimizer_states=True,
+                        load_from_fp32_weights=False,
+                        checkpoint_folder=None,
+                        load_serial=None):
+        r""" Loading the ZeRO-3/MiCS partitioned checkpoints
+        Because the self.dp_process_group is replaced with the communicator for
+        partition group we can call the load_state_dict logic from ZeRO-3.
+        """
+        super().load_state_dict(state_dict_list, load_optimizer_states, load_from_fp32_weights, checkpoint_folder)

deepspeed/runtime/zero/mics_utils.py

@@ -0,0 +1,203 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+from torch import Tensor
+
+from deepspeed import comm as dist
+from deepspeed.accelerator import get_accelerator
+from deepspeed.utils import logger
+from deepspeed.utils.torch import jit_script_compat
+
+
+def _log_rank0(msg):
+    if dist.get_rank() == 0:
+        logger.info(msg)
+
+
+@jit_script_compat
+def scale_tensors(tensors: List[Tensor], scale: int):
+    for t in tensors:
+        t.div_(scale)
+
+
+@dataclass
+class MiCS_CommGroups:
+    """"""
+    param_shard_group = None
+    param_shard_size = -1
+    param_shard_rank = -1
+
+    param_repli_group = None
+    param_repli_size = -1
+    param_repli_rank = -1
+
+    param_intra_node_group = None
+    param_inter_node_shard_group = None
+
+
+def create_mics_comm_groups(
+    shard_size,
+    dp_group,
+    hierarchical_allgather=False,
+    mpu=None,
+):
+    """
+    create shard-group, replicate-group from config_file
+    TODO: consider broadcast the config from rank0
+
+    Returns:
+        MiCS_CommGroups
+    """
+    # env var for debugging purpose
+    ndevices_per_node = int(os.environ.get("NDEV_PER_NODE", get_accelerator().device_count()))
+    _log_rank0(f'creating MiCS communication groups with per node device size {ndevices_per_node}')
+    groups = MiCS_CommGroups()
+
+    if mpu is not None:
+        assert dp_group == mpu.get_data_parallel_group()
+
+    # full size of the world
+    world_size = dist.get_world_size()
+    # global rank
+    global_rank = dist.get_rank()
+
+    config = _generate_mics_config(world_size, ndevices_per_node, shard_size, 1)
+    ranks_of_shard_group = config['shard_groups']
+    ranks_of_repli_group = config['replicate_groups']
+    if len(ranks_of_repli_group) == 0:
+        assert len(ranks_of_shard_group) == 1, "replicate groups are empty only for single shard group"
+        for r in ranks_of_shard_group[0]:
+            ranks_of_repli_group.append([r])
+
+    # for simplicity
+    assert _sizes_all_same(ranks_of_repli_group), "replicate groups must have the same size"
+    assert _sizes_all_same(ranks_of_shard_group), "shard groups must have the same size"
+
+    assert sum([len(g) for g in ranks_of_shard_group]) == dist.get_world_size(), "all sharded ranks "
+    if len(ranks_of_shard_group) > 1:  # if only shard on one group then no need for replicate groups
+        assert len(ranks_of_shard_group) == len(
+            ranks_of_repli_group[0]), "number of shard groups must equal to the size of each replicate group"
+
+    global_rank = dist.get_rank()
+    # create shard groups
+    for shard_ranks in ranks_of_shard_group:
+        _group = dist.new_group(shard_ranks)
+        if global_rank in shard_ranks:
+            groups.param_shard_group = _group
+            groups.param_shard_size = len(shard_ranks)
+            groups.param_shard_rank = dist.get_rank(_group)
+            logger.info(f'rank {global_rank}, shard group'
+                        f' {groups.param_shard_rank}/{dist.get_world_size(group=_group)}')
+
+    # create replicate groups
+    for repli_ranks in ranks_of_repli_group:
+        if len(repli_ranks) > 1:
+            _group = dist.new_group(repli_ranks)
+            if global_rank in repli_ranks:
+                groups.param_repli_group = _group
+                groups.param_repli_size = len(repli_ranks)
+                groups.param_repli_rank = dist.get_rank(group=_group)
+                logger.info(f'rank {global_rank} '
+                            f'replicate group {groups.param_repli_rank}/{dist.get_world_size(group=_group)}')
+        else:
+            groups.param_repli_group = None
+            groups.param_repli_size = 1
+            groups.param_repli_rank = 0
+            logger.info(f'rank {global_rank} replicate group 0/1')
+
+    # assign shard group size as world size
+    assert groups.param_shard_size == len(ranks_of_shard_group[0])
+
+    if hierarchical_allgather:
+        # create hierarchy inter-node, intra-node groups
+        # n_span_nodes = config['shard_span']
+        n_span_nodes = config['span_nodes']
+        assert n_span_nodes > 1, "sharding spans on single node, no need for hierarchy allgather"
+        assert len(ranks_of_shard_group[0]) % n_span_nodes == 0
+
+        n_gpu_per_node = len(ranks_of_shard_group[0]) // n_span_nodes
+        intra_node_ranks_group = []
+        inter_node_ranks_group = []
+        for shard_group in ranks_of_shard_group:
+            _intra_node_ranks = []
+            for i in range(0, len(shard_group), n_gpu_per_node):
+                _intra_node_ranks.append(shard_group[i:i + n_gpu_per_node])
+            _inter_node_ranks = []
+            for i in range(n_gpu_per_node):
+                _ranks = [_g[i] for _g in _intra_node_ranks]
+                _inter_node_ranks.append(_ranks)
+
+            intra_node_ranks_group.append(_intra_node_ranks)
+            inter_node_ranks_group.append(_inter_node_ranks)
+
+        _log_rank0(f"create for hierarchy all-gather groups: intra nodes {intra_node_ranks_group}")
+        _log_rank0(f"create for hierarchy all-gather groups: inter nodes {inter_node_ranks_group}")
+
+        # create communicators
+        for shard_group in intra_node_ranks_group:
+            for intra_node_ranks in shard_group:
+                _group = dist.new_group(intra_node_ranks)
+                if global_rank in intra_node_ranks:
+                    groups.param_intra_node_group = _group
+                _log_rank0(f'create group for intra node ranks {intra_node_ranks}')
+
+        for shard_group in inter_node_ranks_group:
+            for inter_node_ranks in shard_group:
+                _group = dist.new_group(inter_node_ranks)
+                if global_rank in inter_node_ranks:
+                    groups.param_inter_node_shard_group = _group
+                _log_rank0(f'create group for inter node ranks {inter_node_ranks}')
+    return groups
+
+
+def _generate_mics_config(world_size, ndev_per_node, shard_size, pp_size=1):
+    """Generating the configuration for sharding This shard config generation assume
+    that the pipeline stages are partitioned in order, i.e., first ranks
+    hold the stage0, etc.
+
+    Args:
+
+        shard_size (int): zero3 data-parallel shard size, FIXME:
+        change the name later
+
+        pp_size (int): pipeline parallel size, currently, only work with
+        pipeline parallelism + zero
+
+    """
+    assert world_size % pp_size == 0
+    assert (world_size // pp_size) % shard_size == 0, \
+        f"dp group size is not dividable by dp_shard_size, "\
+        f" (world_size {world_size}, pp_size {pp_size}, dp_shard_size {shard_size})"
+
+    config = {}
+    shard_groups = np.arange(world_size).reshape(-1, shard_size)
+    replicate_groups = []
+    for i in range(shard_size):
+        same_shard_ranks = shard_groups[:, i].tolist()
+        n_ranks = len(same_shard_ranks)
+        replicate_size = n_ranks // pp_size
+        replicate_groups.extend([same_shard_ranks[j:j + replicate_size] for j in range(0, n_ranks, replicate_size)])
+
+    config['replicate_groups'] = replicate_groups
+    config['shard_groups'] = shard_groups.tolist()
+    config["span_nodes"] = len(shard_groups[0]) // ndev_per_node
+    return config
+
+
+def _sizes_all_same(groups):
+    """all groups have same length"""
+    all_same = True
+    for g in groups:
+        if len(g) != len(groups[0]):
+            return False
+    return all_same

deepspeed/runtime/zero/muon/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2025 Peng Du and Zhipeng Wang
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team

deepspeed/runtime/zero/muon/muon_optimizer.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 Peng Du and Zhipeng Wang
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+try:
+    from deepspeed.runtime.zero.muon.original_muon import MuonWithAuxAdam as BaseMuonWithAuxAdam
+    from deepspeed.runtime.zero.muon.original_muon import adam_update
+except ImportError:
+    pass
+
+
+class MuonWithAuxAdam(BaseMuonWithAuxAdam):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        for group in self.param_groups:
+            if group["use_muon"]:
+                # we move the muon update part to the deepspeed's optimizer since the parameter here is a flat version
+                # thus not suitable for muon update
+                for p in group["params"]:
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(p.grad.reshape(p.shape), alpha=-group["lr"])
+            else:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+                    state["step"] += 1
+                    update = adam_update(p.grad, state["exp_avg"], state["exp_avg_sq"], state["step"], group["betas"],
+                                         group["eps"])
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss

deepspeed/runtime/zero/muon/original_muon.py

@@ -0,0 +1,443 @@
+# Copyright (c) 2024 Keller Jordan
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+MIT License
+
+Copyright (c) 2024 Keller Jordan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+"""
+
+import torch
+import deepspeed.comm as dist  # replace torch's distributed package with deepspeed.comm to resolve deepspeed check
+from deepspeed.runtime import compiler
+from deepspeed.accelerator import get_accelerator
+
+
+@compiler.compile()
+def zeropower_via_newtonschulz5(G, steps: int):
+    """
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    assert G.ndim >= 2  # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    # Use bf16 when hardware supports it; fp32 otherwise
+    compute_dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32
+    X = G.to(compute_dtype)
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+    # Perform the NS iterations
+    for _ in range(steps):
+        A = X @ X.mT
+        B = b * A + c * A @ A  # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    return X
+
+
+@compiler.compile()
+def zeropower_via_gram_newtonschulz(G, steps: int):
+    """
+    Gram Newton-Schulz iteration for orthogonalization.
+
+    Mathematically equivalent to standard Newton-Schulz but iterates on the
+    small square Gram matrix R = X @ X.T (n x n) instead of the full rectangular
+    X (n x m). This reduces FLOPs significantly when m >> n (typical for
+    transformer weight matrices with aspect ratio ~5).
+
+    Uses fp16 instead of bf16 for better numerical precision at the same
+    compute cost. Includes a restart at iteration 2 to maintain stability
+    in half-precision.
+
+    Falls back to standard Newton-Schulz for square matrices (n == m)
+    where there is no FLOP advantage.
+
+    Reference: https://tridao.me/blog/2026/gram-newton-schulz/
+    """
+    assert G.ndim >= 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    # Use fp16 for better precision than bf16 when hardware supports it; fp32 otherwise
+    compute_dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32
+    X = G.to(compute_dtype)
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    n, m = X.size(-2), X.size(-1)
+
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+
+    # For square matrices, no FLOP advantage; use standard iteration
+    if m <= n:
+        for _ in range(steps):
+            A = X @ X.mT
+            B = b * A + c * A @ A
+            X = a * X + B @ X
+        if G.size(-2) > G.size(-1):
+            X = X.mT
+        return X
+
+    # Gram NS: iterate on R = X @ X.T (n x n) instead of X (n x m)
+    R = X @ X.mT
+    Q = None
+    restart_at = 2
+
+    for i in range(steps):
+        if i == restart_at and i != 0:
+            X = Q @ X
+            R = X @ X.mT
+            Q = None
+
+        Z = b * R + c * R @ R
+
+        if Q is None:
+            Q = Z.clone()
+            if Q.ndim == 2:
+                Q.diagonal().add_(a)
+            else:
+                Q.diagonal(dim1=-2, dim2=-1).add_(a)
+        else:
+            Q = a * Q + Z @ Q
+
+        if i < steps - 1 and (i + 1) != restart_at:
+            RZ = a * R + Z @ R
+            R = a * RZ + Z @ RZ
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT @ Q.mT
+    else:
+        X = Q @ X
+    return X
+
+
+NS_METHODS = {"standard", "gram"}
+
+
+@compiler.compile()
+def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram", is_expert_group=False):
+    orig_dtype = grad.dtype
+    momentum.lerp_(grad, 1 - beta)
+    update = grad.lerp_(momentum, beta) if nesterov else momentum
+    if is_expert_group:
+        ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5
+        scale = max(1, update.size(-2) / update.size(-1))**0.5
+        update = ns_fn(update, steps=ns_steps) * scale
+    else:
+        if update.ndim == 4:  # for the case of conv filters
+            update = update.view(len(update), -1)
+        if ns_method == "gram":
+            update = zeropower_via_gram_newtonschulz(update, steps=ns_steps)
+        else:
+            update = zeropower_via_newtonschulz5(update, steps=ns_steps)
+        update *= max(1, grad.size(-2) / grad.size(-1))**0.5
+    if update.dtype != orig_dtype:
+        update = update.to(orig_dtype)
+    return update
+
+
+class Muon(torch.optim.Optimizer):
+    """
+    Muon - MomentUm Orthogonalized by Newton-schulz
+
+    https://kellerjordan.github.io/posts/muon/
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+    matrix. For efficient orthogonalization we use a Newton-Schulz iteration, which has the
+    advantage that it can be stably run in bfloat16 on the GPU.
+
+    Muon should only be used for hidden weight layers. The input embedding, final output layer,
+    and any internal gains or biases should be optimized using a standard method such as AdamW.
+    Hidden convolutional weights can be trained using Muon by viewing them as 2D and then
+    collapsing their last 3 dimensions.
+
+    Arguments:
+        lr: The learning rate, in units of spectral norm per update.
+        weight_decay: The AdamW-style weight decay.
+        momentum: The momentum. A value of 0.95 here is usually fine.
+        ns_method: Newton-Schulz method. "gram" (default) uses Gram NS for ~2x speedup
+                   on rectangular matrices. "standard" uses the original iteration.
+    """
+
+    def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95, ns_method="gram"):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_method=ns_method)
+        assert isinstance(params, list) and len(params) >= 1 and isinstance(params[0], torch.nn.Parameter)
+        params = sorted(params, key=lambda x: x.size(), reverse=True)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params = group["params"]
+            params_pad = params + [torch.empty_like(params[-1])
+                                   ] * (dist.get_world_size() - len(params) % dist.get_world_size())
+            for base_i in range(len(params))[::dist.get_world_size()]:
+                if base_i + dist.get_rank() < len(params):
+                    p = params[base_i + dist.get_rank()]
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+                    update = muon_update(p.grad,
+                                         state["momentum_buffer"],
+                                         beta=group["momentum"],
+                                         ns_method=group.get("ns_method", "gram"),
+                                         is_expert_group=getattr(p, 'is_expert_group', False))
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update.reshape(p.shape), alpha=-group["lr"])
+                dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()],
+                                params_pad[base_i + dist.get_rank()])
+
+        return loss
+
+
+class SingleDeviceMuon(torch.optim.Optimizer):
+    """
+    Muon variant for usage in non-distributed settings.
+    """
+
+    def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95, ns_method="gram"):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_method=ns_method)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    # continue
+                    p.grad = torch.zeros_like(p)  # Force synchronization
+                state = self.state[p]
+                if len(state) == 0:
+                    state["momentum_buffer"] = torch.zeros_like(p)
+                update = muon_update(p.grad,
+                                     state["momentum_buffer"],
+                                     beta=group["momentum"],
+                                     ns_method=group.get("ns_method", "gram"),
+                                     is_expert_group=getattr(p, 'is_expert_group', False))
+                p.mul_(1 - group["lr"] * group["weight_decay"])
+                p.add_(update.reshape(p.shape), alpha=-group["lr"])
+
+        return loss
+
+
+def adam_update(grad, buf1, buf2, step, betas, eps):
+    buf1.lerp_(grad, 1 - betas[0])
+    buf2.lerp_(grad.square(), 1 - betas[1])
+    buf1c = buf1 / (1 - betas[0]**step)
+    buf2c = buf2 / (1 - betas[1]**step)
+    return buf1c / (buf2c.sqrt() + eps)
+
+
+class MuonWithAuxAdam(torch.optim.Optimizer):
+    """
+    Distributed Muon variant that can be used for all parameters in the network, since it runs an
+    internal AdamW for the parameters that are not compatible with Muon. The user must manually
+    specify which parameters shall be optimized with Muon and which with Adam by passing in a
+    list of param_groups with the `use_muon` flag set.
+
+    The point of this class is to allow the user to have a single optimizer in their code, rather
+    than having both a Muon and an Adam which each need to be stepped.
+
+    You can see an example usage below:
+
+    https://github.com/KellerJordan/modded-nanogpt/blob/master/records/052525_MuonWithAuxAdamExample/b01550f9-03d8-4a9c-86fe-4ab434f1c5e0.txt#L470
+    ```
+    hidden_matrix_params = [p for n, p in model.blocks.named_parameters() if p.ndim >= 2 and "embed" not in n]
+    embed_params = [p for n, p in model.named_parameters() if "embed" in n]
+    scalar_params = [p for p in model.parameters() if p.ndim < 2]
+    head_params = [model.lm_head.weight]
+
+    from muon import MuonWithAuxAdam
+    adam_groups = [dict(params=head_params, lr=0.22), dict(params=embed_params, lr=0.6), dict(params=scalar_params, lr=0.04)]
+    adam_groups = [dict(**g, betas=(0.8, 0.95), eps=1e-10, use_muon=False) for g in adam_groups]
+    muon_group = dict(params=hidden_matrix_params, lr=0.05, momentum=0.95, use_muon=True)
+    param_groups = [*adam_groups, muon_group]
+    optimizer = MuonWithAuxAdam(param_groups)
+    ```
+    """
+
+    def __init__(self, param_groups):
+        for group in param_groups:
+            assert "use_muon" in group
+            if group["use_muon"]:
+                group["params"] = sorted(group["params"], key=lambda x: x.size(), reverse=True)
+                # defaults
+                group["lr"] = group.get("lr", 0.02)
+                group["momentum"] = group.get("momentum", 0.95)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                group["ns_method"] = group.get("ns_method", "gram")
+                assert group[
+                    "ns_method"] in NS_METHODS, f"ns_method must be one of {NS_METHODS}, got {group['ns_method']}"
+                assert set(["params", "lr", "momentum", "weight_decay", "use_muon",
+                            "ns_method"]).issubset(set(group.keys()))
+            else:
+                # defaults
+                group["lr"] = group.get("lr", 3e-4)
+                group["betas"] = group.get("betas", (0.9, 0.95))
+                group["eps"] = group.get("eps", 1e-10)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                assert set(["params", "lr", "betas", "eps", "weight_decay", "use_muon"]).issubset(set(group.keys()))
+        super().__init__(param_groups, dict())
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if group["use_muon"]:
+                params = group["params"]
+                params_pad = params + [torch.empty_like(params[-1])
+                                       ] * (dist.get_world_size() - len(params) % dist.get_world_size())
+                for base_i in range(len(params))[::dist.get_world_size()]:
+                    if base_i + dist.get_rank() < len(params):
+                        p = params[base_i + dist.get_rank()]
+                        if p.grad is None:
+                            # continue
+                            p.grad = torch.zeros_like(p)  # Force synchronization
+                        state = self.state[p]
+                        if len(state) == 0:
+                            state["momentum_buffer"] = torch.zeros_like(p)
+                        update = muon_update(p.grad,
+                                             state["momentum_buffer"],
+                                             beta=group["momentum"],
+                                             ns_method=group.get("ns_method", "gram"),
+                                             is_expert_group=getattr(p, 'is_expert_group', False))
+                        p.mul_(1 - group["lr"] * group["weight_decay"])
+                        p.add_(update.reshape(p.shape), alpha=-group["lr"])
+                    dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()],
+                                    params_pad[base_i + dist.get_rank()])
+            else:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+                    state["step"] += 1
+                    update = adam_update(p.grad, state["exp_avg"], state["exp_avg_sq"], state["step"], group["betas"],
+                                         group["eps"])
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss
+
+
+class SingleDeviceMuonWithAuxAdam(torch.optim.Optimizer):
+    """
+    Non-distributed variant of MuonWithAuxAdam.
+    """
+
+    def __init__(self, param_groups):
+        for group in param_groups:
+            assert "use_muon" in group
+            if group["use_muon"]:
+                # defaults
+                group["lr"] = group.get("lr", 0.02)
+                group["momentum"] = group.get("momentum", 0.95)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                group["ns_method"] = group.get("ns_method", "gram")
+                assert group[
+                    "ns_method"] in NS_METHODS, f"ns_method must be one of {NS_METHODS}, got {group['ns_method']}"
+                assert set(["params", "lr", "momentum", "weight_decay", "use_muon",
+                            "ns_method"]).issubset(set(group.keys()))
+            else:
+                # defaults
+                group["lr"] = group.get("lr", 3e-4)
+                group["betas"] = group.get("betas", (0.9, 0.95))
+                group["eps"] = group.get("eps", 1e-10)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                assert set(["params", "lr", "betas", "eps", "weight_decay", "use_muon"]).issubset(set(group.keys()))
+        super().__init__(param_groups, dict())
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if group["use_muon"]:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+                    update = muon_update(p.grad,
+                                         state["momentum_buffer"],
+                                         beta=group["momentum"],
+                                         ns_method=group.get("ns_method", "gram"),
+                                         is_expert_group=getattr(p, 'is_expert_group', False))
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update.reshape(p.shape), alpha=-group["lr"])
+            else:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+                    state["step"] += 1
+                    update = adam_update(p.grad, state["exp_avg"], state["exp_avg_sq"], state["step"], group["betas"],
+                                         group["eps"])
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss

deepspeed/runtime/zero/offload_config.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from enum import Enum
+from pathlib import Path
+from pydantic import Field, model_validator
+from typing import Optional
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int
+
+
+class OffloadDeviceEnum(str, Enum):
+    """ Enum for valid offload devices """
+    none = "none"
+    cpu = "cpu"
+    nvme = "nvme"
+
+
+class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel):
+    """ Set options for parameter offload. Valid only with stage 3. """
+
+    device: OffloadDeviceEnum = "none"
+    """
+    Device memory to offload model parameters. Supported options are `cpu` and
+    `nvme`.
+    """
+
+    nvme_path: Optional[Path] = None
+    """ Filesystem path for NVMe device for parameter offloading. """
+
+    buffer_count: int = Field(5, ge=0)
+    """ Number of buffers in buffer pool for parameter offloading to NVMe. """
+
+    buffer_size: int = Field(pp_int(1e8), ge=0)
+    """ Size of buffers in buffer pool for parameter offloading to NVMe. """
+
+    max_in_cpu: int = Field(pp_int(1e9), ge=0)
+    """
+    Number of parameter elements to maintain in CPU memory when offloading to
+    NVMe is enabled.
+    """
+
+    pin_memory: bool = False
+    """
+    Offload to page-locked CPU memory. This could boost throughput at the cost
+    of extra memory overhead.
+    """
+
+
+class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel):
+    """ Set options for optimizer offload. Valid with stage 1, 2, and 3. """
+
+    device: OffloadDeviceEnum = "none"
+    """
+    Device memory to offload optimizer state. Supported options are `cpu` and
+    `nvme`. Optimizer computation is offload to CPU regardless of device option.
+    """
+
+    nvme_path: Optional[Path] = None
+    """ Filesystem path for NVMe device for optimizer state offloading. """
+
+    buffer_count: int = Field(4, ge=0)
+    """
+    Number of buffers in buffer pool for optimizer state offloading to NVMe.
+    This should be at least the number of states maintained per parameter by
+    the optimizer. For example, Adam optimizer has 4 states (parameter,
+    gradient, momentum, and variance).
+    """
+
+    pin_memory: bool = False
+    """
+    Offload to page-locked CPU memory. This could boost throughput at the cost
+    of extra memory overhead.
+    """
+
+    pipeline_read: bool = False
+    """
+    For tile-based optimizer step processing, overlap read of next tile with
+    computation of current tile. Used in ZeRO-Infinity.
+    """
+
+    pipeline_write: bool = False
+    """
+    For tile-based optimizer step processing, overlap write of previous tile
+    with computation of current tile.
+    """
+
+    fast_init: bool = False
+    """ Enable fast optimizer initialization when offloading to NVMe. """
+
+    ratio: float = Field(1.0, ge=0.0, le=1.0)
+    """ Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3."""
+
+    super_offload: bool = False
+    """ Enable high performance CPU offloading for Superchips. Only valid with ZeRO Stage 3."""
+
+    cpuadam_cores_perc: float = Field(0.8, ge=0.0, le=1.0)
+    """ Percentage of CPU cores to use for CPU Adam. Only valid with ZeRO Stage 3 and super_offload=True."""
+
+    @model_validator(mode="after")
+    def set_pipeline(self):
+        pipeline = self.pipeline_read or self.pipeline_write
+        self.__dict__["pipeline"] = pipeline
+        return self
+
+
+class OffloadStateTypeEnum(str, Enum):
+    """ Enum for internal buffer types """
+    optim_states = "optim_states"
+    hp_params = "hp_params"
+    lp_params = "lp_params"
+    lp_grads = "lp_grads"
+    contiguous_grad_buffer = "contiguous_grad_buffer"

deepspeed/runtime/zero/offload_states.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Set
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum
+
+
+def _make_offload_state_key(key):
+    return f"{key}_offload_buffer"
+
+
+def offload_optimizer_states(optimizer, device, pin_memory=False, non_blocking=False):
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if torch.is_tensor(v):
+                if pin_memory and v.device.type != 'cpu':
+                    pinned_buffer = torch.empty_like(v, device='cpu').pin_memory()
+                    pinned_buffer.copy_(v, non_blocking=non_blocking)
+                    state[k] = pinned_buffer
+                else:
+                    state[k] = v.to(device, non_blocking=non_blocking)
+
+
+def reload_optimizer_states(optimizer, device, non_blocking=False):
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if torch.is_tensor(v):
+                state[k] = v.to(device, non_blocking=non_blocking)
+
+
+def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False):
+    """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam."""
+
+    def move_key(state, key):
+        offload_buf_key = _make_offload_state_key(key)
+        if offload_buf_key not in state:
+            state[offload_buf_key] = torch.empty_like(state[key], device=device)
+            if pin_memory:
+                state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key])
+        state[offload_buf_key].copy_(state[key], non_blocking=non_blocking)
+        state[key].data = state[offload_buf_key]
+
+    for _, state in optimizer.state.items():
+        if "exp_avg" in state:
+            move_key(state, "exp_avg")
+        if "exp_avg_sq" in state:
+            move_key(state, "exp_avg_sq")
+
+
+def reload_adam_states(optimizer, device, non_blocking: bool = False):
+    """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam."""
+
+    def move_back_key(state, key):
+        state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking)
+
+    for _, state in optimizer.state.items():
+        if "exp_avg" in state:
+            move_back_key(state, "exp_avg")
+        if "exp_avg_sq" in state:
+            move_back_key(state, "exp_avg_sq")
+
+
+def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]:
+    """Retrieve the devices of the specified state of the model.
+
+    Args:
+        model (DeepSpeedEngine): The model whose device allocations are to be checked.
+        state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved.
+
+    Returns:
+        Set[torch.device]: A set of devices of the specified state.
+
+    """
+    if state == OffloadStateTypeEnum.hp_params:
+        return set(model.optimizer.get_hp_param_device(p) for p in model.parameters())
+    elif state == OffloadStateTypeEnum.lp_params:
+        return set(p.ds_tensor.device for p in model.parameters())
+    elif state == OffloadStateTypeEnum.lp_grads:
+        return {model.optimizer.grad_partitions_flat_buffer.device}
+    elif state == OffloadStateTypeEnum.optim_states:
+        return set(model.optimizer.get_hp_param_device(p, "exp_avg") for p in model.parameters()) | \
+               set(model.optimizer.get_hp_param_device(p, "exp_avg_sq") for p in model.parameters())
+    elif state == OffloadStateTypeEnum.contiguous_grad_buffer:
+        return set(bucket.buffer.device for bucket in model.optimizer.ipg_buckets.values()
+                   if bucket.buffer is not None)

deepspeed/runtime/zero/parameter_offload.py

@@ -0,0 +1,649 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import sys
+import torch
+from collections import OrderedDict
+from deepspeed.utils import z3_leaf_module, set_z3_leaf_module
+from deepspeed.runtime.utils import see_memory_usage
+from deepspeed.runtime.zero.utils import apply_to_tensors_only, is_zero_param
+from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
+from deepspeed.runtime.zero.partition_parameters import _init_external_params
+from deepspeed.runtime.zero.partition_parameters import *
+from deepspeed.runtime.zero.partitioned_param_coordinator import PartitionedParameterCoordinator, InflightParamRegistry, iter_params
+from deepspeed.accelerator import get_accelerator
+from deepspeed import utils
+
+FWD_MODULE_STACK = list()
+
+
+#for each tensor in outputs run the forward_function and register backward_function as hook
+def _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, outputs):
+    if type(outputs) is tuple:
+        touched_outputs = []
+        for output in outputs:
+            touched_output = _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function,
+                                                                         output)
+            touched_outputs.append(touched_output)
+        return tuple(touched_outputs)
+    elif type(outputs) is torch.Tensor:
+        forward_function(outputs)
+        if outputs.requires_grad:
+            outputs.register_hook(backward_function)
+        return outputs
+    else:
+        return outputs
+
+
+class ZeROOrderedDict(OrderedDict):
+
+    def __init__(self, parent_module, *args, **kwargs):
+        """A replacement for ``collections.OrderedDict`` to detect external ZeRO params.
+
+        Args:
+            parent_module (``collections.OrderedDict``): the collection to replace
+        """
+
+        super().__init__(*args, **kwargs)
+        self._parent_module = parent_module
+        self._in_forward = False
+
+    def __reduce__(self):
+        r0, _, *r2 = super().__reduce__()
+        return (r0, (self._parent_module, )) + tuple(r2)
+
+    def __getitem__(self, key):
+        param = super().__getitem__(key)
+
+        # Params can be registered as None (e.g., bias)
+        if param is None:
+            return param
+
+        if hasattr(param, "ds_status") and param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
+            if self._parent_module._parameters._in_forward and not torch.compiler.is_compiling():
+                from deepspeed.compile.z3_eager_fallback import get_active_z3_eager_fallback
+                fallback = get_active_z3_eager_fallback()
+                if fallback is None:
+                    register_external_parameter(FWD_MODULE_STACK[-1], param)
+                    param.all_gather()
+                else:
+                    param.all_gather()
+                    fallback.record_gathered_param(param)
+                print_rank_0(f'Registering external parameter from getter {key} ds_id = {param.ds_id}', force=False)
+
+        return param
+
+
+def _inject_parameters(module, cls):
+    for module in module.modules():
+        module._original_parameters = module._parameters
+
+        if cls == ZeROOrderedDict:
+            new_param = cls(parent_module=module)
+        else:
+            new_param = cls()
+
+        for key, param in module._parameters.items():
+            new_param[key] = param
+
+        module._parameters = new_param
+
+
+def ensure_zero_ordered_dict(module):
+    """Wrap ``module._parameters`` in :class:`ZeROOrderedDict` if not already.
+
+    PyTorch 2.5+ defaults ``nn.Module._parameters`` to a plain ``dict``
+    (pytorch/pytorch#129164), which rejects the ``_in_forward`` attribute
+    the forward prologue sets. Modules not converted by ``_inject_parameters``
+    at engine init (e.g. submodules attached after ``deepspeed.initialize``,
+    or restored by ``deepspeed/compile/init_z3.py``) hit issue #6961.
+    Idempotent; no-op if already wrapped, missing, or a non-dict container.
+    """
+    params = getattr(module, "_parameters", None)
+    if isinstance(params, ZeROOrderedDict) or not isinstance(params, dict):
+        return
+    # Preserve the original container only on first wrap so the un-injection
+    # path in ``deepspeed/compile/init_z3.py`` can restore it.
+    if not hasattr(module, "_original_parameters"):
+        module._original_parameters = params
+    new_param = ZeROOrderedDict(parent_module=module)
+    for key, param in params.items():
+        new_param[key] = param
+    module._parameters = new_param
+
+
+class DeepSpeedZeRoOffload(object):
+
+    def __init__(
+        self,
+        module,
+        timers,
+        ds_config,
+        zenflow=False,
+        overlap_comm=True,
+        prefetch_bucket_size=50000000,
+        max_reuse_distance=1000000000,
+        max_live_parameters=1000000000,
+        param_persistence_threshold=100000,
+        model_persistence_threshold=sys.maxsize,
+        dp_process_group=None,
+        offload_param_config=None,
+        mpu=None,
+        zero_param_parallel_group=None,
+        zero_quantized_weights=False,
+        zero_quantized_nontrainable_weights=False,
+        zero_module_granularity_threshold=0,
+        log_trace_cache_warnings=False,
+    ):
+
+        see_memory_usage("DeepSpeedZeRoOffload initialize [begin]", force=False)
+
+        print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False)
+
+        self.module = module
+        self.timers = timers
+        self.zenflow = zenflow
+        self.dtype = list(module.parameters())[0].dtype
+        self.dp_process_group = dp_process_group
+        self.offload_device = None
+        self.offload_param_pin_memory = False
+        self.zero_param_parallel_group = zero_param_parallel_group
+        self.zero_quantized_weights = zero_quantized_weights
+        self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights
+        self.log_trace_cache_warnings = log_trace_cache_warnings
+
+        if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none:
+            self.offload_device = offload_param_config.device
+            self.offload_param_pin_memory = offload_param_config.pin_memory
+
+        self._convert_to_zero_parameters(ds_config, module, mpu)
+
+        for m in module.modules():
+            _init_external_params(m)
+
+        _inject_parameters(module, ZeROOrderedDict)
+
+        self.param_numel_persistence_threshold = int(param_persistence_threshold)
+        self.model_persistence_threshold = int(model_persistence_threshold)
+        self.persistent_parameters = self.mark_persistent_parameters(self.param_numel_persistence_threshold,
+                                                                     self.model_persistence_threshold)
+
+        self._prefetch_bucket_sz = int(prefetch_bucket_size)
+        self._max_reuse_distance_in_numel = int(max_reuse_distance)
+        self._max_available_parameters_in_numel = int(max_live_parameters)
+        self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream(
+        ) if overlap_comm else get_accelerator().default_stream()
+
+        if not hasattr(module, "ds_inflight_param_registry"):
+            module.ds_inflight_param_registry = InflightParamRegistry()
+        self.__inflight_param_registry = module.ds_inflight_param_registry
+
+        self.fast_sharding_for_leaf_module = False
+
+        if zero_module_granularity_threshold > 0:
+            self.min_granularity_value = sys.maxsize
+            self.min_granularity_layer = None
+            self.granularity_info = set()
+            self.z3_leaf_layers = []
+            self._set_z3_leaf_modules_by_threshold(module, zero_module_granularity_threshold)
+            self.fast_sharding_for_leaf_module = True
+
+        self.param_coordinator = PartitionedParameterCoordinator(
+            prefetch_bucket_sz=self._prefetch_bucket_sz,
+            max_reuse_distance_in_numel=self._max_reuse_distance_in_numel,
+            max_available_parameters_in_numel=self._max_available_parameters_in_numel,
+            allgather_stream=self.__allgather_stream,
+            inflight_param_registry=self.__inflight_param_registry,
+            prefetch_nvme=self.offload_device == OffloadDeviceEnum.nvme,
+            timers=self.timers,
+            zero_quantized_weights=self.zero_quantized_weights,
+            zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights,
+            fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module,
+            log_trace_cache_warnings=self.log_trace_cache_warnings,
+        )
+
+        self.forward_hooks = []
+        self.backward_hooks = []
+
+        self.setup_zero_stage3_hooks()
+        print_rank_0(
+            f'Created module hooks: forward = {len(self.forward_hooks)}, backward = {len(self.backward_hooks)}',
+            force=False)
+
+        see_memory_usage("DeepSpeedZeRoOffload initialize [end]", force=False)
+
+    @instrument_w_nvtx
+    def partition_all_parameters(self):
+        """Partitioning Parameters that were not partitioned usually if parameters
+        of modules whose input parameters do not require grad computation do not
+        trigger post call and will therefore will remain unpartitioned"""
+        self.get_param_coordinator().release_and_reset_all(self.module)
+        for param in iter_params(self.module, recurse=True):
+            if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
+                raise RuntimeError(f"{param.ds_summary()} expected to be released")
+
+    def get_param_coordinator(self):
+        return self.param_coordinator
+
+    def empty_partition_cache(self):
+        self.partition_all_parameters()
+
+    def _convert_to_zero_parameters(self, ds_config, module, mpu):
+        non_zero_params = [p for p in module.parameters() if not is_zero_param(p)]
+        if non_zero_params:
+            zero_params = [p for p in module.parameters() if is_zero_param(p)]
+            if zero_params:
+                zero_params[0].convert_to_zero_parameters(param_list=non_zero_params)
+            else:
+                group = None
+                # parallel_state_sp doesn't have get_data_parallel_group
+                if mpu and hasattr(mpu, "get_data_parallel_group"):
+                    group = mpu.get_data_parallel_group()
+
+                Init(module=module,
+                     data_parallel_group=group,
+                     dtype=self.dtype,
+                     config_dict_or_path=ds_config,
+                     remote_device=self.offload_device,
+                     pin_memory=self.offload_param_pin_memory,
+                     mpu=mpu,
+                     zero_param_parallel_group=self.zero_param_parallel_group,
+                     zero_quantized_weights=self.zero_quantized_weights,
+                     zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights)
+
+    def destroy(self):
+        self._remove_module_hooks()
+
+    def _remove_module_hooks(self):
+        num_forward_hooks = len(self.forward_hooks)
+        num_backward_hooks = len(self.backward_hooks)
+
+        for hook in self.forward_hooks:
+            hook.remove()
+
+        for hook in self.backward_hooks:
+            hook.remove()
+
+        self.fwd_pre_hook.remove()
+
+        print_rank_0(f'Deleted module hooks: forward = {num_forward_hooks}, backward = {num_backward_hooks}',
+                     force=False)
+
+    def setup_zero_stage3_hooks(self):
+        self.hierarchy = 0
+
+        #reset step if in inference mode
+        @instrument_w_nvtx
+        def _start_of_forward_hook(module, *args):
+
+            self.get_param_coordinator().reset_step()
+
+        self.fwd_pre_hook = self.module.register_forward_pre_hook(_start_of_forward_hook)
+
+        #likely one of them should be enough but just to be safe
+        self._register_deepspeed_module(self.module)
+
+        # Add top module to stack trace
+        global FWD_MODULE_STACK
+        FWD_MODULE_STACK.append(self.module)
+
+    def mark_persistent_parameters(self, param_threshold, model_threshold):
+        persistent_params = []
+        total_persistent_parameters = 0
+        params_count = 0
+        for name, param in self.module.named_parameters(recurse=True):
+            if param.ds_numel + total_persistent_parameters > model_threshold:
+                continue
+
+            if param.ds_numel <= param_threshold:
+                params_count += 1
+                param.ds_persist = True
+                persistent_params.append(param)
+                total_persistent_parameters += param.ds_numel
+
+        print_rank_0(
+            f"Parameter Offload - Persistent parameters statistics: param_count = {params_count}, numel = {total_persistent_parameters}",
+            force=False)
+
+        return persistent_params
+
+    def _register_deepspeed_module(self, module, count=[0]):
+        # re-registering hooks on the root module leaves the coordinator trace stale;
+        # invalidate so it re-records on the next forward.
+        if module is self.module:
+            coordinator = self.get_param_coordinator()
+            if coordinator is not None and not coordinator.is_invalid_trace():
+                coordinator._invalidate_trace()
+        my_count = count[0]
+        module.ds_id = my_count
+
+        #print(f"{module.__class__} : {module.ds_id}")
+
+        if z3_leaf_module(module):
+            for param in module.parameters():
+                param.ds_z3_leaf_module = module
+        else:
+            for child in module.children():
+                count[0] = count[0] + 1
+                self._register_deepspeed_module(child, count=count)
+
+        @torch.compiler.disable
+        def _pre_forward_module_hook(module, *args):
+            self.pre_sub_module_forward_function(module)
+
+        @instrument_w_nvtx
+        def _post_forward_module_hook(module, input, output):
+
+            global FWD_MODULE_STACK
+            FWD_MODULE_STACK.pop()
+            if output is None:
+                output = []
+            elif not isinstance(output, (list, tuple)):
+                if torch.is_tensor(output):
+                    output = [output]
+                else:
+                    #print(f'got UNKNOWN type {type(output)}')
+                    outputs = []
+                    output = output if isinstance(output, dict) else vars(output)
+                    for name, val in output.items():
+                        if not name.startswith('__') and torch.is_tensor(val):
+                            outputs.append(val)
+                    output = outputs
+
+            for item in filter(lambda item: is_zero_param(item) or hasattr(item, 'ds_param_alias'), output):
+                key = id(item) if hasattr(item, 'ds_id') else id(item.ds_param_alias)
+                actual_external_param = item if hasattr(item, 'ds_id') else item.ds_param_alias
+
+                if not any(key in m._external_params for m in FWD_MODULE_STACK):
+                    actual_external_param.is_external_param = True
+                    module_to_register = FWD_MODULE_STACK[-1]
+                    register_external_parameter(module_to_register, actual_external_param)
+                    print_rank_0(
+                        f'Registering dangling parameter for module {module_to_register.__class__.__name__}, ds_id = {actual_external_param.ds_id}.',
+                        force=False)
+
+                    # It's possible that the parameter was already external to the completed module. If so, remove it the
+                    # registration as it will be covered by the outer module instead.
+                    if key in module._external_params:
+                        print_rank_0(
+                            f'  Unregistering nested dangling parameter from module {module.__class__.__name__}, ds_id = {actual_external_param.ds_id}',
+                            force=False)
+                        unregister_external_parameter(module, actual_external_param)
+
+                    actual_external_param.all_gather()
+
+            self.post_sub_module_forward_function(module)
+
+        def _bwd_hook_unexpected_inputs_msg(value):
+            return f"A module has unknown inputs or outputs type ({type(value)}) and the tensors embedded in it cannot be detected. " \
+                "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and " \
+                "output tensors and therefore may not get triggered properly."
+
+        def _pre_backward_module_hook(module, inputs, output):
+
+            return apply_to_tensors_only(module.pre_bwd_fn.apply,
+                                         output,
+                                         warning_msg_fn=_bwd_hook_unexpected_inputs_msg)
+
+        #This is an alternate to doing _post_backward_module_hook
+        #it uses tensor.register_hook instead of using torch.autograd.Function
+        def _alternate_post_backward_module_hook(module, inputs):
+            module.ds_grads_remaining = 0
+
+            #print(f"Before Forward {module.__class__.__name__}")
+
+            def _run_after_backward_hook(*unused):
+                module.ds_grads_remaining = module.ds_grads_remaining - 1
+                if module.ds_grads_remaining == 0:
+                    #print(f"After backward {module.__class__.__name__}")
+                    self.post_sub_module_backward_function(module)
+
+            def _run_before_forward_function(input):
+                if input.requires_grad:
+                    module.ds_grads_remaining += 1
+
+            return _apply_forward_and_backward_to_tensors_only(module, _run_before_forward_function,
+                                                               _run_after_backward_hook, inputs)
+
+        @torch.compiler.disable
+        def _post_backward_module_hook(module, inputs):
+            module.ds_grads_remaining = 0
+
+            return apply_to_tensors_only(module.post_bwd_fn.apply,
+                                         inputs,
+                                         warning_msg_fn=_bwd_hook_unexpected_inputs_msg)
+
+        # Pre forward hook
+        self.forward_hooks.append(module.register_forward_pre_hook(_pre_forward_module_hook))
+
+        # Post forward hook
+        self.forward_hooks.append(module.register_forward_hook(_post_forward_module_hook))
+
+        # Pre backward hook
+        if not hasattr(module, "pre_bwd_fn"):
+
+            @instrument_w_nvtx
+            def _run_before_backward_function(sub_module):
+                # some models (e.g. Albert) may run multiple forwards on the same layer in a loop
+                # before doing backwards, so each backward will need a pre-fetch - using reference
+                # counting to support this scenario
+                #print(f"COUNTER before: {sub_module.applied_pre_backward_ref_cnt}")
+                if sub_module.applied_pre_backward_ref_cnt > 0:
+                    self.pre_sub_module_backward_function(sub_module)
+                    sub_module.applied_pre_backward_ref_cnt -= 1
+                #print(f"COUNTER after: {sub_module.applied_pre_backward_ref_cnt}")
+
+            class PreBackwardFunctionForModule(torch.autograd.Function):
+
+                @staticmethod
+                def forward(outputs):
+                    return outputs.detach()
+
+                @staticmethod
+                def setup_context(ctx, inputs, output):
+                    ctx.module = module
+                    ctx.pre_backward_function = _run_before_backward_function
+                    if not hasattr(ctx.module, "applied_pre_backward_ref_cnt"):
+                        ctx.module.applied_pre_backward_ref_cnt = 0
+                    ctx.module.applied_pre_backward_ref_cnt += 1
+
+                @staticmethod
+                def backward(ctx, *args):
+                    ctx.pre_backward_function(ctx.module)
+                    return args
+
+            module.pre_bwd_fn = PreBackwardFunctionForModule
+
+        self.backward_hooks.append(module.register_forward_hook(_pre_backward_module_hook))
+
+        # post backward hook
+        if not hasattr(module, "post_bwd_fn"):
+
+            @instrument_w_nvtx
+            def _run_after_backward_function(sub_module):
+                if sub_module.ds_grads_remaining == 0:
+                    self.post_sub_module_backward_function(sub_module)
+
+            class PostBackwardFunctionModule(torch.autograd.Function):
+
+                @staticmethod
+                def forward(output):
+                    return output.detach()
+
+                @staticmethod
+                def setup_context(ctx, inputs, output):
+                    (output_in, ) = inputs
+                    ctx.module = module
+                    if output_in.requires_grad:
+                        #TODO SOME TIMES post backward does not seem to be triggered debug in detail
+                        #Should only cause increase in memory not correctness issue
+                        #if output.grad_fn.__class__.__name__ == 'ViewBackward':
+                        #    ctx.view=True
+                        #    print(f"Warning view tensor for input to module : {module.__class__.__name__}. Backward hooks may not trigger properly")
+                        #assert len(module.parameters(recurse=False)), "The input tensor to the module is a view, and autograd Function or register_hook is not triggered with view tensors."
+                        #if module.ds_grads_remaining == 0:
+                        #    print(f"Before Forward: {ctx.module.__class__.__name__}")
+                        module.ds_grads_remaining += 1
+                        ctx.post_backward_function = _run_after_backward_function
+
+                @staticmethod
+                def backward(ctx, *args):
+                    ctx.module.ds_grads_remaining = ctx.module.ds_grads_remaining - 1
+                    if ctx.module.ds_grads_remaining == 0:
+                        ctx.post_backward_function(ctx.module)
+                    return args
+
+            module.post_bwd_fn = PostBackwardFunctionModule
+
+        self.backward_hooks.append(module.register_forward_pre_hook(_post_backward_module_hook))
+
+    @torch.no_grad()
+    def pre_sub_module_forward_function(self, sub_module):
+        see_memory_usage(f"Before sub module function {sub_module.__class__.__name__}", force=False)
+
+        global FWD_MODULE_STACK
+        FWD_MODULE_STACK.append(sub_module)
+
+        param_coordinator = self.get_param_coordinator()
+        param_coordinator.trace_prologue(sub_module)
+        if param_coordinator.is_record_trace():
+            param_coordinator.record_module(sub_module)
+        param_coordinator.fetch_sub_module(sub_module, forward=True)
+
+        if self.zenflow:
+            params_to_fetch = set(iter_params(sub_module, recurse=z3_leaf_module(sub_module)))
+            for param in params_to_fetch:
+                param.data = param.data.t() if len(param.ds_shape) != 1 else param.data
+
+        see_memory_usage(f"Before sub module function {sub_module.__class__.__name__} after fetch", force=False)
+
+    @torch.no_grad()
+    def post_sub_module_forward_function(self, sub_module):
+        see_memory_usage(
+            f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} before release",
+            force=False)
+
+        if self.zenflow:
+            params_to_fetch = set(iter_params(sub_module, recurse=z3_leaf_module(sub_module)))
+            for param in params_to_fetch:
+                param.data = param.data.t() if len(param.ds_shape) != 1 else param.data
+
+        param_coordinator = self.get_param_coordinator()
+        param_coordinator.release_sub_module(sub_module, forward=True)
+
+        see_memory_usage(
+            f"After sub module function {sub_module.__class__.__name__}  {sub_module.ds_id} after release",
+            force=False)
+
+    @torch.no_grad()
+    def pre_sub_module_backward_function(self, sub_module):
+        # assert sub_module.training, "backward pass is invalid for module in evaluation mode"
+        param_coordinator = self.get_param_coordinator()
+        param_coordinator.trace_prologue(sub_module)
+        if param_coordinator.is_record_trace():
+            param_coordinator.record_module(sub_module)
+        param_coordinator.fetch_sub_module(sub_module, forward=False)
+
+        if self.zenflow:
+            params_to_fetch = set(iter_params(sub_module, recurse=z3_leaf_module(sub_module)))
+            for param in params_to_fetch:
+                param.data = param.data.t() if len(param.ds_shape) != 1 else param.data
+
+    @torch.no_grad()
+    def post_sub_module_backward_function(self, sub_module):
+        # assert sub_module.training, "backward pass is invalid for module in evaluation mode"
+        see_memory_usage(
+            f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} before release",
+            force=False)
+
+        if self.zenflow:
+            params_to_fetch = set(iter_params(sub_module, recurse=z3_leaf_module(sub_module)))
+            for param in params_to_fetch:
+                param.data = param.data.t() if len(param.ds_shape) != 1 else param.data
+
+        self.get_param_coordinator().release_sub_module(sub_module, forward=False)
+
+        see_memory_usage(
+            f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} after release",
+            force=False)
+
+    def _set_z3_leaf_modules_by_threshold(self, module, zero_module_granularity_threshold):
+
+        self._get_granularity_recursively(module)
+        print_rank_0(f"{'MODULE NAME'.ljust(30)}|{'GRANULARITY VALUE'.rjust(20)}", force=False)
+        for granularity in self.granularity_info:
+            print_rank_0(granularity, force=False)
+
+        if self.min_granularity_value <= zero_module_granularity_threshold:
+            self._set_leaf_by_threshold_preorder(module, zero_module_granularity_threshold)
+            utils.logger.info(
+                f"z3_leaf_module was set by stage3_module_granularity_threshold:{zero_module_granularity_threshold}")
+            for layer in self.z3_leaf_layers:
+                print_rank_0(f"{layer.__class__.__name__}:{layer.ds_model_granularity}", force=False)
+        else:
+            utils.logger.warning(
+                f"The smallest module granularity is [{self.min_granularity_layer}:{self.min_granularity_value}]. "\
+                f"To make stage3_module_granularity_threshold effective, you need to set stage3_module_granularity_threshold >= {self.min_granularity_value}. "\
+                f"Current Value:{zero_module_granularity_threshold}"
+            )
+
+    def _get_granularity_recursively(self, module):
+        """This function is used to recursively obtain the granularity of each module."""
+
+        # avoid setting as leaf for particularly large models, even if the granularity is very small
+        # an oversized leaf module increases the number of live parameters, introducing memory overhead
+        Z3_MAX_LEAF_SIZE = 1e9
+
+        if not list(module.parameters()):
+            # skip Modules without parameters, such as GELU, etc.
+            module.ds_model_granularity = sys.maxsize
+            return 0, 0
+
+        num_layers = 0
+        num_params = 0
+        num_params += sum(p.ds_numel for p in module.parameters(recurse=False))
+        if not any(module.children()):
+            # torch leaf module
+            module.ds_model_granularity = sys.maxsize
+            return 1, num_params
+
+        for child in module.children():
+            layers_in_child, params_in_child = self._get_granularity_recursively(child)
+            num_layers += layers_in_child
+            num_params += params_in_child
+
+        if module.__class__.__name__ in torch.nn.modules.container.__all__:
+            # Do not set container modules like ModuleList as leaf modules
+            # as this will prevent hooks from being set on their children
+            # and they may do not invoke the forward method
+            module.ds_model_granularity = sys.maxsize
+            return num_layers, num_params
+
+        num_layers += 1
+        ds_model_granularity = (num_params // num_layers) if num_params <= Z3_MAX_LEAF_SIZE else sys.maxsize
+        module.ds_model_granularity = ds_model_granularity
+        # module.ds_model_num_layers = num_layers
+        # module.ds_model_num_params = num_params
+        if self.min_granularity_value > ds_model_granularity:
+            self.min_granularity_value = ds_model_granularity
+            self.min_granularity_layer = module.__class__.__name__
+        self.granularity_info.add(f"{module.__class__.__name__.ljust(30)}|{str(ds_model_granularity).rjust(20)}")
+
+        return num_layers, num_params
+
+    def _set_leaf_by_threshold_preorder(self, module, granularity_treshhold):
+        '''Set modules as leaf modules based on the threshold, prioritizing parent nodes.'''
+
+        num_params = sum(p.ds_numel for p in module.parameters())
+        if num_params == 0:
+            # skip Modules without parameters, such as GELU, etc.
+            return
+        if module.ds_model_granularity <= granularity_treshhold:
+            set_z3_leaf_module(module, True)
+            self.z3_leaf_layers.append(module)
+            return
+
+        for sub_module in module.children():
+            self._set_leaf_by_threshold_preorder(sub_module, granularity_treshhold)

deepspeed/runtime/zero/partitioned_param_coordinator.py

@@ -0,0 +1,633 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from dataclasses import dataclass
+import collections
+from collections import UserDict
+import threading
+from typing import Deque, Set
+
+from deepspeed import comm as dist
+from deepspeed.utils import z3_leaf_module
+from deepspeed.utils.logging import logger
+from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
+from deepspeed.runtime.zero.partition_parameters import *
+from deepspeed.runtime.zero.partitioned_param_profiler import PartitionedParameterProfiler
+from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus
+from deepspeed.utils.debug import debug_param2name_id_shape
+from deepspeed.accelerator import get_accelerator
+import deepspeed.runtime.compiler as compiler
+from deepspeed.runtime.compiler import is_compiling
+
+import logging
+
+ENABLE_PROFILER = False
+
+
+def debug_rank0(message: str) -> None:
+    if dist.get_rank() == 0:
+        logger.debug(message)
+
+
+@instrument_w_nvtx
+def get_all_parameters(sub_module, recurse=False):
+    return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters())
+
+
+@compiler.enable(min_version="2.7.0")
+def iter_params(module: Module, recurse=False) -> Iterable[Parameter]:
+    return map(lambda pair: pair[1], get_all_parameters(module, recurse))
+
+
+class ZeRoTraceMode(Enum):
+    # Record trace of the network during a single forward+backward (for training) or forward (for inference)
+    RECORD = 1
+    # Use recorded network trace to optimize current forward+backward or forward
+    COMPLETE = 2
+    # Recorded trace does not match current forward+backward or forward pass.
+    INVALID = 3
+
+
+class InflightParamRegistry(UserDict):
+    """registry for parameters in flight"""
+
+    def __setitem__(self, param: Parameter, handle: AllGatherCoalescedHandle) -> None:
+        if param in self.data:
+            raise RuntimeError(f"{param.ds_summary()} already in registry")
+        if param.ds_status != ZeroParamStatus.INFLIGHT:
+            raise RuntimeError(f"attempted to add non-inflight parameter to registry {param.ds_summary()}")
+        self.data[param] = handle
+
+
+class PartitionedParameterCoordinator:
+    FORWARD_FETCH_SUBMIT = 'forward_fetch_submit'
+    FORWARD_FETCH_WAIT = 'forward_fetch_wait'
+    FORWARD_PREFETCH_SUBMIT = 'forward_prefetch_submit'
+    BACKWARD_FETCH_SUBMIT = 'backward_fetch_submit'
+    BACKWARD_FETCH_WAIT = 'backward_fetch_wait'
+    BACKWARD_PREFETCH_SUBMIT = 'backward_prefetch_submit'
+    FORWARD_ALL_GATHER = 'forward_all_gather'
+    BACKWARD_ALL_GATHER = 'backward_all_gather'
+    """Handles partitioning and gathering of parameters."""
+
+    @dataclass
+    class __ParamInTrace:
+        param: Parameter
+        step_id_last_used_at: int
+
+    def __init__(
+        self,
+        prefetch_bucket_sz: int,
+        max_reuse_distance_in_numel: int,
+        max_available_parameters_in_numel: int,
+        allgather_stream: get_accelerator().Stream,
+        inflight_param_registry: InflightParamRegistry,
+        prefetch_nvme: bool = False,
+        timers=None,
+        zero_quantized_weights=False,
+        zero_quantized_nontrainable_weights=False,
+        fast_sharding_for_leaf_module=False,
+        log_trace_cache_warnings=False,
+    ) -> None:
+        # mapping of param -> handle for each param that is currently in flight
+        self.__inflight_param_registry = inflight_param_registry
+        # keeps track of the number of submodules invoked so far.
+        self.__step_id: int = 0
+        # network tracing mode
+        self.__trace_mode: ZeRoTraceMode = ZeRoTraceMode.INVALID
+        # sequence of submodules/parameters in forward pass + backward pass
+        self.__submodule_order: Iterable[Module] = []
+        self.__param_order: Iterable[__class__.__ParamInTrace] = []
+        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10))
+        self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque())
+        # number of available params, and max number of available params
+        self.__n_available_params: int = 0
+        self.__max_n_available_params: int = max_available_parameters_in_numel
+        # max distance between two use of the module beyond which module is released
+        self.__max_reuse_dist_in_numel: int = max_reuse_distance_in_numel
+        # queue for parameters to fetch. parameters will be popped off the left
+        # side of the dequeue as they are fetched
+        self.__param_queue: Deque[__class__.__ParamInTrace] = None
+        self.__prefetch_bucket_sz: int = prefetch_bucket_sz
+        self.__prefetch_nvme: bool = prefetch_nvme
+        self.hierarchy: int = 0
+        self.zero_quantized_weights = zero_quantized_weights
+        self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights
+
+        # stream that will be used for allgather operations
+        self.__allgather_stream: get_accelerator().Stream = allgather_stream
+
+        # limit the number of fetch events that can be queued at once
+        # otherwise, what happens is memory is allocated by the host thread at the
+        # time of the call, but not used until later by the asynchronous cuda stream.
+        # allowing an infinite number of these to queue up causes a lot of memory
+        # pressure that then becomes detrimental to performance.
+        # this is a much less elegant way of fixing this vs something like using
+        # cudaMallocAsync/cudaFreeAsync. Choosing to not expose this to the user now
+        # because ideally in the future its replaced by an async allocation
+        # mechanism which doesn't require any configuration by the user.
+        self.__ongoing_fetch_events: Deque[get_accelerator().Event] = collections.deque()
+        # TODO. make this configurable via JSON
+        self.__max_ongoing_fetch_events: int = 2
+        self.__profiler = PartitionedParameterProfiler(timers if ENABLE_PROFILER else None)
+
+        # Whether to log trace cache warnings, e.g. invalidation events
+        self.__log_trace_cache_warnings = log_trace_cache_warnings
+
+        # whether to enable fast fetch for the z3 leaf module.
+        # this will improve fetch speed but will not break down leaf module parameters to alleviate memory pressure.
+        self.fast_sharding_for_leaf_module = fast_sharding_for_leaf_module
+
+        # Thread synchronization for leaf module fetches during backward pass.
+        # When autograd executes hooks in multiple threads (e.g., for modules returning multiple tensors),
+        # we need to ensure only one thread fetches parameters for a given leaf module at a time.
+        # This is only needed during backward pass; forward pass is single-threaded.
+        self.__ongoing_fetch_leaf_module_events = collections.defaultdict(threading.Event)
+        self.__leaf_module_lock = threading.Lock()
+
+    """Tracing and Tracking
+    TODO. consider performing trace before initializing PartitionedParameterCoordinator
+    and passing trace results into constructor. This way all the code in here can
+    just assume that the trace is complete and the results can be entirely
+    immutable.
+
+    Bookkeeping operations used to track where we are in the forward/backward pass
+    """
+
+    def _clear_trace_structures(self) -> None:
+        self.__submodule_order = []
+        self.__param_order = []
+        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10))
+        # clear the fetch-step deque too; a stale entry here causes record_parameters() to
+        # pop an empty deque (IndexError) after trace invalidation.
+        self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque())
+        self.__param_queue = None
+
+    def is_complete_trace(self) -> bool:
+        return self.__trace_mode == ZeRoTraceMode.COMPLETE
+
+    def is_invalid_trace(self) -> bool:
+        return self.__trace_mode == ZeRoTraceMode.INVALID
+
+    def is_record_trace(self) -> bool:
+        return self.__trace_mode == ZeRoTraceMode.RECORD
+
+    def _clean_inflight_param_registry(self) -> None:
+        for param, handle in self.__inflight_param_registry.items():
+            handle.wait()
+            self.__release_param(param)
+        self.__inflight_param_registry.clear()
+
+    def _invalidate_trace(self) -> None:
+        if self.is_invalid_trace():
+            raise RuntimeError("attempted to invalidate already invalid trace")
+        self.__trace_mode = ZeRoTraceMode.INVALID
+        self._clear_trace_structures()
+        self._clean_inflight_param_registry()
+
+    def trace_prologue(self, sub_module: Module) -> None:
+        if self.is_complete_trace():
+            # sub_module must match expectation else invalidate trace cache
+            if len(self.__submodule_order) <= self.__step_id:
+                print_rank_0(
+                    f"Invalidate trace cache @ step {self.__step_id} and module {sub_module.ds_id}: "
+                    f"cache has only {len(self.__submodule_order)} modules",
+                    force=self.__log_trace_cache_warnings)
+                self._invalidate_trace()
+                return
+
+            if sub_module != self.__submodule_order[self.__step_id]:
+                expected_module_id = self.__submodule_order[self.__step_id].ds_id
+                print_rank_0(
+                    f"Invalidate trace cache @ step {self.__step_id}: "
+                    f"expected module {expected_module_id}, but got module {sub_module.ds_id}",
+                    force=self.__log_trace_cache_warnings)
+                self._invalidate_trace()
+
+    @compiler.enable(min_version="2.7.0")
+    def record_module(self, sub_module: Module) -> None:
+        """adds sub module to trace"""
+        if is_compiling():
+            return
+
+        if not self.is_record_trace():
+            raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}")
+
+        self.__submodule_order.append(sub_module)
+        self.__step_id_module_fetched_for[sub_module.ds_id].append(self.__step_id)
+
+    def record_parameters(self, sub_module: Module) -> None:
+        if is_compiling():
+            return
+        """adds sub module to trace"""
+        if not self.is_record_trace():
+            raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}")
+
+        step_id = self.__step_id_module_fetched_for[sub_module.ds_id].popleft()
+        for param in sorted(set(iter_params(sub_module, recurse=z3_leaf_module(sub_module))), key=lambda p: p.ds_id):
+            self.__param_order.append(__class__.__ParamInTrace(param=param, step_id_last_used_at=step_id))
+
+    def construct_parameter_trace_from_module_trace(self):
+        """use module trace to construct parameter trace"""
+        self.__param_order = []
+        for sub_module in self.__submodule_order:
+            self.record_parameters(sub_module)
+
+    @compiler.disable
+    def reset_step(self) -> None:
+        """indicate that we have completed one fwd+bwd for the model"""
+        if is_compiling():
+            return
+
+        self._clean_inflight_param_registry()
+
+        if not self.is_complete_trace():  # not self.trace_complete:
+            # Make sure that recorded submodule orders are identical across ranks
+            assert_ints_same_as_other_ranks([m.ds_id for m in self.__submodule_order])
+
+            if self.is_record_trace():
+                # Successfully recorded a trace
+                self.construct_parameter_trace_from_module_trace()
+                # Make sure that recorded parameter orders are identical across ranks
+                assert_ints_same_as_other_ranks([p.param.ds_id for p in self.__param_order])
+                assert_ints_same_as_other_ranks([p.step_id_last_used_at for p in self.__param_order])
+
+                self.__submodule_order = tuple(self.__submodule_order)  # freeze
+                self.__param_order = tuple(self.__param_order)  # freeze
+                self.__trace_mode = ZeRoTraceMode.COMPLETE
+                print_rank_0(
+                    f"completed record trace of {len(self.__submodule_order)} sub modules: {[m.ds_id for m in self.__submodule_order]}",
+                    force=False)
+            else:
+                # Enable trace recording for next forward/backward pass
+                self.__trace_mode = ZeRoTraceMode.RECORD
+
+        else:
+            if self.__profiler is not None:
+                self.__profiler.log_events()
+
+        self.__param_queue = collections.deque(self.__param_order)  # reset fetch queue
+        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10))
+        self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque())
+        self.__step_id = 0
+        self.__n_available_params = 0
+        self.__profiler.reset_events()
+        # Clear leaf module fetch events for clean state
+        self.__ongoing_fetch_leaf_module_events.clear()
+
+    def _dump_params(self, tag, sub_module, params, step_id=None):
+        if step_id is None:
+            step_id = self.__step_id
+        param_names = [debug_param2name_id_shape(p) for p in params]
+        print_rank_0(f'{tag} step = {step_id} p_names = {param_names}', force=False)
+
+    def _dump_param_ids(self, tag, mod_id, p_ids, step_id=None):
+        if step_id is None:
+            step_id = self.__step_id
+        print_rank_0(f'{tag} mod = {mod_id}, step = {step_id}, p_ids = {p_ids}', force=False)
+
+    """Fetch and Release
+    Fetching, prefetching, and releasing parameters
+    """
+
+    @compiler.disable
+    @instrument_w_nvtx
+    @torch.no_grad()
+    def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None:
+        """This method does the following (in order):
+        1. kick off fetch for parameters in immediately required sub module
+        2. kick off fetch for next few parameters we will need later (prefetch)
+        3. block on parameters in immediately required sub module
+        """
+        # For leaf modules during backward pass, autograd may trigger hooks from multiple
+        # threads concurrently (e.g., when a module returns multiple tensors). We need to
+        # serialize access to prevent race conditions in parameter state management.
+        # Forward pass is single-threaded, so no synchronization is needed there.
+        is_leaf = z3_leaf_module(current_submodule)
+        needs_sync = is_leaf and not forward
+        if needs_sync:
+            event_to_wait = None
+            with self.__leaf_module_lock:
+                event = self.__ongoing_fetch_leaf_module_events.get(current_submodule.ds_id)
+                if event is not None:
+                    # Another thread is already fetching this leaf module, wait for it
+                    event_to_wait = event
+                else:
+                    # Mark that we're starting a fetch for this leaf module
+                    new_event = threading.Event()
+                    self.__ongoing_fetch_leaf_module_events[current_submodule.ds_id] = new_event
+
+            if event_to_wait is not None:
+                # Wait outside the lock to avoid deadlock
+                event_to_wait.wait()
+                return
+
+        try:
+            self._fetch_sub_module_impl(current_submodule, forward, is_leaf)
+        finally:
+            if needs_sync:
+                # Signal that we're done fetching this leaf module and remove the event
+                with self.__leaf_module_lock:
+                    event = self.__ongoing_fetch_leaf_module_events.pop(current_submodule.ds_id, None)
+                    if event is not None:
+                        event.set()
+
+    def _fetch_sub_module_impl(self, current_submodule: Module, forward: bool, is_leaf: bool) -> None:
+        """Implementation of fetch_sub_module, separated for thread synchronization."""
+        if logger.isEnabledFor(logging.DEBUG):
+            debug_rank0(
+                f"{self.__step_id}: M{current_submodule.ds_id}({type(current_submodule).__name__}) P{[p.ds_id for p in iter_params(current_submodule, recurse=is_leaf)]} "
+                + str({
+                    "avail": f"{self.__n_available_params:.1e}",
+                    "queue_sz": f"{len(self.__param_queue or [])}",
+                    "inflight": [p.ds_id for p in self.__inflight_param_registry],
+                }))
+
+        params_to_fetch = set(iter_params(current_submodule, recurse=is_leaf))
+        fetch_numel = sum(
+            [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE])
+
+        if fetch_numel > 0:
+            event_name = __class__.FORWARD_FETCH_SUBMIT if forward else __class__.BACKWARD_FETCH_SUBMIT
+            self._dump_param_ids(event_name, current_submodule.ds_id,
+                                 [(p.ds_id, p.ds_shape)
+                                  for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE])
+            # self._dump_params(event_name, current_submodule, [p for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE])
+
+            self.__profiler.start_event(event_name)
+            # kick off all gather for params in the immediately required submodule
+            #for param in params_to_fetch:
+            if logger.isEnabledFor(logging.DEBUG):
+                for param in params_to_fetch:
+                    debug_rank0(f"-fetch: {param.ds_summary()}")
+            self.__all_gather_params(params_to_fetch, forward)
+            self.__profiler.stop_event(event_name, fetch_numel)
+
+        wait_numel = 0
+        wait_event_name = __class__.FORWARD_FETCH_WAIT if forward else __class__.BACKWARD_FETCH_WAIT
+        self.__profiler.start_event(wait_event_name)
+        fast_fetch = self.fast_sharding_for_leaf_module and is_leaf
+        # wait for parameters in the immediately needed submodule to become available
+        for param in params_to_fetch:
+            param.ds_active_sub_modules.add(current_submodule.ds_id)
+            if logger.isEnabledFor(logging.DEBUG):
+                debug_rank0(f"-wait: {param.ds_summary()}")
+            if param in self.__inflight_param_registry:
+                wait_numel += param.partition_numel()
+                with get_accelerator().stream(self.__allgather_stream):
+                    while self.__ongoing_fetch_events and self.__ongoing_fetch_events[0].query():
+                        self.__ongoing_fetch_events.popleft()
+                    if len(self.__ongoing_fetch_events) > self.__max_ongoing_fetch_events:
+                        self.__ongoing_fetch_events.popleft().synchronize()
+
+                    self.__inflight_param_registry.pop(param).wait(handle_dependency=not fast_fetch)
+
+                    if not get_accelerator().handles_memory_backpressure() and not fast_fetch:
+                        event = get_accelerator().Event()
+                        event.record()
+                        self.__ongoing_fetch_events.append(event)
+
+            assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
+        if not get_accelerator().resolves_data_dependency():
+            get_accelerator().current_stream().wait_stream(self.__allgather_stream)
+        if fast_fetch:
+            AllGatherCoalescedHandle.free_buffer()
+        self.__profiler.stop_event(wait_event_name, wait_numel)
+
+        # kick off parameter prefetches for upcoming modules
+        # don't prefetch if we dont have a completed model trace
+        if self.is_complete_trace():
+            # go through the parameters we need for the current module and pop them
+            # off the fetch queue so that they aren't prefetched later.
+            # if params have already been popped off the fetch queue by earlier
+            # prefetches we won't look for them here
+            discarded_from_prefetch_queue = set()
+            params_not_already_fetched = set(
+                filter(lambda p: self.__most_recent_step_id_param_fetched_for[p] < self.__step_id, params_to_fetch))
+            while self.__param_queue and len(discarded_from_prefetch_queue) < len(params_not_already_fetched):
+                param_in_trace = self.__param_queue.popleft()
+                self.__most_recent_step_id_param_fetched_for[
+                    param_in_trace.param] = param_in_trace.step_id_last_used_at
+                discarded_from_prefetch_queue.add(param_in_trace.param)
+
+            if discarded_from_prefetch_queue != params_not_already_fetched:
+                raise RuntimeError(
+                    f"tracing error at step {self.__step_id}: \n"
+                    f"module id: {current_submodule.ds_id}, training: {current_submodule.training}\n"
+                    f"expected the next {len(params_not_already_fetched)} parameters in the "
+                    f"parameter fetch queue to be {tuple(p.ds_summary(use_debug_name=True) for p in params_not_already_fetched)} \n"
+                    f"but got \n {tuple(p.ds_summary(use_debug_name=True) for p in discarded_from_prefetch_queue)}.")
+
+            def _is_currently_on_nvme(param):
+                if param.nvme_swapper is None:
+                    return False
+
+                return param.ds_tensor.final_location == OffloadDeviceEnum.nvme \
+                    and param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE
+
+            # kick off all gather for params in the next few submodules (prefetch)
+            if self.__prefetch_bucket_sz > 0:
+                max_params_to_prefetch = min(self.__max_n_available_params - self.__n_available_params,
+                                             self.__prefetch_bucket_sz)
+                params_to_prefetch = set()
+                numel_prefetching = 0
+                while self.__param_queue and numel_prefetching < max_params_to_prefetch:
+                    param_in_trace: __class__.__ParamInTrace = self.__param_queue.popleft()
+
+                    if _is_currently_on_nvme(param_in_trace.param):
+                        # nvme prefetch is handled elsewhere. Need to break here to preserve fetch order
+                        self.__param_queue.appendleft(param_in_trace)
+                        break
+
+                    do_prefetch = param_in_trace.param.ds_status == ZeroParamStatus.NOT_AVAILABLE
+                    if param_in_trace.param in params_to_prefetch:
+                        # Avoid duplicates
+                        do_prefetch = False
+
+                    self.__most_recent_step_id_param_fetched_for[param_in_trace.param] = \
+                        max(self.__most_recent_step_id_param_fetched_for[param_in_trace.param],
+                            param_in_trace.step_id_last_used_at)
+
+                    if do_prefetch:
+                        params_to_prefetch.add(param_in_trace.param)
+                        numel_prefetching += param_in_trace.param.ds_numel
+
+                if numel_prefetching > 0:
+                    event_name = __class__.FORWARD_PREFETCH_SUBMIT if forward else __class__.BACKWARD_PREFETCH_SUBMIT
+                    self.__profiler.start_event(event_name)
+                    if logger.isEnabledFor(logging.DEBUG):
+                        for param in params_to_prefetch:
+                            debug_rank0(f"-prefetch: {param.ds_summary()}")
+                    self.__all_gather_params(params_to_prefetch, forward)
+                    self.__profiler.stop_event(event_name, numel_prefetching)
+
+                if self.__prefetch_nvme:
+                    self.__prefetch_nvme_param_partitions()
+
+        self.__step_id += 1
+
+    @instrument_w_nvtx
+    @torch.no_grad()
+    def release_sub_module(self, submodule: Module, forward=False) -> None:
+        """release the parameters of a sub module, assuming they meet conditions to
+        be released."""
+        #print_rank_0(f"release_sub_module {'fwd' if forward else 'bwd'}: {debug_module2name_id(submodule)}", force=False)
+        params_to_release = (self.__params_to_release(submodule, self.__step_id) if self.is_complete_trace() else set(
+            p.ds_id for p in iter_params(submodule, recurse=z3_leaf_module(submodule))))
+
+        free_data = not z3_leaf_module(submodule) or not self.fast_sharding_for_leaf_module
+        if not free_data:
+            # wait for the computation to finish and launch as early as possible.
+            empty_buffer = torch.empty(1, device=torch.device(get_accelerator().current_device_name()))
+
+        for param in iter_params(submodule, recurse=z3_leaf_module(submodule)):
+            param.ds_active_sub_modules.discard(submodule.ds_id)
+            if param.ds_id in params_to_release and not param.is_external_param:
+                self.__release_param(param, free_data)
+            if not free_data:
+                if param.ds_id in params_to_release and not param.is_external_param:
+                    # empty buffer ensures that all computations are complete
+                    param.data = empty_buffer
+
+    @instrument_w_nvtx
+    @torch.no_grad()
+    def release_and_reset_all(self, module: Module) -> None:
+        """release all module parameters"""
+        for param in iter_params(module, recurse=True):
+            if param in self.__inflight_param_registry:
+                self.__inflight_param_registry.pop(param).wait()
+
+            # TODO. make this throw if if there are still active submodules. currently
+            # there's a hook execution issue
+            param.ds_active_sub_modules.clear()
+            self.__release_param(param)
+
+        for param in iter_params(module, recurse=True):
+            if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
+                raise RuntimeError(f"{param.ds_summary()} expected to be released")
+
+    @instrument_w_nvtx
+    def __all_gather_params(self, params: Set[Parameter], forward: bool) -> None:
+        quantized_params = []
+        nonquantized_params = []
+        for param in params:
+            if hasattr(param.ds_tensor, 'ds_quant_scale'):
+                quantized_params.append(param)
+            else:
+                nonquantized_params.append(param)
+        if quantized_params:
+            self.__all_gather_params_(quantized_params, forward, quantize=True)
+        if nonquantized_params:
+            self.__all_gather_params_(nonquantized_params, forward, quantize=self.zero_quantized_weights)
+
+    def __all_gather_params_(self, params: Set[Parameter], forward: bool, quantize: bool = False) -> None:
+        """for each partitioned parameter, kick off an async allgather and store
+        the work handle for the in flight parameters."""
+        partitioned_params = []
+        all_gather_numel = 0  # numel = num of elements
+        for param in params:
+            if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
+                partitioned_params.append(param)
+                all_gather_numel += param.ds_numel
+
+        if partitioned_params:
+            self.__n_available_params += all_gather_numel
+            # here we need to handle a special case where some of the parameters have a valid hpz secondary tensor (e.g. they are not trainable so their secondary tensor never expire) but others do not.
+            partitioned_params_with_secondary_tensors = [
+                p for p in partitioned_params if p.ds_secondary_tensor is not None
+            ]
+            partitioned_params_without_secondary_tensors = [
+                p for p in partitioned_params if p.ds_secondary_tensor is None
+            ]
+            for param_group in [
+                    partitioned_params_with_secondary_tensors, partitioned_params_without_secondary_tensors
+            ]:
+                if not param_group:
+                    continue
+                with get_accelerator().stream(self.__allgather_stream):
+                    event_name = __class__.FORWARD_ALL_GATHER if forward else __class__.BACKWARD_ALL_GATHER
+                    self.__profiler.start_event(event_name)
+                    handle = param_group[0].all_gather_coalesced(param_group, quantize=quantize)
+                    self.__profiler.stop_event(event_name, all_gather_numel)
+                for param in param_group:
+                    assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary()
+                    self.__inflight_param_registry[param] = handle
+
+            # Release swap buffers for persisted params on nvme since they will never be partitioned or evicted from GPU
+            swap_persisted_params = [
+                p for p in partitioned_params if p.ds_persist and p.ds_tensor.final_location == OffloadDeviceEnum.nvme
+            ]
+            if swap_persisted_params:
+                swap_persisted_params[0].nvme_swapper.remove_partition_and_release_buffers(swap_persisted_params)
+
+    @compiler.disable
+    @instrument_w_nvtx
+    def __release_param(self, param: Parameter, free_data: bool = True) -> None:
+        if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules:
+            if logger.isEnabledFor(logging.DEBUG):
+                debug_rank0(f"-release: {param.ds_summary()}")
+                print_rank_0(f"release: {debug_param2name_id_shape(param)}", force=False)
+            param.partition(free_data=free_data)
+            self.__n_available_params -= param.ds_numel
+
+    @instrument_w_nvtx
+    @functools.lru_cache(maxsize=None)
+    def __params_to_release(self, submodule_to_release: Module, step_id: int) -> Set[int]:
+        if not self.is_complete_trace():
+            raise RuntimeError("expected trace to be complete")
+
+        params_to_release = set(
+            p.ds_id for p in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release))
+            if not p.ds_persist)
+
+        # Problem: When prefetcher scans the param trace, it skips AVAILABLE params.
+        # This creates issues if those params are released before the skipped uses:
+        # 1) It hurts performance as the skipped uses are never prefetched.
+        # 2) For nvme params, we run out of swap buffers because the prefetch order
+        # diverges from the trace.
+        # Solution: Don't release params whose reuse was skipped by prefetch. This is
+        # possible because we detect such skips during prefetch and mark those params.
+        for param in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)):
+            if self.__most_recent_step_id_param_fetched_for[param] > step_id:
+                params_to_release.discard(param.ds_id)
+
+        # examine all modules within `max_reuse_dist_in_numel` of the current step,
+        # if we see any of the candidate parameters to be released reoccur while
+        # doing this, remove them from the set of parameters to release.
+        params_traversed = 0
+        for module in self.__submodule_order[step_id:]:
+            if params_traversed >= self.__max_reuse_dist_in_numel:
+                break
+            for param in iter_params(module, recurse=z3_leaf_module(submodule_to_release)):
+                params_to_release.discard(param.ds_id)
+                params_traversed += param.ds_numel
+
+        return params_to_release
+
+    @instrument_w_nvtx
+    def __prefetch_nvme_param_partitions(self) -> None:
+        """swap in parameter partitions from nvme for those parameters that will be used
+        after the ones that are already being prefetched into full parameters
+        """
+        if not self.is_complete_trace():
+            return
+
+        numel_in_flight = sum(param.ds_numel for param in self.__inflight_param_registry)
+
+        numel_considered = 0
+        swap_in_params = []
+        for param_in_trace in self.__param_queue:
+            param = param_in_trace.param
+            if param.nvme_swapper is None:
+                continue
+            if (numel_considered > 2 * numel_in_flight
+                    or len(swap_in_params) >= param.nvme_swapper.available_swap_in_buffers()):
+                break
+            if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE:
+                swap_in_params.append(param)
+            numel_considered += param.ds_numel
+
+        if swap_in_params:
+            swap_in_params[0].nvme_swapper.swap_in(swap_in_params, async_op=True)

deepspeed/runtime/zero/partitioned_param_profiler.py

@@ -0,0 +1,63 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from dataclasses import dataclass
+from deepspeed.utils import log_dist
+
+
+class PartitionedParameterProfiler(object):
+
+    @dataclass
+    class EventCounter:
+        name: str
+        count: int
+        num_elem: int
+
+        def reset(self):
+            self.count = 0
+            self.num_elem = 0
+
+        def increment(self, numel):
+            self.count += 1
+            self.num_elem += numel
+
+    def __init__(self, timers):
+        self.timers = timers
+        self.event_counters = {}
+
+    def reset_events(self):
+        for event_ctr in self.event_counters.values():
+            event_ctr.reset()
+
+    def start_event(self, name):
+        if self.timers is None:
+            return
+
+        if name not in self.event_counters:
+            self.event_counters[name] = __class__.EventCounter(name=name, count=0, num_elem=0)
+        self.timers(name).start()
+
+    def stop_event(self, name, num_elem):
+        if self.timers is None:
+            return
+        assert name in self.event_counters, f'unknown event {name}'
+        self.event_counters[name].increment(num_elem)
+        self.timers(name).stop()
+
+    def _log_timers(self):
+        if self.timers is None:
+            return
+        self.timers.log(names=list(self.event_counters.keys()))
+
+    def _log_event_counters(self):
+        for event_ctr in self.event_counters.values():
+            log_dist(
+                f'{event_ctr.name}: count = {event_ctr.count}, numel = {event_ctr.num_elem}',
+                #f'{event_ctr.name}: time = {self._log_timers()},count = {event_ctr.count}, numel = {event_ctr.num_elem}',
+                ranks=[0])
+
+    def log_events(self):
+        self._log_event_counters()
+        self._log_timers()

deepspeed/runtime/zero/test.py

@@ -0,0 +1,77 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from deepspeed.runtime.zero.contiguous_memory_allocator import ContiguousMemoryAllocator
+
+
+def test1():
+    mem = ContiguousMemoryAllocator(1024, torch.half, 'cpu')
+    mem.print_allocation(resolution=100)
+    a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0)
+    mem.print_allocation(resolution=100)
+    mem.release_tensor(a1)
+    mem.print_allocation(resolution=100)
+    a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0)
+    a3 = mem.allocate_tensor(256).mul_(0.0).add_(3.0)
+    a4 = mem.allocate_tensor(128).mul_(0.0).add_(4.0)
+    mem.print_allocation(resolution=100)
+    mem.release_tensor(a3)
+    mem.print_allocation(resolution=100)
+    a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0)
+    a6 = mem.allocate_tensor(256).mul_(0.0).add_(6.0)
+    a7 = mem.allocate_tensor(128).mul_(0.0).add_(7.0)
+    mem.print_allocation(resolution=100)
+    a8 = mem.allocate_tensor(256).mul_(0.0).add_(8.0)
+    a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0)
+    mem.print_allocation(resolution=100)
+    mem.release_tensor(a9)
+    mem.release_tensor(a6)
+    mem.release_tensor(a2)
+    mem.release_tensor(a5)
+
+    a10 = mem.allocate_tensor(512).mul_(0.0).add_(10.0)
+    mem.print_allocation(resolution=100)
+    #print(f"a4:{a4}")
+    #print(f"a7:{a7}")
+    #print(f"a8:{a8}")
+    #print(f"a10:{a10}")
+    assert (a4.norm() + a7.norm() + a8.norm() + a10.norm()).item() == 474.50, "Test failed"
+
+
+def test2():
+    mem = ContiguousMemoryAllocator(512, torch.half, 'cpu')
+    a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0)
+    a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0)
+    a3 = mem.allocate_tensor(64).mul_(0.0).add_(3.0)
+    a4 = mem.allocate_tensor(64).mul_(0.0).add_(4.0)
+    a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0)
+    a6 = mem.allocate_tensor(64).mul_(0.0).add_(6.0)
+    a7 = mem.allocate_tensor(64).mul_(0.0).add_(7.0)
+    a8 = mem.allocate_tensor(64).mul_(0.0).add_(8.0)
+    mem.release_tensor(a2)
+    mem.release_tensor(a4)
+    mem.release_tensor(a6)
+    mem.release_tensor(a8)
+    mem.print_allocation(resolution=100)
+
+    a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0)
+    a10 = mem.allocate_tensor(64).mul_(0.0).add_(10.0)
+    a11 = mem.allocate_tensor(64).mul_(0.0).add_(11.0)
+    mem.release_tensor(a1)
+    mem.release_tensor(a5)
+    mem.print_allocation(resolution=100)
+    a12 = mem.allocate_tensor(128).mul_(0.0).add_(12.0)
+    mem.print_allocation(resolution=100)
+    print(f"a7:{a7}")
+    print(f"a9:{a9}")
+    print(f"a10:{a10}")
+    print(f"a11:{a11}")
+    print(f"a12:{a12}")
+    assert (a7.norm() + a9.norm() + a10.norm() + a11.norm() + a12.norm()) == 460.75, "TestFailed"
+
+
+test1()
+test2()

deepspeed/runtime/zero/tiling.py

@@ -0,0 +1,296 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from deepspeed.runtime.utils import partition_uniform as partition
+
+
+def split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False):
+    """Split a tensor along its last dimension. Adapted from Megatron-LM.
+
+    Arguments:
+        tensor: input tensor.
+        partitions: list of partition sizes to supply to torch.split
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                 in memory.
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+
+    # Split.
+    tensor_list = torch.split(tensor, partitions, dim=last_dim)
+    # Note: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tensor_list
+
+
+class TiledLinear(torch.nn.Module):
+
+    def __init__(self,
+                 in_features,
+                 out_features,
+                 bias=True,
+                 in_splits=1,
+                 out_splits=1,
+                 input_is_already_split=False,
+                 combine_out_splits=True,
+                 linear_cls=torch.nn.Linear,
+                 init_linear=None,
+                 **kwargs):
+        """A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce
+        memory requirements via tiling.
+
+        TiledLinear breaks the input and output dimensions of a linear layer
+        into tiles that are processed in sequence. This class enables huge
+        linear layers when combined with ZeRO-3 because inactive tiles can be
+        partitioned and offloaded.
+
+        .. note::
+            We recommend using as few tiles as necessary. Tiling
+            significantly reduces memory usage, but can reduce throughput
+            for inexpensive layers. This due to the smaller kernels having
+            less parallelism and lower arithmetic intensity, while
+            introducing more frequent synchronization and communication.
+
+        Args:
+            in_features (int): See ``torch.nn.Linear``
+            out_features (int): See ``torch.nn.Linear``
+            bias (bool, optional): See ``torch.nn.Linear``
+            in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1.
+            out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1.
+            input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in
+                to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``.
+            combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs
+                into a single tensor. Defaults to ``True``.
+            linear_cls (class, optional): The underlying class to build individual tiles.
+                Defaults to ``torch.nn.Linear``.
+            init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of
+                ``init_linear``. Useful for debugging. Defaults to ``None``.
+            kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``.
+
+        Raises:
+            RuntimeError: ``in_splits`` must be within the range [1, in_features).
+            RuntimeError: ``out_splits`` must be within the range of [1, out_features).
+        """
+
+        super().__init__()
+
+        if (in_splits < 1) or (in_splits > in_features):
+            raise RuntimeError('in splits must be in range [1, in_features].')
+        if (out_splits < 1) or (out_splits > out_features):
+            raise RuntimeError('out splits must be in range [1, out_features].')
+
+        # global, not necessarily local
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.out_splits = out_splits
+        self.in_splits = in_splits
+        self.input_is_already_split = input_is_already_split
+        self.combine_out_splits = combine_out_splits
+
+        # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features]
+        # For example, row_parts[p] gives the start of partition p and row_parts[p+1]
+        # is the exclusive end.
+        self.in_parts = partition(num_items=in_features, num_parts=in_splits)
+        self.out_parts = partition(num_items=out_features, num_parts=out_splits)
+
+        assert len(self.out_parts) == out_splits + 1
+        assert len(self.in_parts) == in_splits + 1
+        assert self.out_parts[0] == 0
+        assert self.out_parts[out_splits] == out_features
+        assert self.in_parts[in_splits] == in_features
+
+        self.linears = torch.nn.ModuleList()
+        for out_id in range(out_splits):
+            self.linears.append(torch.nn.ModuleList())
+
+            local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id]
+
+            for in_id in range(in_splits):
+                #if input_size is split, we only need one bias
+                local_bias = bias if in_id == (in_splits - 1) else False
+
+                local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id]
+                local = linear_cls(local_in_dim, local_out_dim, bias=local_bias, **kwargs)
+                self.linears[out_id].append(local)
+
+        # Optionally initialize with a known tensor
+        if init_linear is not None:
+            self.copy_params_from(init_linear)
+
+    def forward(self, input_):
+        if self.in_splits > 1 and not self.input_is_already_split:
+            input_parts = partition(input_.shape[-1], self.in_splits)
+            split_sizes = [input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)]
+            inputs = self._split_global_input(input_, split_sizes)
+        elif self.in_splits > 1:
+            inputs = input_
+            assert len(
+                inputs) == self.in_splits, f"Col splits {self.in_splits} does not match input splits {len(inputs)}"
+        else:
+            # no splits
+            inputs = [input_]
+
+        outputs = [None] * self.out_splits
+        for out_id in range(self.out_splits):
+            for in_id in range(self.in_splits):
+                local_output = self.linears[out_id][in_id](inputs[in_id])
+
+                outputs[out_id] = self._reduce_local_output(in_id=in_id,
+                                                            out_id=out_id,
+                                                            current_out=outputs[out_id],
+                                                            new_out=local_output)
+
+        if self.combine_out_splits:
+            return self._combine_output_splits(outputs)
+
+        return outputs
+
+    def _split_global_input(self, input, split_sizes):
+        """Partition an input tensor along the last dimension, aligned with given splits.
+
+        Subclasses should override this method to account for new input types.
+
+        Args:
+            input (List[Tensor]): The tensor to partition along the last dimension.
+            split_sizes (List[int]): The size of each partition.
+
+        Returns:
+            List[Any]: A list of the chunks of ``input``.
+        """
+        return split_tensor_along_last_dim(input, split_sizes)
+
+    def _reduce_local_output(self, in_id, out_id, current_out, new_out):
+        """Reduce (sum) a new local result into the existing local results.
+
+        Subclasses should override this method.
+
+        For a given ``out_id``, this method is called ``in_id-1`` times. The first input
+        split is a simple assignment.
+
+        Args:
+            in_id (int): The input split that produced ``new_out``.
+            out_id (int): The output split that produced ``new_out``.
+            current_out (Any): The reduced form of all previous ``out_id`` results.
+            new_out (Any): The local result from forward (``in_id``, ``out_id``)e
+
+        Returns:
+            Any: The combined result of ``current_out`` and ``new_out``.
+        """
+
+        if current_out is None:
+            #this clone is necessary to preserve auto grad
+            #there is some issue with inplace update for outputs that are views
+            return new_out.clone()
+        else:
+            return current_out + new_out
+
+    def _combine_output_splits(self, outputs):
+        """Join the splits of the output into a single result.
+
+        Args:
+            outputs (List[Any]): The reduced outputs for each output split.
+
+        Returns:
+            Any: The combined outputs.
+        """
+        assert len(outputs) == self.out_splits
+        return torch.cat(outputs, dim=-1)
+
+    @torch.no_grad()
+    def copy_params_from(self, other):
+        """Copy the weight and bias data from ``other``.
+
+        This is especially useful for reproducible initialization and testing.
+
+        Equivalent to:
+
+        .. code-block:: python
+
+            with torch.no_grad():
+                self.weight.copy_(other.weight)
+                if self.bias is not None:
+                    self.bias.copy_(other.bias)
+
+        .. note::
+            If ZeRO-3 is enabled, this is a collective operation and the
+            updated parameters of data-parallel rank 0 will be visible on all
+            ranks. See :class:`deepspeed.zero.GatheredParameters` for more
+            information.
+
+
+        Args:
+            other (``torch.nn.Linear``): the linear layer to copy from.
+        """
+        assert hasattr(other, 'weight')
+        assert other.weight.size() == (self.out_features, self.in_features)
+        if self.use_bias:
+            assert hasattr(other, 'bias')
+            assert other.bias is not None
+            assert other.bias.size() == (self.out_features, )
+        else:
+            assert other.bias is None
+
+        for row in range(self.out_splits):
+            rstart = self.out_parts[row]
+            rstop = self.out_parts[row + 1]
+
+            for col in range(self.in_splits):
+                cstart = self.in_parts[col]
+                cstop = self.in_parts[col + 1]
+
+                local = self.linears[row][col]
+                global_weight = other.weight[rstart:rstop, cstart:cstop]
+                with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0):
+                    local.weight.copy_(global_weight)
+
+            if local.bias is not None:
+                with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0):
+                    local.bias.data.copy_(other.bias[rstart:rstop].data)
+
+
+class TiledLinearReturnBias(TiledLinear):
+    """Wrapper for a Linear class that returns its own bias parameter, such as
+    used by Megatron-LM.
+    """
+
+    def _reduce_local_output(self, in_id, out_id, current_out, new_out):
+        """Reduces output tensors, but not the returned bias. """
+        if current_out is not None:
+            old_tensor, old_bias = current_out
+        else:
+            old_tensor, old_bias = None, None
+
+        assert isinstance(new_out, tuple)
+        assert len(new_out) == 2
+
+        tensor, bias = new_out
+        assert tensor is not None
+
+        tensor = super()._reduce_local_output(in_id=in_id, out_id=out_id, current_out=old_tensor, new_out=tensor)
+
+        if bias is None:
+            bias = old_bias
+
+        return tensor, bias
+
+    def _combine_output_splits(self, outputs):
+        # stack output tensors
+        tensors = [o[0] for o in outputs]
+        tensor = super()._combine_output_splits(tensors)
+
+        # stack biases if applicable
+        biases = [o[1] for o in outputs if o[1] is not None]
+        if len(biases) > 0:
+            bias = super()._combine_output_splits(biases)
+        else:
+            bias = None
+
+        return tensor, bias

deepspeed/runtime/zero/utils.py

@@ -0,0 +1,235 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import gc
+from typing import List, Tuple
+
+import torch
+from deepspeed import comm as dist
+from deepspeed.utils import logger
+from deepspeed.ops.adam import DeepSpeedCPUAdam, ZenFlowCPUAdam
+from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad
+from deepspeed.ops.adam import FusedAdam
+from deepspeed.ops.lion import DeepSpeedCPULion, FusedLion
+from deepspeed.utils.nvtx import instrument_w_nvtx
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.utils import get_only_unique_item
+
+# ensure we only warn once, otherwise every iteration will trigger a warning
+warned = False
+
+
+def _initialize_parameter_parallel_groups(parameter_parallel_size=None):
+    data_parallel_size = int(dist.get_world_size())
+    parameter_parallel_size = parameter_parallel_size or data_parallel_size
+    logger.info("data_parallel_size: %s, parameter_parallel_size: %s", data_parallel_size, parameter_parallel_size)
+    assert data_parallel_size % parameter_parallel_size == 0, \
+        'world size should be divisible by parameter parallel size'
+    rank = dist.get_rank()
+    my_group = None
+    for i in range(data_parallel_size // parameter_parallel_size):
+        ranks = range(i * parameter_parallel_size, (i + 1) * parameter_parallel_size)
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            my_group = group
+    return my_group
+
+
+class ZeRORuntimeException(Exception):
+    pass
+
+
+ZERO_SUPPORTED_OPTIMIZERS = [
+    torch.optim.Adam, torch.optim.AdamW, FusedAdam, DeepSpeedCPUAdam, ZenFlowCPUAdam, torch.optim.Adagrad,
+    DeepSpeedCPUAdagrad, DeepSpeedCPULion, FusedLion
+]
+
+# Add MuonWithAuxAdam to supported list if muon is installed
+try:
+    from deepspeed.runtime.zero.muon.muon_optimizer import MuonWithAuxAdam
+    ZERO_SUPPORTED_OPTIMIZERS.append(MuonWithAuxAdam)
+except ImportError:
+    pass
+
+# Add apex FusedAdam to supported list if apex is installed
+try:
+    import apex
+    if hasattr(apex, 'optimizers') and hasattr(apex.optimizers, 'FusedAdam'):
+        ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam)
+except ImportError:
+    pass
+
+
+def is_zero_supported_optimizer(optimizer):
+    if dist.get_rank() == 0:
+        logger.info(f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}')
+    return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS
+
+
+@instrument_w_nvtx
+def assert_lst_len_same_as_other_ranks(lst: List[int]) -> None:
+    rank0_len_tensor = torch.tensor(
+        len(lst) if dist.get_rank() == 0 else -1,
+        dtype=int,
+        device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])),
+        requires_grad=False,
+    )
+    local_list_length = len(lst)
+    dist.broadcast(rank0_len_tensor, src=0, async_op=False)
+    rank0_list_length = rank0_len_tensor.cpu().item()
+    if rank0_list_length != local_list_length:
+        raise RuntimeError(f"Detected a disagreement on list length between rank0 and rank{dist.get_rank()}: "
+                           f"\n rank0: {rank0_list_length} "
+                           f"\n rank{dist.get_rank()}: {local_list_length}")
+
+
+def get_lst_from_rank0(lst: List[int]) -> None:
+    """
+    NOTE: creates both communication and synchronization overhead so should be used
+    sparingly
+    """
+    lst_tensor = torch.tensor(
+        lst if dist.get_rank() == 0 else [-1] * len(lst),
+        dtype=int,
+        device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])),
+        requires_grad=False,
+    )
+    dist.broadcast(lst_tensor, src=0, async_op=False)
+
+    return [t.item() for t in lst_tensor.cpu()]
+
+
+@instrument_w_nvtx
+def assert_ints_same_as_other_ranks(ints: List[int]) -> None:
+    """
+    NOTE: creates both communication and synchronization overhead so should be
+    used sparingly
+
+    takes a list of ints from each rank and ensures that they are the same
+    across ranks, throwing an exception if they are not.
+    """
+    assert_lst_len_same_as_other_ranks(ints)
+    rank0_ints = get_lst_from_rank0(ints)
+    if ints != rank0_ints:
+        raise RuntimeError(f"Detected a disagreement on list contents between rank0 and rank{dist.get_rank()}: "
+                           f"\n list length: {len(ints)}"
+                           f"\n rank0: {rank0_ints} "
+                           f"\n rank{dist.get_rank()}: {ints}")
+
+
+def is_builtin_type(obj):
+    # https://stackoverflow.com/a/17795199
+    return obj.__class__.__module__ == '__builtin__' or obj.__class__.__module__ == "builtins"
+
+
+def isinstance_namedtuple(obj: object) -> bool:
+    """
+    Is this an instance of namedtuple/NamedTuple?
+    From: https://stackoverflow.com/a/62692640
+
+    Args:
+        obj (object): An object.
+
+    Returns:
+        bool: True if namedtuple/NamedTuple else False.
+    """
+    return isinstance(obj, tuple) and hasattr(obj, '_asdict') and hasattr(obj, '_fields')
+
+
+def is_zero_param(parameter):
+    if not torch.is_tensor(parameter):
+        return False
+    return hasattr(parameter, 'ds_id')
+
+
+def apply_to_tensors_only(function, value, warning_msg_fn=None):
+    """
+    Apply `function` to every Tensor in `value`.
+
+    Args:
+        functional: The function class to apply.
+        value (Any): Target object to apply `function` to.
+
+    Returns:
+        Any: Output of `function`.
+    """
+    if isinstance(value, (tuple, list)):
+        touched_outputs = []
+        for elem in value:
+            touched_output = apply_to_tensors_only(function, elem)
+            touched_outputs.append(touched_output)
+
+        if isinstance_namedtuple(value):
+            # namedtuples require a slightly different syntax.
+            return value.__class__(*touched_outputs)
+
+        return value.__class__(touched_outputs)
+    elif isinstance(value, dict):
+        # apply inplace to avoid recreating dict inherited objects
+        for key in value.keys():
+            value[key] = apply_to_tensors_only(function, value[key])
+        return value
+
+    elif isinstance(value, torch.Tensor):
+        # this also applies to torch.Tensor's subclasses like torch.nn.parameter.Parameter
+        touched_output = function(value)
+
+        # restore zero param attributes if those get stripped by `backward_function`
+        if not is_zero_param(touched_output) and is_zero_param(value):
+            touched_output.ds_param_alias = value
+
+        return touched_output
+    else:
+        if not is_builtin_type(value):
+            global warned
+            if warning_msg_fn and not warned and dist.get_rank() == 0:
+                logger.warning(warning_msg_fn(value))
+                warned = True
+        return value
+
+
+def get_mapping_to_flat_buffer(tensors: List[torch.Tensor]) -> List[Tuple[torch.Tensor, int, int]]:
+    tensor_infos: List[Tuple[torch.Tensor, int, int]] = []
+
+    offset = 0
+    for tensor in tensors:
+        tensor_numel = tensor.numel()
+        # record some data so we can restore the device tensor later
+        tensor_infos.append((tensor, offset, tensor_numel))
+        offset += tensor_numel
+
+    return tensor_infos
+
+
+def defragment(tensors: List[torch.Tensor]) -> torch.Tensor:
+    """move provided tensors into a contiguous flat buffer, with some additional
+    measures taken to reduce memory fragmentation"""
+    assert len(set(t.dtype for t in tensors)) == 1
+    assert len(set(t.device for t in tensors)) == 1
+
+    cpu_buffer = torch.empty(sum(p.numel() for p in tensors),
+                             dtype=get_only_unique_item(t.dtype for t in tensors),
+                             device="cpu")
+    tensor_infos: List[Tuple[torch.Tensor, int, int]] = get_mapping_to_flat_buffer(tensors)
+    orig_device = get_only_unique_item(t.device for t in tensors)
+
+    offset = 0
+    for tensor, offset, tensor_numel in tensor_infos:
+        # move the tensor from device memory to host memory
+        cpu_buffer.narrow(0, offset, tensor_numel).copy_(tensor)
+        tensor.data = torch.empty(0, dtype=tensor.dtype, device=tensor.device)
+
+    gc.collect()
+    get_accelerator().empty_cache()
+
+    # copy tensors (now flattened and contiguous) back to GPU
+    device_buffer = cpu_buffer.to(orig_device)
+
+    # restore device tensors
+    for tensor, offset, tensor_numel in tensor_infos:
+        tensor.data = device_buffer.narrow(0, offset, tensor_numel)
+
+    return device_buffer