From b2b56454bed92a48362d1046278789d549e0e333 Mon Sep 17 00:00:00 2001 From: "877825076@qq.com" <877825076@qq.com> Date: Tue, 24 Oct 2023 11:20:08 +0800 Subject: [PATCH] refactor lins --- deepspeed/runtime/engine.py | 50 + deepspeed/runtime/zero/config.py | 7 + deepspeed/runtime/zero/lins.py | 1216 +++++++++++++++++ deepspeed/runtime/zero/lins_utils.py | 286 ++++ .../runtime/zero/partition_parameters.py | 5 + .../zero/partitioned_param_coordinator.py | 18 +- deepspeed/runtime/zero/stage3.py | 46 +- test_lins.py | 201 +++ 8 files changed, 1813 insertions(+), 16 deletions(-) create mode 100644 deepspeed/runtime/zero/lins.py create mode 100644 deepspeed/runtime/zero/lins_utils.py create mode 100644 test_lins.py diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 8a8193ddd8f5..b46cc844ac3e 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -708,6 +708,18 @@ def zero_optimization_stage(self): def mics_shard_size(self): return self._config.mics_shard_size + + def lins_enable(self): + return self._config.zero_config.enable_lins + + def lins_param_partition_num(self): + return self._config.zero_config.lins_param_partition_num + + def lins_os_partition_num(self): + return self._config.zero_config.lins_os_partition_num + + def lins_grad_partition_num(self): + return self._config.zero_config.lins_grad_partition_num def zero_reduce_bucket_size(self): return self._config.zero_config.reduce_bucket_size @@ -1532,6 +1544,8 @@ def _configure_zero_optimizer(self, optimizer): ranks=[0]) if mics_shard_size > 0: return self._return_mics_optimizer(optimizer, timers) + elif self.lins_enable(): + return self._return_lins_optimizer(optimizer, timers) log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 @@ -1608,6 +1622,40 @@ def _return_mics_optimizer(self, basic_optimizer, timers): communication_data_type=self.communication_data_type) return optimizer + def _return_lins_optimizer(self, basic_optimizer, timers): + from deepspeed.runtime.zero.lins import LinS_Optimizer + model_dtype, gradient_accumulation_dtype = self.get_data_types() + print(f"enable LinS_Optimizer!!!!", flush=True) + optimizer = LinS_Optimizer(self.module, + basic_optimizer, + timers=timers, + ds_config=self.config, + static_loss_scale=self.loss_scale(), + dynamic_loss_scale=self.dynamic_loss_scale(), + dynamic_loss_args=self.dynamic_loss_scale_args(), + clip_grad=self.gradient_clipping(), + contiguous_gradients=self.zero_contiguous_gradients(), + reduce_bucket_size=self.zero_reduce_bucket_size(), + prefetch_bucket_size=self.zero_prefetch_bucket_size(), + max_reuse_distance=self.zero_max_reuse_distance(), + max_live_parameters=self.zero_max_live_parameters(), + param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), + dp_process_group=self.seq_data_parallel_group, + reduce_scatter=self.zero_reduce_scatter(), + overlap_comm=self.zero_overlap_comm(), + offload_optimizer_config=self.zero_offload_optimizer(), + offload_param_config=self.zero_offload_param(), + sub_group_size=self.zero_sub_group_size(), + mpu=self.mpu, + postscale_gradients=self.postscale_gradients(), + gradient_predivide_factor=self.gradient_predivide_factor(), + gradient_accumulation_steps=self.gradient_accumulation_steps(), + aio_config=self.aio_config(), + gradient_accumulation_dtype=gradient_accumulation_dtype, + communication_data_type=self.communication_data_type) + return optimizer + def _configure_eigenvalue(self): eigenvalue = Eigenvalue( verbose=self.eigenvalue_verbose(), @@ -1659,6 +1707,8 @@ def deepspeed_io(self, data_sampler=None, collate_fn=None, num_local_io_workers=None): + + print(type(dataset)) if not (self.is_map_style_dataset(dataset) or self.is_iterable_style_dataset(dataset)): raise ValueError("Training data must be a torch Dataset") diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py index 35d60b5b3290..c9994d10535f 100644 --- a/deepspeed/runtime/zero/config.py +++ b/deepspeed/runtime/zero/config.py @@ -293,6 +293,13 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel): Override nn.Module apply function, for Stage 3. """ + enable_lins: bool = False + lins_param_partition_num: int = Field(-1, new_param="lins_param_partition_num") + lins_os_partition_num: int = Field(-1, new_param="lins_os_partition_num") + lins_grad_partition_num: int = Field(-1, new_param="lins_grad_partition_num") + hierarchical_allgather: bool = False + + # Validators @validator("overlap_comm") def overlap_comm_valid(cls, field_value, values): diff --git a/deepspeed/runtime/zero/lins.py b/deepspeed/runtime/zero/lins.py new file mode 100644 index 000000000000..6e9a4ecaa383 --- /dev/null +++ b/deepspeed/runtime/zero/lins.py @@ -0,0 +1,1216 @@ +import os +import sys +import math +from typing import Iterable, Dict +# from enum import Enum +from typing import List +import logging +import torch +from torch import Tensor +# from torch.nn import Module +from torch.nn import Parameter +from functools import partial + +import deepspeed +import itertools +from deepspeed.utils import instrument_w_nvtx, logger, log_dist + +from deepspeed import comm as dist +from deepspeed.runtime.zero.mics_utils import (create_mics_comm_groups, scale_tensors) +# from deepspeed.runtime.zero.parameter_offload import (is_zero_param) +from deepspeed.runtime.zero.partition_parameters import Init, AllGatherCoalescedHandle, ZeroParamStatus, free_param, print_rank_0, _no_gather_coalesced, _dist_allgather_fn, AllGatherHandle, assert_ints_same_as_other_ranks +from deepspeed.accelerator import get_accelerator +from ..swap_tensor.partitioned_param_swapper import PartitionedParamStatus +from ..utils import get_only_unique_item, see_memory_usage +from deepspeed.utils.debug import (debug_param2name_id_shape_device, debug_param2name_id_shape_status) +from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3, reuse_buffers +from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced +from deepspeed.runtime.zero.lins_utils import zero35_g_p_reduce_scatter_coalesced, \ + zero35_g_p_all_gather_coalesced, zero35_debug, set_lins_parition_type, zero35_judge_gahter_boundary, global_parition_type, global_lins_utils, LinSProcessGroup + +# Copyright Shanghai AI Laboratory, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +ENBALE_MEM_DEBUG = False +ENBALE_COMM_DEBUG = False + + +class LinS_AllGatherHandle: + + def __init__(self, handle, param: Parameter, quantization=None) -> None: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to be available") + + self.__handle = handle + self.__param = param + self.__quantization = quantization + + def wait(self) -> None: + instrument_w_nvtx(self.__handle.wait)() + if self.__quantization: + instrument_w_nvtx(self.__quantization.quant_handle.wait)() + self.__param.data = self.__quantization.backend.dequantize( + self.__quantization.quantized_param, self.__quantization.scale_buffer).to(self.__param.device) + self.__param.ds_status = ZeroParamStatus.AVAILABLE + # 恢复原有param的shape等元信息 + self.__param = self.__param.zero35_restore_allgahter_ds_tensor(self.__param) + + +class LinS_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, + use_secondary_tensor=False, + forward=False, + quantization=None, + ) -> None: + self.allgather_handle = allgather_handle + self.params = params + self.partitions = partitions + self.world_size = world_size + self.use_secondary_tensor = use_secondary_tensor + self.forward = forward + self.complete = False + self.quantization = quantization + + for param in self.params: + if param.ds_status != ZeroParamStatus.INFLIGHT: + raise RuntimeError(f"expected param {param.ds_summary()} to not be available") + + @instrument_w_nvtx + def wait(self) -> None: + if self.complete: + return + + instrument_w_nvtx(self.allgather_handle.wait)() + + if self.quantization: + instrument_w_nvtx(self.quantization.quant_handle.wait)() + flat_tensor = self.quantization.backend.dequantize( + self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device) + + self.partitions: List[Parameter] = [] + for i in range(self.world_size): + self.partitions.append( + flat_tensor.narrow(0, self.quantization.partition_sz * i, self.quantization.partition_sz)) + + # split the single tensor out into individual tensors + param_offset = 0 + for param in self.params: + assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" + partitions: List[Tensor] = [] + ds_tensor_numel = param.ds_tensor.ds_numel + if self.use_secondary_tensor and not self.forward: + ds_tensor_numel *= param.ds_secondary_tensor_num_of_groups + for rank in range(self.world_size): + param_start = rank * ds_tensor_numel + if param_start < param.ds_numel: + part_to_copy = self.partitions[rank].narrow(0, param_offset, + min(param.ds_numel - param_start, ds_tensor_numel)) + partitions.append(part_to_copy) + + param.zero35_restore_allgahter_ds_tensor(param) + param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape) + param.ds_status = ZeroParamStatus.AVAILABLE + + for part_to_copy in partitions: + if not get_accelerator().is_synchronized_device(): + part_to_copy.record_stream(get_accelerator().current_stream()) + + param_offset += ds_tensor_numel + + if self.now_mico_step_id() == 0 and self.forward: + # 在每个micro_step 的第一次fwd时需要做全局的 all-gather + # zero35_debug(f"now_mico_step_id:{self.now_mico_step_id()}, forward:{forward}, skip at count: {self.all_gahter_count} get: {param}!", flush=True) + # TODO: 参数从 from os-partition to param-partition，只影响正确性，不影响性能测试 + pass + else: + # 在后续的 fwd/bwd 只需要做节点内的 all-gahter + zero35_g_p_all_gather_coalesced([param]) # partition_type + # zero35_debug(f"now_mico_step_id:{self.now_mico_step_id()}, forward:{forward}, do reshape finish at count: {self.all_gahter_count} get: {param}!", flush=True) + + self.complete = True + + +class LinS_Init(Init): + + def __init__(self, + module=None, + 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): + 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" + + self._dp_process_group = dist.get_world_group() + self.hierarchical_allgather = _ds_config.zero_config.hierarchical_allgather + self.world_size = dist.get_world_size(self._dp_process_group) + # self.local_device = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])) + + if mpu is not None: + self.model_parallel_group = mpu.get_model_parallel_group() + self.model_parallel_rank = mpu.get_model_parallel_rank() + self.model_parallel_size = dist.get_world_size(self.model_parallel_group) + self.data_parallel_size = self.world_size // self.model_parallel_size + else: + self.model_parallel_size = 1 + self.model_parallel_group = None + self.model_parallel_rank = 0 + self.data_parallel_size = self.world_size + + assert self.model_parallel_size == 1, "zero35现在不支持模型并行" + + global global_lins_utils + if global_lins_utils is None: + global_lins_utils = LinSProcessGroup( + self._dp_process_group, + lins_param_partition_num=_ds_config.zero_config.lins_param_partition_num, + lins_grad_partition_num=_ds_config.zero_config.lins_grad_partition_num, + lins_os_partition_num=_ds_config.zero_config.lins_os_partition_num) + + self._grad_process_group = global_lins_utils._grad_process_group + self._param_process_group = global_lins_utils._param_process_group + + # hack掉 partition param 的方法 + self._allgather_param = self._allgather_param_lins + self._allgather_params_coalesced = self._allgather_params_coalesced_lins + + super().__init__(module, data_parallel_group, mem_efficient_linear, remote_device, pin_memory, + config_dict_or_path, config, enabled, dtype, mpu) + + def _update_persist_config(self, ds_config): + # persistence_threshold 我们以节点内切分大小为准 + with set_lins_parition_type(partition_type="param"): + num_parition = self.num_partitions + + Init.apply_param_persistence = True + Init.param_persistence_threshold = ds_config.zero_config.param_persistence_threshold + Init.model_persistence_threshold = ds_config.zero_config.model_persistence_threshold // num_parition + + def _convert_to_deepspeed_param(self, param): + super()._convert_to_deepspeed_param(param) + print_rank_0("Lins: _convert_to_deepspeed_param", force=True) + # 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) + # return self.lins_all_gather_coalesced(params, param_buffers=None, **kwargs) + + def partition(param_list=None, backward=False, hierarchy=0, has_been_updated=False): + cls = param + print_rank_0(f"{'--'*hierarchy}----Zero35 Partitioning param {debug_param2name_id_shape_device(cls)}", + force=False) + if param_list is None: + param_list = [cls] + self._lins_partition(param_list, has_been_updated=has_been_updated) + + def padding_size(partition_type): + return self._lins_padding_size(param, partition_type) + + def aligned_size(partition_type): + return self._lins_aligned_size(param, partition_type) + + def partition_numel(partition_type=None): + return self._lins_partition_numel(param, partition_type) + + def zero35_hack_allgahter_ds_tensor(mico_step, forward): + + see_memory_usage(f"before zero35_hack_allgahter_ds_tensor, mico_step:{mico_step}, forward:{forward}") + if zero35_judge_gahter_boundary(mico_step, forward): + # gather boundary + partition_type = "os" + + assert hasattr(param, 'ds_numel'), 'zero35_hack_allgahter_ds_tensor input must be ds_param' + parition_num = self.get_world_size(partition_type) + + node_id = dist.get_rank() // self.zero35_parallel_size + partition_unit_size = param.ds_numel // parition_num + param_ds_tensor = param.ds_tensor.view(-1, partition_unit_size) + + # backup ds_tensor + param.ds_numel_backup = param.ds_tensor.ds_numel + param.ds_tensor_backup = param.ds_tensor.data + + # hack ds_tensor + param.ds_tensor = param_ds_tensor[node_id] + assert param.ds_tensor.storage().data_ptr() == param_ds_tensor.storage().data_ptr() + + param.ds_tensor.ds_numel = partition_unit_size + param.ds_tensor.status = PartitionedParamStatus.AVAILABLE + param.ds_tensor.final_location = None + param.ds_tensor.is_first_fwd_all_gahter = True + + # zero35_debug(f"zero35_hack_allgahter_ds_tensor DEBUG: mico_step: {mico_step}, forward:{forward}, param.ds_numel : {param.ds_numel}, get : {param_ds_tensor}, partition_type:{partition_type}, partition_unit_size:{partition_unit_size}", flush=True) + else: + partition_type = "param" + partition_unit_size = param.ds_tensor.ds_numel + # zero35_debug(f"zero35_hack_allgahter_ds_tensor DEBUG: mico_step: {mico_step}, forward:{forward}, SKIP hack, partition_unit_size:{partition_unit_size}", flush=True) + + see_memory_usage(f"after zero35_hack_allgahter_ds_tensor, mico_step:{mico_step}, forward:{forward}") + return partition_unit_size + + def zero35_restore_allgahter_ds_tensor(param): + # TODO:(wgt) + if param.ds_tensor.is_first_fwd_all_gahter == True: + # zero35_debug(f"now ds_tensor numel: {param.ds_tensor.ds_numel}, backup numel: {param.ds_numel_backup}") + # zero35_debug(f"now ds_tensor data: {param.ds_tensor.data}, backup data: {param.ds_tensor_backup}") + + param.ds_tensor.data = param.ds_tensor_backup + param.ds_tensor.is_first_fwd_all_gahter = False + param.ds_tensor.ds_numel = param.ds_numel_backup + + # def get_partition_group(param): + # global + # return param. + + def lins_all_gahter_coalesced(params: Iterable[Parameter], + forward: bool = True, + safe_mode: bool = False, + quantize: bool = False, + mico_step: int = 0): + if self.zero35_judge_gahter_boundary(mico_step, forward) and \ + self.hierarchical_allgather: + return self.zero35_hierarchical_all_gather_params(params, forward, safe_mode, quantize, mico_step) + else: + return self.lins_all_gather_coalesced(params, forward, safe_mode, quantize, mico_step) + + # change the all_gather_coalesced method + param.all_gather_coalesced = lins_all_gahter_coalesced + param.partition = partition + param.padding_size = padding_size + param.aligned_size = aligned_size + param.partition_numel = partition_numel + # param.get_partition_group = get_partition_group + + param.zero35_hack_allgahter_ds_tensor = zero35_hack_allgahter_ds_tensor + param.zero35_restore_allgahter_ds_tensor = zero35_restore_allgahter_ds_tensor + + def _lins_partition(self, param_list, force=False, has_been_updated=False): + for param in param_list: + print_rank_0(f"Before Zero35 Partitioning Param {param.ds_id}", force=False) + self._lins_partition_param(param, has_been_updated=has_been_updated) + param.ds_status = ZeroParamStatus.NOT_AVAILABLE + + def _lins_partition_numel(self, param, partition_type=None): + if partition_type is not None: + with set_lins_parition_type(partition_type=partition_type): + tensor_size = self._lins_aligned_size(param, partition_type) + return tensor_size // self.num_partitions + else: + tensor_size = self._lins_aligned_size(param, partition_type) + return tensor_size // self.num_partitions + + def _lins_padding_size(self, param, partition_type=None): + # param group 切分的 parition_unit 的大小要大于 os group + # 如何返回正确的 padding ? + if partition_type is not None: + with set_lins_parition_type(partition_type): + remainder = param.ds_numel % self.num_partitions + return (self.num_partitions - remainder) if remainder else 0 + else: + remainder = param.ds_numel % self.num_partitions + return (self.num_partitions - remainder) if remainder else 0 + + def _lins_aligned_size(self, param, partition_type=None): + if partition_type is not None: + with set_lins_parition_type(partition_type): + return param.ds_numel + self._lins_padding_size(param, partition_type) + else: + return param.ds_numel + self._lins_padding_size(param, partition_type) + + def get_partition_dp_group(self, param): + return param.ds_process_group + + def get_partition_rank(self): + """subclass can overload to specify different relative rank in + parameter partition group""" + return dist.get_rank(self.get_process_group()) + + @property + def num_partitions(self): + global global_parition_type + _parition_type = global_parition_type['type'] + if _parition_type == "os": + return dist.get_world_size(self._dp_process_group) + elif _parition_type == "grad": + return dist.get_world_size(self._grad_process_group) + elif _parition_type == "param": + return dist.get_world_size(self._param_process_group) + else: + assert False, f"unknown partition_type: {_parition_type}" + + def get_dp_process_group(self): + """ Return the communication group with all data-parallel ranks """ + return self.ds_process_group + + def get_process_group(self): + """ Return the communication group with all data-parallel ranks """ + global global_parition_type + _parition_type = global_parition_type['type'] + if _parition_type == "os": + return self._dp_process_group + elif _parition_type == "grad": + return self._grad_process_group + elif _parition_type == "param": + return self._param_process_group + else: + assert False, f"unknown partition_type: {_parition_type}" + + def get_param_process_group(self, param): + """ Return the communication group with all data-parallel ranks """ + return param.get_param_process_group() + + """ + 通信相关需要重载的方法，包括 partition param 和 stage3 的 + """ + + def _allgather_param_lins(self, param, async_op=False, hierarchy=0): + partition_size = param.ds_tensor.ds_numel + + with set_lins_parition_type(partition_type="param"): + num_partitions = param.partition_numel + aligned_param_size = param.aligned_size + partition_all_gather_group = self.get_partition_group(param) + + tensor_size = partition_size * num_partitions + assert tensor_size == aligned_param_size, f'param id {param.ds_id} aligned size {aligned_param_size} does not match tensor size {tensor_size}' + print_rank_0( + f"{'--'* hierarchy}---- Before allocating allgather param {debug_param2name_id_shape_status(param)} partition size={partition_size}" + ) + + see_memory_usage( + f'Before allocate allgather param {debug_param2name_id_shape_status(param)} partition_size={partition_size} ', + force=False) + flat_tensor = torch.zeros(aligned_param_size, dtype=param.dtype, device=param.device).view(-1) + see_memory_usage( + f'After allocate allgather param {debug_param2name_id_shape_status(param)} {aligned_param_size} {partition_size} ', + force=False) + + get_accelerator().synchronize() + + print_rank_0( + f"{'--'* hierarchy}----allgather param with {debug_param2name_id_shape_status(param)} partition size={partition_size}" + ) + handle = dist.all_gather_into_tensor(flat_tensor, + param.ds_tensor.to(get_accelerator().device_name()), + group=partition_all_gather_group, + async_op=async_op) + replicated_tensor = flat_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape) + param.data = replicated_tensor.data + return handle + + # all_gather_coalesced 和 _allgather_params_coalesced 是两个名字很像，但是干不同事情的 all-gather + def lins_all_gather_coalesced(self, + params: Iterable[Parameter], + forward: bool = True, + safe_mode: bool = False, + quantize: bool = False, + mico_step: int = 1) -> LinS_AllGatherCoalescedHandle: + # fetches from nvme if the partition is not available and in nvme + self._ensure_availability_of_partitioned_params(params) + + with set_lins_parition_type(partition_type="param"): + num_partitions = self.num_partitions + + if num_partitions == 1: + return _no_gather_coalesced(params) + + for param in params: + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(param.ds_summary()) + param.ds_status = ZeroParamStatus.INFLIGHT + + #use appropriate all gather process group + partition_type = "os" if zero35_judge_gahter_boundary(mico_step, forward) else "param" + with set_lins_parition_type(partition_type): + ds_process_group = self.get_process_group() + rank_in_group = self.get_partition_rank() + world_size = self.num_partitions + + use_secondary_tensor = False + + params = sorted(params, key=lambda p: p.ds_id) + + if logger.isEnabledFor(logging.DEBUG): + print_rank_0(f"-allgather_coalesced: {[p.ds_id for p in params]}") + + if safe_mode: + assert_ints_same_as_other_ranks([p.ds_id for p in params]) + assert_ints_same_as_other_ranks([p.ds_tensor.ds_numel for p in params]) + + if len(params) == 1: + # have an opportunity to avoid some intermediate memory allocations + param, = params + self.zero35_hack_allgahter_ds_tensor(param, mico_step, forward) + + buffer_size = math.ceil(param.ds_numel / world_size) * world_size + param_ds_tensor = param.ds_tensor + param_buffer = torch.empty( + buffer_size, + dtype=param_ds_tensor.dtype if not quantize else torch.int8, + device=get_accelerator().current_device_name(), + requires_grad=False, + ) + handles = _dist_allgather_fn( + param_ds_tensor.to(get_accelerator().current_device_name()), + param_buffer, + ds_process_group, + ) + param.data = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device) + return AllGatherHandle(handles, param) + else: + partition_sz = 0 + for param in params: + partition_sz += self.zero35_hack_allgahter_ds_tensor(param, mico_step, forward) + + flat_tensor = torch.empty(partition_sz * world_size, + dtype=get_only_unique_item(p.ds_tensor.dtype + for p in params) if not quantize else torch.int8, + device=get_accelerator().current_device_name(), + requires_grad=False) + partitions: List[Parameter] = [] + for i in range(world_size): + partitions.append(flat_tensor.narrow(0, partition_sz * i, partition_sz)) + instrument_w_nvtx(torch.cat)([p.ds_tensor.to(get_accelerator().current_device_name()) for p in params], + out=partitions[rank_in_group]) + handle = _dist_allgather_fn(partitions[rank_in_group], flat_tensor, ds_process_group) + + return LinS_AllGatherCoalescedHandle( + allgather_handle=handle, + params=params, + partitions=partitions, + world_size=world_size, + use_secondary_tensor=use_secondary_tensor, + forward=forward, + ) + + def _allgather_params_coalesced_lins(self, param_list, hierarchy=0, quantize=False): + """ blocking call + avoid explicit memory copy in _allgather_params + """ + assert False + + @instrument_w_nvtx + def _lins_partition_param(self, param, buffer=None, has_been_updated=False): + """ + zero35 进行 partition的基本单元是按照 dp 范围切分，这些切分部分我们称之为 'partition_unit' + 每个rank可能会分到多个 'partition_unit' + + Args: + param (_type_): _description_ + buffer (_type_, optional): _description_. Defaults to None. + has_been_updated (bool, optional): _description_. Defaults to False. + """ + assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" + global reuse_buffers + print_rank_0(f"Param id {param.ds_id} status is {param.ds_status}", force=False) + # # zero35_debug(f"do _lins_partition_param!") + + param_comm_group = self._param_process_group + dp_comm_group = self._dp_process_group + param_num_partitions = dist.get_world_size(param_comm_group) + dp_num_partitions = dist.get_world_size(dp_comm_group) + + # assert param_num_partitions == 8, "zero35 split param in local node device" + # if param_num_partitions != 8: + # print(f"zero35 split param worldisze: {param_num_partitions}", flush=True) + + if param.ds_status is ZeroParamStatus.AVAILABLE: + print_rank_0(f"Partitioning param id {param.ds_id} reuse buffers {reuse_buffers}", force=False) + + if param.ds_tensor is not None and not has_been_updated: ##param already partitioned + see_memory_usage(f'Before partitioning param 2:{param.ds_id} {param.shape}', force=False) + # param.data does not store anything meaningful in partitioned state + free_param(param) + see_memory_usage(f'After partitioning param 2:{param.ds_id} {param.shape}', force=False) + return + + tensor_size_param = self._lins_aligned_size(param, partition_type="param") + tensor_size_dp = self._lins_aligned_size(param, partition_type="os") + + # assert tensor_size_dp == tensor_size_param, f"different padding size: tensor_size_dp:{tensor_size_dp} = tensor_size_param: {tensor_size_param}" + tensor_size = tensor_size_dp + + partition_size = tensor_size // param_num_partitions + unit_partition_size = tensor_size // dp_num_partitions + + assert partition_size % unit_partition_size == 0 + + if param.ds_tensor is None: + final_location = None + # assert param.ds_persist is False, "ds_persist可能会有ug" + if param.ds_persist: + device = self.local_device + else: + device = self.remote_device + + # buffer 大小仍然是 partition_size 这么大 + partitioned_tensor = torch.empty(partition_size, dtype=param.dtype, device=device) + partitioned_tensor.requires_grad = False + param.ds_tensor = partitioned_tensor # 被切分后的tensor + param.ds_tensor.ds_numel = partition_size # ds_numel 是 buffer大小，等于partition_size，但每个param实际上仍然是被切分了 1/dp 份 + param.unit_partition_size = unit_partition_size # 这里存一份 unit_partition_size，给 _unflatten_partitioned_parameters 用 + param.ds_tensor.is_first_fwd_all_gahter = False + + param.ds_tensor.status = PartitionedParamStatus.AVAILABLE + param.ds_tensor.final_location = final_location + + partition_unit_num = partition_size // unit_partition_size + + assert tensor_size % partition_unit_num == 0 + partition_stride = tensor_size // partition_unit_num + + # 这里需要改成分段拷贝 + try: + with set_lins_parition_type(partition_type="param"): + print( + f"partition type: {global_parition_type}, self.get_partition_rank():{self.get_partition_rank()}", + flush=True) + offset = unit_partition_size * self.get_partition_rank() + one_dim_param = param.contiguous().view(-1) + + # # zero35_debug(f"Rank: {os.environ['SLURM_PROCID']}, partition_unit_num: {partition_unit_num}, partition_stride:{partition_stride}, offset:{offset}, ", flush=True) + for pdx in range(partition_unit_num): + start = offset + partition_stride * pdx + sub_start = unit_partition_size * pdx + # # zero35_debug(f"Rank: {os.environ['SLURM_PROCID']}, pdx:{pdx},start: {start}, sub_start:{sub_start}" , flush=True) + + if self.get_partition_rank() == param_num_partitions -1 \ + and pdx == partition_unit_num - 1: # 只有最后一块 1/dp 的 partition unit 需要补齐 padding + param.ds_tensor[sub_start:] = one_dim_param[start:] + else: + param.ds_tensor[sub_start:sub_start + + unit_partition_size] = one_dim_param[start:start + unit_partition_size] + except Exception as e: + print(f"catch exception: {e}", flush=True) + import pdb + pdb.set_trace() + + # if os.environ['SLURM_PROCID'] == '0': + # # zero35_debug(f"Rank: {os.environ['SLURM_PROCID']}, partition param done {param.ds_tensor}", flush=True) + + see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=ENBALE_MEM_DEBUG) + # # zero35_debug(f"Before partitioning param ID {param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}") + free_param(param) + # # zero35_debug(f"After partitioning param ID {param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}") + see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=ENBALE_MEM_DEBUG) + + @instrument_w_nvtx + def zero35_hierarchical_all_gather_params(self, + params: Iterable[Parameter], + forward: bool = True, + safe_mode: bool = False, + quantize: bool = False, + mico_step: int = 1): + + params_buffers = None + + # self._ensure_availability_of_partitioned_params(params) + from deepspeed.runtime.zero.mics import MiCS_AllGatherCoalescedHandle + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + + 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) + + local_rank = dist.get_rank(group=self.zero35_group) + inter_node_comm_group = self.zero35_hierarchical_group + intra_node_comm_group = self.zero35_group + intra_param_shard_size = dist.get_world_size(intra_node_comm_group) + assert intra_param_shard_size == 8 + + dp_param_shard_size = dist.get_world_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): + self.zero35_hack_allgahter_ds_tensor(p, mico_step, forward) + param_size = p.ds_tensor.ds_numel * dp_param_shard_size # unit_size * dp world size + #zero35_debug(F"param_size :{param_size / dp_param_shard_size}, intra_param_shard_size: {intra_param_shard_size}, dp_param_shard_size, :{dp_param_shard_size/dp_param_shard_size}, p.ds_tensor.ds_numel:{p.ds_tensor.ds_numel/dp_param_shard_size}, p.ds_tensor.ds_shape:{p.ds_shape}", force=True) + 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 = [] + try: + for i, p in enumerate(params): + inter_size = p.ds_tensor.ds_numel * inter_node_size + #zero35_debug(f"p.ds_tensor.ds_numel:{p.ds_tensor.ds_numel/dp_param_shard_size}, inter_node_size:{inter_node_size}, inter_size: {inter_size/dp_param_shard_size}, p.ds_tensor.ds_shape:{p.ds_shape}", force=True) + _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)) + #zero35_debug(f"inter input unit parma size: {p.ds_tensor.data.numel()/dp_param_shard_size}, inter_size: {inter_size/dp_param_shard_size}", force=True) + except Exception as e: + import time + if dist.get_rank() == 0: + print(e, flush=True) + time.sleep(10000) + + # 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 + try: + #zero35_debug(f"inter_node_size:{inter_node_size}, intra_node_size:{intra_node_size}, p.ds_tensor.ds_numel:{p.ds_tensor.ds_numel}, local_rank:{local_rank}", force=True) + 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 + #zero35_debug(f"intra input unit parma size: {p.ds_tensor.data.numel()/dp_param_shard_size}, intra_chunk_size: {intra_chunk_size/dp_param_shard_size}", force=True) + _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) + except Exception as e: + import time + if dist.get_rank() == 0: + print(e, flush=True) + time.sleep(10000) + + 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 + #zero35_debug(f"finish pre allgather param.data:{param.data.numel()/dp_param_shard_size}, ds_tensor.numel: {param.ds_tensor.numel()/dp_param_shard_size}", force=True) + + # import time + # time.sleep(10000) + return MiCS_AllGatherCoalescedHandle( + allgather_handle=all_gather_handle, + params=params, + partitions=[], + world_size=intra_param_shard_size, + ) + + +class LinS_Optimizer(DeepSpeedZeroOptimizer_Stage3): + """ + MiCS Optimizer + """ + + def __init__(self, + module, + init_optimizer, + timers, + ds_config, + static_loss_scale=1, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + contiguous_gradients=True, + reduce_bucket_size=500000000, + 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, + reduce_scatter=True, + overlap_comm=False, + offload_optimizer_config=None, + offload_param_config=None, + sub_group_size=1000000000000, + mpu=None, + clip_grad=0, + gradient_accumulation_dtype=torch.float16, + communication_data_type=torch.float16, + postscale_gradients=True, + gradient_predivide_factor=1, + gradient_accumulation_steps=1, + elastic_checkpoint=False, + aio_config=None): + + log_dist("Init LinS optimizer", ranks=[0]) + + self._param_partition_unit_size = [] + + global global_lins_utils + self._grad_process_group = global_lins_utils._grad_process_group + self._param_process_group = global_lins_utils._param_process_group + + super().__init__(module, init_optimizer, timers, ds_config, static_loss_scale, dynamic_loss_scale, + dynamic_loss_args, verbose, contiguous_gradients, reduce_bucket_size, prefetch_bucket_size, + max_reuse_distance, max_live_parameters, param_persistence_threshold, + model_persistence_threshold, dp_process_group, reduce_scatter, overlap_comm, + offload_optimizer_config, offload_param_config, sub_group_size, mpu, clip_grad, + gradient_accumulation_dtype, communication_data_type, postscale_gradients, + gradient_predivide_factor, gradient_accumulation_steps, elastic_checkpoint, aio_config) + first_param = next(module.parameters()) + # overload the dp_process_group and partition_count + # self.dp_process_group = first_param.comm.param_shard_group + # self.partition_count = first_param.comm.param_shard_size + + def cal_tensor_size(self, param): + ds_t_size = 0 + dtype = param.dtype + if dtype == torch.float32: + unit_byte = 4 + elif dtype == torch.float16 or dtype == torch.bfloat16: + unit_byte = 2 + else: + assert False, f"Unexpected dtype: {dtype}!" + + if hasattr(param, 'ds_tensor'): + ds_t_size = param.ds_tensor.numel() + torch_t_size = param.numel() + assert ds_t_size >= torch_t_size, f"{ds_t_size} >= {torch_t_size}" + assert param.ds_tensor.dtype == param.dtype + if param.grad is not None: + assert not hasattr(param.grad, 'ds_tensor') + assert param.grad.numel() == 0 + else: + ds_t_size = param.numel() + + return ds_t_size * unit_byte + + def format_size(self, size): + if size < 1024: + return f"{size} B" + elif size >= 1024 and size < 1024**2: + return f"{size / 1024:.2f} KB" + elif size >= 1024**2 and size < 1024**3: + return f"{size / 1024**2:.2f} MB" + else: + return f"{size / 1024**3:.2f} GB" + + def cal_stage3_mem_usage(self, all_params): + grad_partitions_flat_buffer_size = self.cal_tensor_size(self.grad_partitions_flat_buffer) + ipg_bucket_flat_buffer_size = self.cal_tensor_size(self.__ipg_bucket_flat_buffer) + + fp32_partitioned_groups_flat_size = 0 + for i in range(len(self.fp32_partitioned_groups_flat)): + partitioned_groups_flat = self.fp32_partitioned_groups_flat[i] + fp32_partitioned_groups_flat_size += self.cal_tensor_size(partitioned_groups_flat) + + # zero35_debug(f"Grad size: {self.format_size(grad_partitions_flat_buffer_size)}", force=True) + # zero35_debug(f"fp32 os size: {self.format_size(fp32_partitioned_groups_flat_size)}", force=True) + # zero35_debug(f"os size total: {self.format_size(fp32_partitioned_groups_flat_size * 3)}", force=True) + # zero35_debug(f"ipg_bucket_flat_buffer_size: {self.format_size(ipg_bucket_flat_buffer_size)}", force=True) + + param_size = 0 + os_size = 0 + for param in all_params: + param_size += self.cal_tensor_size(param) + + # zero35_debug(f"param_size: {self.format_size(param_size)}", force=True) + # zero35_debug(f"all size: {self.format_size(param_size + grad_partitions_flat_buffer_size + ipg_bucket_flat_buffer_size + fp32_partitioned_groups_flat_size)}", force=True) + + """ + 切分相关需要重载的函数 + """ + + def bookkeeping_param_group(self, param_groups): + print("bookkeeping_param_group!!!", flush=True) + with set_lins_parition_type(partition_type="os"): + for param_group_idx, param_group in enumerate(param_groups): + for sub_group in param_group: + sub_group_idx = len(self.fp16_groups) + + # record sub group and partitions + self.fp16_groups.append(sub_group) + self.fp16_partitioned_groups.append([param.ds_tensor for param in sub_group]) + + # record total elements of parameter partitions in sub group + assert self._param_partition_unit_size is not None + self._param_partition_unit_size.append([param.unit_partition_size for param in sub_group]) + self.fp16_partitioned_groups_flat_numel.append(sum(p.partition_numel() for p in sub_group)) + param_process_group = self._param_process_group + + # record sub group -> group mapping + self.sub_group_to_group_id[sub_group_idx] = param_group_idx + + # record padding required to align group to world size (only applies to last rank) + rank_requires_padding = dist.get_rank( + param_process_group) == dist.get_world_size(param_process_group) - 1 + self.groups_padding.append( + [p.padding_size(partition_type="os") if rank_requires_padding else 0 for p in sub_group]) + + def get_sub_p_g_parition(self, ds_param, grad=None): + assert hasattr(ds_param, 'ds_numel'), 'get_sub_p_g_parition input must be ds_param' + zero35_rank = dist.get_rank() // self.zero35_parallel_size # TODO, remove hard code + partition_unit_size = self.get_partition_unit_size(ds_param.ds_numel) + if grad is not None: + reshape_grad = grad.view(-1, partition_unit_size)[zero35_rank] + assert reshape_grad.storage().data_ptr() == grad.storage().data_ptr() + return reshape_grad + else: + reshape_param = ds_param.ds_tensor.view(-1, partition_unit_size)[zero35_rank] + assert reshape_param.storage().data_ptr() == ds_param.ds_tensor.storage().data_ptr() + return reshape_param + + @instrument_w_nvtx + def independent_gradient_partition_epilogue(self): + self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) + self.__reduce_and_partition_ipg_grads() + self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) + + if not get_accelerator().is_synchronized_device(): + self.reduce_and_partition_stream.synchronize() + + #in case of cpu offload, averaged gradients are already in fp32_partitioned_groups_flat.grad + #TODO: use a similar code path for both cpu_offload and non-cpu offload + if not self.offload_optimizer: + for i, sub_group in enumerate(self.fp16_groups): + #TODO: This is redundant + self.averaged_gradients[i] = [ + self.get_sub_p_g_parition(param, self.__param_id_to_grad_partition[param.ds_id]) + if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group + ] + # this method gets called after every backward. need to increment + # here because if it gets incremented in backward() the micro step + # id will be off by one when we do the reduce and partition at the. + # start of this method. + # TODO. make this less error prone + self.micro_step_id += 1 + self._get_param_coordinator(training=True).micro_step_id += 1 + + def _setup_for_real_optimizer(self): + see_memory_usage("Before creating fp32 partitions", force=True) + self._create_fp32_partitions() + see_memory_usage("After creating fp32 partitions", force=True) + dist.barrier() + + # To support pipelined optimizer swapping + self._create_next_swappable_fp32_groups() + + see_memory_usage("Before initializing optimizer states", force=True) + + self.initialize_optimizer_states() + see_memory_usage("After initializing optimizer states", force=True) + dist.barrier() + + if dist.get_rank() == 0: + logger.info(f"optimizer state initialized") + + # IPG + if self.contiguous_gradients: + self.__ipg_bucket_flat_buffer: Tensor = torch.empty(self.reduce_bucket_size, + dtype=self.dtype, + device=get_accelerator().current_device_name()) + + self.grad_partitions_flat_buffer = None + self.__param_id_to_grad_partition: Dict[int, Tensor] = {} + + all_params = list(itertools.chain.from_iterable(self.fp16_groups)) + + with set_lins_parition_type(partition_type="grad"): + self.grad_partitions_flat_buffer: Tensor = torch.zeros(sum(p.partition_numel() for p in all_params), + dtype=self.gradient_accumulation_dtype, + device=self.device) + if self.offload_optimizer_pin_memory: + self.grad_partitions_flat_buffer = get_accelerator().pin_memory(self.grad_partitions_flat_buffer) + + offset = 0 + for param in all_params: + # self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( + # 0, offset, param.partition_numel()) + # offset += param.partition_numel() + self.__param_id_to_grad_partition[param.ds_id] = \ + self.grad_partitions_flat_buffer.narrow(0, offset, param.partition_numel()) + offset += param.partition_numel() + + see_memory_usage("End of _setup_for_real_optimizer", force=True) + return all_params + + def _create_fp16_sub_groups(self, params_group): + params_group_numel = sum([param.partition_numel(partition_type="os") for param in params_group]) + sub_group_size = self.sub_group_size + + if sub_group_size is None or sub_group_size >= params_group_numel: + return [params_group] + + sub_groups = [] + sub_group = [] + local_sub_group_size = 0 + for param in params_group: + + sub_group.append(param) + local_sub_group_size += param.partition_numel(partition_type="os") + + if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): + + sub_groups.append(sub_group) + + sub_group = [] + local_sub_group_size = 0 + + return sub_groups + + """ + 通信相关需要重载的函数 + """ + + @instrument_w_nvtx + def __avg_scatter_grads(self, params_to_reduce: List[Parameter]) -> List[Tensor]: + """average gradients and scatter partitions across ranks""" + see_memory_usage(f"before __avg_scatter_grads, boundary: {self.zero35_judge_grad_boundary()}", + force=ENBALE_MEM_DEBUG) + + full_grads_for_rank = [p.grad for p in params_to_reduce] + if self.communication_data_type != self.dtype: + full_grads_for_rank = [g.to(self.communication_data_type) for g in full_grads_for_rank] + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0: + full_grads_for_rank = [g.div(self.gradient_predivide_factor) for g in full_grads_for_rank] + + if not self.zero35_judge_grad_boundary(): + # 梯度累加 + # 假设 dp_world_size = 8, 每个节点只有4张卡 + # os 切分: + # 节点1: [[0], [1], [2], [3]], 节点2: [[4], [5], [6], [7]] + # grad/param切分: + # 节点1：[[0, 4], [1, 5] ,[2, 6], [3, 7]], 节点2: [[0, 4], [1, 5] ,[2, 6], [3, 7]] + # [[0, 1, 2, 3], [4, 5, 6, 7]] -> [[0], 1, 2, 3, [16], 5, 6, 7] + # [[0, 1, 2, 3], [4, 5, 6, 7]] -> [0, [4], 2, 3, 4, [20], 6, 7] + # [[0, 1, 2, 3], [4, 5, 6, 7]] -> [0, 1, [8], 3, 4, 5, [24], 7] + # [[0, 1, 2, 3], [4, 5, 6, 7]] -> [0, 1, 2, [12], 4, 5, 6, [28]]] + full_grads_for_rank, scatter_comm_group = zero35_g_p_reduce_scatter_coalesced(full_grads_for_rank, + partition_type="grad") + + # zero35_debug(f"Rank: {os.environ['SLURM_PROCID']}, mico_step: {self.micro_step_id}, before __avg_scatter_grads, after grad_tensor_list : {full_grads_for_rank}", flush=True) # before grad_tensor_list: {tensor_list_debug}, \ + else: + # boundary + # 在 boundary 阶段，进行 dp 范围的 reduce-scatter， + # [0, 1, 2, 3, 4, 5, 6, 7] -> [[0], 1, 2, 3, 4, 5, 6, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, [8], 2, 3, 4, 5, 6, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, [16], 3, 4, 5, 6, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 2, [24], 4, 5, 6, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 2, 3, [32], 5, 6, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 2, 3, 4, [40], 6, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 2, 3, 4, 5, [48], 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 2, 3, 4, 5, 6, [56]] + # zero35_debug(f"Rank: {os.environ['SLURM_PROCID']}, mico_step: {self.micro_step_id}, skip zero35_g_p_reduce_scatter_coalesced, before __avg_scatter_grads, before grad_tensor_list", flush=True) + scatter_comm_group = self.get_dp_process_group(partition_type="os") + + local_world_size = get_accelerator().device_count() + global_world_size = dist.get_world_size() + num_nodes = global_world_size // local_world_size + + if ENBALE_COMM_DEBUG: + if dist.get_rank() == 0: + numels = sum([p.numel() for p in full_grads_for_rank]) + print(f"__avg_scatter_grads, \ +comm_group:{dist.get_world_size(scatter_comm_group)} \ +nums:{numels}, \ +size: {numels* full_grads_for_rank[0].element_size()/ (1024**2):.4f} MB", + flush=True) + + # # zero35_debug(f"before __avg_scatter_grads: {full_grads_for_rank}", flush=True) + grad_partitions_for_rank = reduce_scatter_coalesced(full_grads_for_rank, scatter_comm_group) + # # zero35_debug(f"after __avg_scatter_grads: {full_grads_for_rank}", flush=True) + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0 and self.gradient_predivide_factor != dist.get_world_size( + self.dp_process_group): + grad_partitions_for_rank = [g.mul(self.gradient_predivide_factor) for g in grad_partitions_for_rank] + + if self.communication_data_type != self.dtype: + grad_partitions_for_rank = [g.to(self.dtype) for g in grad_partitions_for_rank] + + see_memory_usage(f"after __avg_scatter_grads, boundary: {self.zero35_judge_grad_boundary()}", + force=ENBALE_MEM_DEBUG) + + return grad_partitions_for_rank + + @instrument_w_nvtx + def __avg_scatter_grads_hierarchical(self, params_to_reduce: List[Parameter]) -> List[Tensor]: + see_memory_usage(f"before __avg_scatter_grads, boundary: {self.zero35_judge_grad_boundary()}", + force=ENBALE_MEM_DEBUG) + full_grads_for_rank = [p.grad for p in params_to_reduce] + + if self.communication_data_type != self.dtype: + full_grads_for_rank = [g.to(self.communication_data_type) for g in full_grads_for_rank] + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0: + full_grads_for_rank = [g.div(self.gradient_predivide_factor) for g in full_grads_for_rank] + + inter_comm_group = self.zero35_hierarchical_group + intra_comm_group = self.zero35_group + + assert dist.get_world_size(intra_comm_group) == 8 + + def count_list_params_numel(grads_list): + all_numel = 0 + for grad in grads_list: + all_numel += grad.numel() + return all_numel + + if not self.zero35_judge_grad_boundary(): + # full_grads_for_rank, scatter_comm_group = zero35_g_p_reduce_scatter_coalesced(full_grads_for_rank, partition_type="grad") + # Replace + grad_partitions_for_rank = reduce_scatter_coalesced(full_grads_for_rank, intra_comm_group) + else: + # boundary + # inter reduce-scatter + zero35_debug(f"before first educe_scatter: {count_list_params_numel(full_grads_for_rank)}", force=False) + grad_partitions_for_rank = reduce_scatter_coalesced(full_grads_for_rank, inter_comm_group) + zero35_debug(f"after first educe_scatter: {count_list_params_numel(grad_partitions_for_rank)}", + force=False) + + # Replace + # intra reduce-scatter + grad_partitions_for_rank = reduce_scatter_coalesced(grad_partitions_for_rank, intra_comm_group) + zero35_debug(f"after sec educe_scatter: {count_list_params_numel(grad_partitions_for_rank)}", force=False) + + if self.postscale_gradients and self.gradient_predivide_factor != 1.0 and self.gradient_predivide_factor != dist.get_world_size( + self.dp_process_group): + grad_partitions_for_rank = [g.mul(self.gradient_predivide_factor) for g in grad_partitions_for_rank] + + if self.communication_data_type != self.dtype: + grad_partitions_for_rank = [g.to(self.dtype) for g in grad_partitions_for_rank] + + see_memory_usage(f"after __avg_scatter_grads, boundary: {self.zero35_judge_grad_boundary()}", + force=ENBALE_MEM_DEBUG) + + return grad_partitions_for_rank + + def _unflatten_partitioned_parameters(self, sub_group_id): + + def get_sub_p_g_parition_from_torch_tensor(torch_tensor, i): + zero35_rank = dist.get_rank() // self.zero35_parallel_size # TODO, remove hard code + # zero35_debug(f"self._param_partition_unit_size[sub_group_id]: {self._param_partition_unit_size[sub_group_id]}", flush=ENBALE_MEM_DEBUG) + return torch_tensor.reshape(-1, self._param_partition_unit_size[sub_group_id][i])[zero35_rank] + + sub_fp16_partitioned_groups_flat = self.fp16_partitioned_groups_flat[sub_group_id] + sub_fp16_partitioned_groups = [ + get_sub_p_g_parition_from_torch_tensor(param, i) + for i, param in enumerate(self.fp16_partitioned_groups[sub_group_id]) + ] + + updated_params = self.unflatten(sub_fp16_partitioned_groups_flat, sub_fp16_partitioned_groups) + + # (TODO):wgt 为了简单，这里直接注释掉了，为了防止 intra 划分的 ds_tensor 覆盖 dp 划分的 ds_tensor + # for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): + # partitioned_param.data = q.data + + @instrument_w_nvtx + def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: + offload_fp32_gradients = {} + offload_fp32_offsets = {} + buffers = [] + see_memory_usage(f"before partition_grads", force=ENBALE_MEM_DEBUG) + for param, grad_partition in zip(params_to_release, grad_partitions): + + if self.micro_step_id == self.gradient_accumulation_steps - 1: # bounary + grad_rank = dist.get_rank(self.get_dp_process_group(partition_type="os")) + else: + grad_rank = dist.get_rank(self.get_dp_process_group(partition_type="grad")) + + with set_lins_parition_type(partition_type="grad"): + contains_real_data = param.partition_numel() * grad_rank < param.ds_numel + + if not contains_real_data: + # this grad partition is empty - don't need to do anything + param.grad = None + continue + + # move or accumulate gradient partition to target buffer + grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel()) + buffers.append(grad_buffer) + if self.micro_step_id == 0: # don't accumulate + grad_buffer.copy_(grad_partition, non_blocking=True) + # ensure grad buffer is a CUDA buffer to speed up the next few + # operations and so it can be used asynchronously + grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) + elif get_accelerator().on_accelerator(grad_buffer): + grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(grad_buffer.shape)) + else: + # if dst is CPU, copy first to src device, do the addition + # there, then move back to dst. adding directly to cpu is very slow + cuda_grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) + cuda_grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(cuda_grad_buffer.shape)) + grad_buffer.copy_(cuda_grad_buffer, non_blocking=True) + # ensure grad buffer is a CUDA buffer to speed up the next few + # operations and so it can be used asynchronously + grad_buffer = cuda_grad_buffer + + # offload the gradient partition if applicable + if self.offload_optimizer: + i, dest_offset, _ = self.grad_position[self.get_param_id(param)] + offload_fp32_gradients = {} + offload_fp32_offsets = {} + + if self.is_gradient_accumulation_boundary: + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) + + if self._swappable_optimizer_subgroup(i): + if not i in offload_fp32_gradients.keys(): + offload_fp32_gradients[i] = [] + offload_fp32_offsets[i] = [] + + offload_fp32_gradients[i].append(grad_buffer.float()) + offload_fp32_offsets[i].append(dest_offset) + else: + fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( + 0, dest_offset, grad_buffer.numel()) + fp32_grad_tensor.copy_(grad_buffer) + + # free the gradient + if not get_accelerator().is_synchronized_device(): + param.grad.record_stream(get_accelerator().current_stream()) + param.grad = None + + if self.offload_optimizer and self.swap_optimizer: + for i in offload_fp32_gradients.keys(): + self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i], + gradient_offsets=offload_fp32_offsets[i], + gradient_tensors=offload_fp32_gradients[i]) + see_memory_usage(f"after partition_grads", force=ENBALE_MEM_DEBUG) + return buffers diff --git a/deepspeed/runtime/zero/lins_utils.py b/deepspeed/runtime/zero/lins_utils.py new file mode 100644 index 000000000000..91dfe993ef85 --- /dev/null +++ b/deepspeed/runtime/zero/lins_utils.py @@ -0,0 +1,286 @@ +import os +import gc +import psutil + +# if global_zero35_manager is None: +global_zero35_manager = None + +from deepspeed.accelerator import get_accelerator +from deepspeed import comm as dist +from deepspeed.utils import logger +from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus +from deepspeed.utils import instrument_w_nvtx, logger +from typing import Callable, Iterable +from torch.nn import Parameter +from contextlib import contextmanager + +import torch + +init_count = 0 + +FORCE = False + +torch_memory_reserved = get_accelerator().memory_reserved +torch_max_memory_reserved = get_accelerator().max_memory_reserved + +# global param process group +global_lins_utils = None + + +class LinSProcessGroup: + + def __init__(self, dp_process_group, lins_param_partition_num: int, lins_grad_partition_num: int, + lins_os_partition_num: int) -> None: + + self.world_size = dist.get_world_size(dp_process_group) + self.data_parallel_size = self.world_size + self.model_parallel_size = 1 # not support model parallel + + self.lins_param_partition_num = lins_param_partition_num + self.lins_grad_partition_num = lins_grad_partition_num + if lins_os_partition_num == -1: + self.lins_os_partition_num = self.data_parallel_size + + assert lins_param_partition_num == lins_grad_partition_num and lins_grad_partition_num == 8, \ + "Only test lins_param_partition_num==lins_grad_partition_num == 8" + assert self.lins_os_partition_num == self.data_parallel_size, \ + f"Only test lins_os_partition_num:{self.lins_os_partition_num} == self.data_parallel_size:{self.data_parallel_size}" + + self.zero35_hierarchical_group = self.init_zero35_hierarchical_process_group() + self._grad_process_group = self.init_lins_process_group( + self.lins_grad_partition_num, self.data_parallel_size // self.lins_grad_partition_num) + self._param_process_group = self.init_lins_process_group( + self.lins_param_partition_num, self.data_parallel_size // self.lins_param_partition_num) + + self._param_rank = dist.get_rank(group=self._param_process_group) + self._param_world_size = dist.get_world_size(group=self._param_process_group) + + self._grad_rank = dist.get_rank(group=self._grad_process_group) + self._grad_world_size = dist.get_world_size(group=self._grad_process_group) + + # if dist.get_rank() < 8: + print(f"Rank:{dist.get_rank()} \ +zero35_hierarchical_group: {dist.get_all_ranks_from_group(self.zero35_hierarchical_group)}, \ +zero35_group:{dist.get_all_ranks_from_group(self._grad_process_group)} \ +self._param_rank: {self._param_rank}, \ +self._param_world_size: {self._param_world_size}, \ +self._grad_rank {self._grad_rank}, \ +self._grad_world_size: {self._grad_world_size}", + flush=True) + + def init_lins_process_group(self, parallel_size, parallel_num): + my_zero35_group = None + + for i in range(self.model_parallel_size): + for j in range(parallel_num): + ranks = [i + (j * parallel_size + k) * self.model_parallel_size for k in range(parallel_size)] + group = dist.new_group(ranks) + + if dist.get_rank() in ranks: + my_zero35_group = group + + return my_zero35_group + + def init_zero35_hierarchical_process_group(self): + my_hierarchical_zero35_group = None + gpus_per_node = 8 + nnodes = self.world_size // 8 + + for i in range(gpus_per_node): + ranks = [i + 8 * j for j in range(nnodes)] + group = dist.new_group(ranks) + if dist.get_rank() in ranks: + my_hierarchical_zero35_group = group + + return my_hierarchical_zero35_group + + +global_parition_type = {"type": "os"} + + +@contextmanager +def set_lins_parition_type(partition_type): + global global_parition_type + old_parition_type = global_parition_type["type"] + try: + global_parition_type["type"] = partition_type + yield + finally: + global_parition_type["type"] = old_parition_type + + +def zero35_judge_gahter_boundary(mico_step, forward): + """判断是不是需要做dp范围的 all-gahter + + Args: + mico_step (_type_): _description_ + forward (_type_): _description_ + + Returns: + _type_: _description_ + """ + return mico_step == 0 and forward + + +# avoid circular reference +def see_memory_usage(message, force=False): + if not force: + return + if dist.is_initialized() and not dist.get_rank() == 0: + return + + # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports + gc.collect() + + # Print message except when distributed but not rank 0 + logger.info(message) + logger.info(f"MA {round(get_accelerator().memory_allocated() / (1024 * 1024 * 1024),2 )} GB \ + Max_MA {round(get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \ + CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \ + Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ") + + vm_stats = psutil.virtual_memory() + used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2) + logger.info(f'CPU Virtual Memory: used = {used_GB} GB, percent = {vm_stats.percent}%') + + # get the peak memory to report correct data, so reset the counter for the next call + get_accelerator().reset_peak_memory_stats() + + +def zero35_debug(msg, rank=None, force=FORCE, flush=True): + if force: + msg = f"Rank: {os.environ['SLURM_PROCID']}, " + msg + if rank is None: + if flush: + print(msg, flush=True) + else: + logger.info(msg) + elif os.environ['SLURM_PROCID'] == str(rank): + if flush: + print(msg, flush=True) + else: + logger.info(msg) + + +def zero35_g_p_reduce_scatter_coalesced(tensor_list, partition_type): + # reshape 的逻辑 + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 4, 1, 5, 2, 6, 3, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 4, 1, 5, 2, 6, 3, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 4, 1, 5, 2, 6, 3, 7] + # [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 4, 1, 5, 2, 6, 3, 7] + + see_memory_usage(f"before zero35_g_p_reduce_scatter_coalesced, partition_type:{partition_type}") + + do_reshape = partition_type == "grad" or partition_type == "param" + dtype = tensor_list[0].dtype + + # if do_reshape: + dp_comm_group = global_zero35_manager._dp_process_group + param_comm_group = global_zero35_manager._param_process_group + + if do_reshape: + scatter_comm_group = param_comm_group + + # dp_world_size = dist.get_world_size(dp_comm_group) + # param_world_size = dist.get_world_size(param_comm_group) + + # new_tensor_list = [] + # _undo_indexs_for_per_tensor = [] + # for grad in tensor_list: + # assert grad.numel() % dp_world_size == 0 + # assert grad.numel() % param_world_size == 0 + + # dp_partition_size = int(grad.numel() / dp_world_size) # 按照dp范围划分的最小part大小 + # param_partition_size = int(grad.numel() / param_world_size) + # assert param_partition_size % dp_partition_size == 0 + # param_partition_num = int(param_partition_size / dp_partition_size) # 每个节点内包含的 dp_partition_size 的数量 + # grad = grad.reshape(-1, dp_partition_size) + # indexs = [] + # for idx in range(param_world_size): + # for jdx in range(param_partition_num): + # indexs.append(idx + jdx * param_world_size) + + # # zero35_debug(f"scatter index : {indexs}") + + # indexs=torch.tensor(indexs).to(get_accelerator().device_name()) + # _, undo_indices = torch.sort(indexs, dim=0, descending=False) + # _undo_indexs_for_per_tensor.append(undo_indices) + + # reshape_grad = torch.index_select(grad, 0, indexs) + # assert reshape_grad.is_contiguous() + # new_tensor_list.append(reshape_grad.view(-1)) + # tensor_list = new_tensor_list + else: + scatter_comm_group = dp_comm_group + + # if do_reshape: + # new_tensor_list = [] + # for i, grad in enumerate(tensor_list): + # new_tensor_list.append(torch.index_select(grad, 0, _undo_indexs_for_per_tensor[i])) + # tensor_list = new_tensor_list + + see_memory_usage(f"after zero35_g_p_reduce_scatter_coalesced, partition_type:{partition_type}") + + return tensor_list, scatter_comm_group + + +def zero35_g_p_all_gather_coalesced(tensor_list, partition_type=None): + # reshape 的逻辑 + # [0, 4, 1, 5, 2, 6, 3, 7] -> [0, 1, 2, 3, 4, 5, 6, 7] + # [0, 4, 1, 5, 2, 6, 3, 7] -> [0, 1, 2, 3, 4, 5, 6, 7] + # [0, 4, 1, 5, 2, 6, 3, 7] -> [0, 1, 2, 3, 4, 5, 6, 7] + # [0, 4, 1, 5, 2, 6, 3, 7] -> [0, 1, 2, 3, 4, 5, 6, 7] + # + # [0, , 1, ] + # do_reshape = partition_type == "grad" or partition_type == "param" + dtype = tensor_list[0].dtype + + # if do_reshape: + see_memory_usage(f"before zero35_g_p_all_gather_coalesced, partition_type:{partition_type}") + + dp_comm_group = global_zero35_manager._dp_process_group + param_comm_group = global_zero35_manager._param_process_group + + all_gather_comm_group = param_comm_group + + # dp_world_size = dist.get_world_size(dp_comm_group) + # param_world_size = dist.get_world_size(param_comm_group) + + # new_tensor_list = [] + + # for t_data in tensor_list: + # param_full_tensor = t_data.data + # indexs = [] + + # partition_unit_size = t_data.numel() // dp_world_size # 按照dp范围划分的最小part大小 + + # # zero35_debug(f"param_full_tensor.numel() :{t_data.numel()}, dp_world_size:{dp_world_size},partition_unit_size:{partition_unit_size}", flush=True) + + # partition_unit_num = t_data.numel() // partition_unit_size + + # partition_unit_size_per_rank = t_data.numel() // param_world_size + # partition_unit_num_per_rank = partition_unit_size_per_rank // partition_unit_size # 2 + # partition_unit_num_per_node = partition_unit_num_per_rank * param_world_size # 2 * 4 -> 8 + + # param_full_tensor = param_full_tensor.reshape(-1, partition_unit_size) + + # for idx in range(partition_unit_num_per_rank): # 8 + # indexs.extend([idx + jdx * partition_unit_num_per_rank for jdx in range(param_world_size)]) + + # # zero35_debug(f"gather index : {indexs}") + # # indexs = [0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15] + # indexs=torch.tensor(indexs).to(get_accelerator().device_name()) + # reshape_t_data = torch.index_select(param_full_tensor, 0, indexs) + # reshape_t_data = reshape_t_data.view(t_data.data.shape) + # assert reshape_t_data.is_contiguous() + + # # param_full_tensor.ds_tensor = reshape_t_data + # t_data.data = reshape_t_data + + # new_tensor_list.append(reshape_t_data) + # tensor_list = new_tensor_list + # # else: + # # all_gather_comm_group = dp_comm_group + see_memory_usage(f"before zero35_g_p_all_gather_coalesced, partition_type:{partition_type}") + + return tensor_list, all_gather_comm_group diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index de94d757c3b7..043ab5eff936 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -1028,6 +1028,11 @@ def _convert_to_deepspeed_param(self, param): # The group that the parameter is scattered across. param.ds_process_group = self.ds_process_group + def get_param_process_group(): + return param.ds_process_group + + param.get_param_process_group = get_param_process_group + # Stores the secondary partitioned copy of the tensor param.ds_secondary_tensor = None diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py index 9bcf5a91bc95..b83d33f80665 100644 --- a/deepspeed/runtime/zero/partitioned_param_coordinator.py +++ b/deepspeed/runtime/zero/partitioned_param_coordinator.py @@ -18,6 +18,9 @@ from deepspeed.accelerator import get_accelerator import logging +from deepspeed.runtime.zero.lins_utils import zero35_debug, set_lins_parition_type + + ENABLE_PROFILER = False @@ -125,6 +128,12 @@ def __init__( self.__max_ongoing_fetch_events: int = 2 self.__profiler = PartitionedParameterProfiler(timers if ENABLE_PROFILER else None) + # （TODO):wgt + self.is_gradient_accumulation_boundary = True + self.micro_step_id = 0 + self.gradient_accumulation_steps = None + self.all_gahter_count = 0 + """Tracing and Tracking TODO. consider performing trace before initializing PartitionedParameterCoordinator and passing trace results into constructor. This way all the code in here can @@ -269,8 +278,10 @@ def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None: })) params_to_fetch = frozenset(iter_params(current_submodule)) - fetch_numel = sum( - [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + + with set_lins_parition_type(partition_type = "param"): + 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.id, @@ -293,7 +304,8 @@ def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None: 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 set_lins_parition_type(partition_type = "param"): + 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() diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index ed44af08a686..eb24b7ba5aa9 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -26,6 +26,7 @@ from deepspeed.runtime.swap_tensor.pipelined_optimizer_swapper import PipelinedOptimizerSwapper from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FP32_FLAT_GROUPS, PARTITION_COUNT, ZERO_STAGE from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.lins_utils import zero35_debug, set_lins_parition_type # Toggle this to true to enable correctness test # with gradient partitioning and without @@ -188,6 +189,7 @@ def __init__( zero_quantized_weights=zero_quantized_weights, zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights) + self._get_param_coordinator(training=True).gradient_accumulation_steps = gradient_accumulation_steps self.persistent_parameters = self.parameter_offload.persistent_parameters self._configure_offloading(offload_optimizer_config, offload_param_config) @@ -339,7 +341,7 @@ def __init__( print_rank_0(f'Largest partitioned param numel = {largest_partitioned_param_numel}', force=False) - self._setup_for_real_optimizer() + all_params = self._setup_for_real_optimizer() self.grad_position = {} self.set_grad_positions() @@ -375,6 +377,10 @@ def __init__( if dist.get_rank(group=self.dp_process_group) == 0: see_memory_usage(f"After initializing ZeRO optimizer", force=True) + if all_params is not None: + self.cal_stage3_mem_usage(all_params) + + def destroy(self): self.parameter_offload.destroy() del self.__ipg_bucket_flat_buffer @@ -631,13 +637,7 @@ def _create_param_groups_fp16_flat_cpu_memory(self): self.param_groups_fp16_flat_cpu_memory.append(torch.empty(1, dtype=self.dtype)) - def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): - dist.barrier() - - param_groups: List[List[Parameter]] = tuple( - self._create_fp16_sub_groups(param_group["params"]) for param_group in fp16_param_groups) - - # bookkeeping related to param groups + def bookkeeping_param_group(self, param_groups): for param_group_idx, param_group in enumerate(param_groups): for sub_group in param_group: sub_group_idx = len(self.fp16_groups) @@ -657,6 +657,15 @@ def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): self.dp_process_group) == dist.get_world_size(self.dp_process_group) - 1 self.groups_padding.append([p.padding_size() if rank_requires_padding else 0 for p in sub_group]) + def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): + dist.barrier() + + param_groups: List[List[Parameter]] = tuple( + self._create_fp16_sub_groups(param_group["params"]) for param_group in fp16_param_groups) + + # bookkeeping related to param groups + self.bookkeeping_param_group(param_groups) + # move parameters to flattened buffer if not self.offload_param: # partitioned params remain in GPU during training # move parameter partitions into a single contiguous flat buffer @@ -670,8 +679,16 @@ def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): # contiguous flat buffer for all parameters that we created earlier offset = 0 for sub_group in self.fp16_groups: - sub_group_numel = sum(param.partition_numel() for param in sub_group) - self.fp16_partitioned_groups_flat.append(device_buffer.narrow(0, offset, sub_group_numel)) + with set_lins_parition_type(partition_type="os"): + sub_group_numel = sum(param.partition_numel() for param in sub_group) + # import pdb + # pdb.set_trace() + try: + self.fp16_partitioned_groups_flat.append(device_buffer.narrow(0, offset, sub_group_numel)) + except Exception as e: + print("catch: {e}", flush=True) + import pdb + pdb.set_trace() offset += sub_group_numel else: # partitioned params offloaded to CPU when not in use # create a flat CPU memory allocation for each param group @@ -854,7 +871,7 @@ def _create_fp32_partitions(self): def _create_fp16_sub_groups(self, params_group): - params_group_numel = sum([param.partition_numel() for param in params_group]) + params_group_numel = sum([param.partition_numel(partition_type="os") for param in params_group]) sub_group_size = self.sub_group_size if sub_group_size is None or sub_group_size >= params_group_numel: @@ -866,7 +883,7 @@ def _create_fp16_sub_groups(self, params_group): for param in params_group: sub_group.append(param) - local_sub_group_size += param.partition_numel() + local_sub_group_size += param.partition_numel(partition_type="os") if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): @@ -1241,7 +1258,8 @@ def set_grad_positions(self): current_offset = 0 for param in group: param_id = self.get_param_id(param) - num_elements = param.partition_numel() + with set_lins_parition_type(partition_type="param"): + num_elements = param.partition_numel() self.grad_position[param_id] = [int(i), int(current_offset), int(num_elements)] #print(f"param id {param_id} i:{i}, ds_tensor {num_elements} numel {param.numel()}") @@ -1574,6 +1592,7 @@ def zero_grad(self, set_to_none=True): Zero FP16 parameter grads. """ self.micro_step_id = 0 + self._get_param_coordinator(training=True).micro_step_id = 0 # FP32 grad should never exist. # For speed, set model fp16 grad to None by default @@ -1703,6 +1722,7 @@ def reset_cpu_buffers(self): def _pre_step(self): self.micro_step_id = 0 + self._get_param_coordinator(training=True).micro_step_id = 0 print_rank_0(f"Inside Step function") see_memory_usage(f"In step before checking overflow", force=False) diff --git a/test_lins.py b/test_lins.py new file mode 100644 index 000000000000..83243fd0feb6 --- /dev/null +++ b/test_lins.py @@ -0,0 +1,201 @@ +import torch +import os +import torch.nn.functional as F + +from torch.distributed import GroupMember +import torch.nn as nn +import deepspeed +from deepspeed.runtime.utils import see_memory_usage +from torch.optim import Adam +from deepspeed.comm import init_distributed +import torch.distributed as dist +import numpy as np +from torch.utils.data import Dataset + +json_path = "stage3.json" + + +def print_rank_0(message): + """If distributed is initialized, print only on rank 0.""" + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() == 0: + print(message, flush=True) + else: + print(message, flush=True) + + +def get_master_node(): + import subprocess + + if os.getenv("SLURM_JOB_ID") is None: + raise RuntimeError("get_master_node can only used in Slurm launch!") + result = subprocess.check_output('scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1', shell=True) + result = result.decode("utf8").strip() + return result + + +class ToyModel2(torch.nn.Module): + + def __init__(self) -> None: + super(ToyModel2, self).__init__() + self.step_count = 0 + self.dense1 = torch.nn.Parameter( + data=torch.tensor([-i for i in range(64)], dtype=torch.bfloat16, requires_grad=True)) + + def forward(self, x, labels=None): + x = x.to(f"cuda:{int(os.environ['SLURM_PROCID']) % 8}") + if os.environ['SLURM_PROCID'] == '0': + print(f"Rank: {os.environ['SLURM_PROCID']}:\ +do ToyModel2-forward!,self.step_count: {self.step_count}, \ +self.dense:{self.dense1}", + flush=True) + + self.step_count += 1 + y = self.dense1.mul(x) + # x = self.dense2.mul(y) + return y + + def set_input_tensor(self, input_tensor): + self.input_tensor = input_tensor + + +class ToyModel(torch.nn.Module): + + def __init__(self) -> None: + # 64 / 16 = 4 + # 64 / 8 = 8 + super(ToyModel, self).__init__() + self.step_count = 0 + # self.dense = torch.nn.Linear(64, 1, bias=False, dtype=torch.bfloat16) + + # 我们让 reduce_bucket_size 正好等于 64，这样在做完 ToyModel2 的 bwd 后就可以直接进行 bucket 的 allreduce + self.dense1 = torch.nn.Parameter( + data=torch.tensor([i for i in range(64)], dtype=torch.bfloat16, requires_grad=True)) + self.sub_module = ToyModel2() + self.dense2 = torch.nn.Parameter( + data=torch.tensor([-i for i in range(64)], dtype=torch.bfloat16, requires_grad=True)) + + def forward(self, x, labels=None): + x = x.to(f"cuda:{int(os.environ['SLURM_PROCID']) % 8}") + if os.environ['SLURM_PROCID'] == '0': + print(f"Rank: {os.environ['SLURM_PROCID']}:\ +do forward!,self.step_count: {self.step_count}, \ +self.dense:{self.dense1}", + flush=True) + + self.step_count += 1 + y = self.sub_module(x) + y = self.dense1 + y + x = self.dense2.mul(y) + return x + + def set_input_tensor(self, input_tensor): + self.input_tensor = input_tensor + + +def model_provider(pre_process=True, post_process=True): + """Build the model.""" + + print_rank_0('building llama model ...') + see_memory_usage(f"Before Building Model", force=True) + + + init_distributed(dist_backend="nccl", \ + auto_mpi_discovery=False, \ + init_method=f"tcp://[{get_master_node()}]:12349", \ + rank=int(os.environ['SLURM_PROCID']), \ + world_size=16) + + local_group1 = dist.new_group([i for i in range(8)]) + local_group2 = dist.new_group([i for i in range(8, 16, 1)]) + + os.environ['LOCAL_RANK'] = f"{int(os.environ['SLURM_PROCID']) % 8}" + + def get_param_group(): + if int(os.environ['SLURM_PROCID']) < 8: + return local_group1 + else: + return local_group2 + + from deepspeed.runtime.zero.lins import LinS_Init + with LinS_Init(data_parallel_group=GroupMember.WORLD, + remote_device=None, + config_dict_or_path=json_path, + enabled=True, + mpu=None): + model = ToyModel() + see_memory_usage(f"After Building Model", force=True) + return model + + +def _get_params_for_weight_decay_optimization_one_subgourp(modules): + weight_decay_params = {'params': [], 'name': 'weight_decay_params'} + for module in modules: + for module_ in module.modules(): + for num, param in list(module_._parameters.items()): + weight_decay_params['params'].extend([param]) + + return weight_decay_params, + + +class RandomDataset(torch.utils.data.IterableDataset): + + def __init__(self, num_samples=1000000, seq_len=1024) -> None: + super().__init__() + self.len = 1024 + self.data = [i for i in range(64)] + + def __getitem__(self, index): + return np.array(self.data, dtype=int) + + def get_dataset_name(self): + return "test_toy" + + def __len__(self): + return self.len + + +class SimpleBatchSampler: + + def __init__(self, total_samples) -> None: + self.total_samples = total_samples + + def __len__(self): + return self.total_samples + + def __iter__(self): + yield list(range(0, 64)) + + +# srun -p llm_t --quotatype=spot -n16 -N2 --ntasks-per-node=8 --gpus-per-task=1 python test_lins.py +if __name__ == "__main__": + model = model_provider() + model = [model] + param_groups = _get_params_for_weight_decay_optimization_one_subgourp(model) + optimizer = Adam(param_groups, lr=1e-4, weight_decay=1e-2, betas=(1e-2, 1e-2), eps=1e-2) + + lr_scheduler = None + train_ds = RandomDataset() + # batch_sampler = SimpleBatchSampler(total_samples=1024) + # train_dataloader = torch.utils.data.DataLoader(train_ds, + # batch_sampler=None, + # num_workers=0, + # pin_memory=False) + ( + model, + optimizer, + deepspeed_dataloader, + lr_scheduler, + ) = deepspeed.initialize( + model=model[0], + optimizer=optimizer, + args=None, + lr_scheduler=lr_scheduler, + training_data=train_ds, + mpu=None, + config=json_path, + ) + + data_iterator = iter(deepspeed_dataloader) + for i in range(10): + loss = model[0].train_batch(data_iter=data_iterator)