chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -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