mirrors/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deepspeed/runtime/zero/mics.py
Executable
+447
View File
@@ -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)