From ea372da285fc758a7330f9f55074b9d611749ba3 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Sat, 26 Oct 2024 17:59:04 +0000 Subject: [PATCH] Remove unused adaptor files --- opensora/adaptor/bf16_optimizer.py | 458 ---- opensora/adaptor/engine.py | 3568 ---------------------------- opensora/adaptor/stage_1_and_2.py | 2522 -------------------- opensora/adaptor/utils.py | 1042 -------- opensora/adaptor/zp_manager.py | 31 - opensora/sample/sample_t2v_sp.py | 1 - t2v_debug_multi.sh | 1 - 7 files changed, 7623 deletions(-) delete mode 100644 opensora/adaptor/bf16_optimizer.py delete mode 100644 opensora/adaptor/engine.py delete mode 100644 opensora/adaptor/stage_1_and_2.py delete mode 100644 opensora/adaptor/utils.py delete mode 100644 opensora/adaptor/zp_manager.py diff --git a/opensora/adaptor/bf16_optimizer.py b/opensora/adaptor/bf16_optimizer.py deleted file mode 100644 index e0401c0b3..000000000 --- a/opensora/adaptor/bf16_optimizer.py +++ /dev/null @@ -1,458 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from collections import OrderedDict -import torch -import sys -import os -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from deepspeed import comm as dist -from deepspeed.runtime.constants import PIPE_REPLICATED -from deepspeed.runtime import ZeROOptimizer -from packaging import version as pkg_version - -from deepspeed.git_version_info import version -from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim, - align_dense_tensors, all_gather_dp_groups, bwc_tensor_model_parallel_rank, - is_model_parallel_parameter, see_memory_usage, graph_process) - -from deepspeed.utils import link_hp_params, fragment_address -from deepspeed.checkpoint import enable_universal_checkpoint -from deepspeed.checkpoint.constants import (DS_VERSION, PARTITION_COUNT, BASE_OPTIMIZER_STATE, - SINGLE_PARTITION_OF_FP32_GROUPS, CLIP_GRAD, GROUP_PADDINGS, - PARAM_SLICE_MAPPINGS) - -setattr(sys.modules[__name__], 'fragment_address', fragment_address) - - -def contigous_flatten(tensors): - return _flatten_dense_tensors([tensor.contiguous() for tensor in tensors]) - - -class BF16_Optimizer(ZeROOptimizer): - - def __init__(self, - init_optimizer, - param_names, - mpu=None, - clip_grad=0.0, - norm_type=2, - allgather_bucket_size=5000000000, - dp_process_group=None, - timers=None, - grad_acc_dtype=None, - graph_harvesting=False): - # super().__init__() - # base_class = ZeROOptimizer.__bases__[0] - # # 直接调用基类的 __init__ 方法 - # base_class.__init__() - see_memory_usage('begin bf16_optimizer', force=True) - self.timers = timers - self.optimizer = init_optimizer - self.param_names = param_names - self.using_real_optimizer = not isinstance(self.optimizer, DummyOptim) - - assert grad_acc_dtype in [torch.float32, torch.bfloat16 - ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}" - self.grad_acc_dtype = grad_acc_dtype - - self.clip_grad = clip_grad - self.norm_type = norm_type - self.mpu = mpu - self.allgather_bucket_size = int(allgather_bucket_size) - self.dp_process_group = dp_process_group - self.dp_rank = dist.get_rank(group=self.dp_process_group) - self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] - - # Use torch (un)flatten ops - self.flatten = contigous_flatten - self.unflatten = _unflatten_dense_tensors - - #align nccl all-gather send buffers to 4-bye boundary - self.nccl_start_alignment_factor = 16 - - # Build BF16/FP32 groups - self.bf16_groups = [] - self.bf16_groups_flat = [] - self.bf16_partitioned_groups = [] - - self.fp32_groups_flat_partition = [] - - # Maintain different fp32 gradients views for convenience - self.fp32_groups_gradients = [] - self.fp32_groups_gradient_dict = {} - self.fp32_groups_gradients_flat = [] - self.fp32_groups_actual_gradients_flat = [] - self.fp32_groups_gradient_flat_partition = [] - self.fp32_groups_has_gradients = [] - - self.group_paddings = [] - self.graph_harvesting = graph_harvesting - if self.using_real_optimizer: - self._setup_for_real_optimizer() - - see_memory_usage('end bf16_optimizer', force=True) - - def _setup_for_real_optimizer(self): - dp_world_size = dist.get_world_size(group=self.dp_process_group) - self.partition_count = [dp_world_size for i in range(len(self.optimizer.param_groups))] - - for i, param_group in enumerate(self.optimizer.param_groups): - see_memory_usage(f'before initializing group {i}', force=True) - - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - - # grab the original list - trainable_parameters = [param for param in param_group['params'] if param.requires_grad] - self.bf16_groups.append(trainable_parameters) - - # create flat bf16 params - self.bf16_groups_flat.append( - self._flatten_dense_tensors_aligned(self.bf16_groups[i], - self.nccl_start_alignment_factor * dp_world_size)) - - # Make bf16 params point to flat tensor storage - self._update_storage_to_flattened_tensor(tensor_list=self.bf16_groups[i], - flat_tensor=self.bf16_groups_flat[i]) - - # divide flat weights into equal sized partitions - partition_size = self.bf16_groups_flat[i].numel() // dp_world_size - bf16_dp_partitions = [ - self.bf16_groups_flat[i].narrow(0, dp_index * partition_size, partition_size) - for dp_index in range(dp_world_size) - ] - self.bf16_partitioned_groups.append(bf16_dp_partitions) - - # create fp32 params partition - self.fp32_groups_flat_partition.append(bf16_dp_partitions[partition_id].clone().float().detach()) - self.fp32_groups_flat_partition[i].requires_grad = True - - num_elem_list = [t.numel() for t in self.bf16_groups[i]] - - # create fp32 gradients - self.fp32_groups_gradients_flat.append( - torch.zeros_like(self.bf16_groups_flat[i], dtype=self.grad_acc_dtype)) - - # track individual fp32 gradients for entire model - fp32_gradients = self._split_flat_tensor(flat_tensor=self.fp32_groups_gradients_flat[i], - num_elem_list=num_elem_list) - self.fp32_groups_gradients.append(fp32_gradients) - self.fp32_groups_gradient_dict[i] = fp32_gradients - - # flat tensor corresponding to actual fp32 gradients (i.e., minus alignment padding) - length_without_padding = sum(num_elem_list) - self.fp32_groups_actual_gradients_flat.append( - torch.narrow(self.fp32_groups_gradients_flat[i], 0, 0, length_without_padding)) - - # flat tensor corresponding to gradient partition - self.fp32_groups_gradient_flat_partition.append( - torch.narrow(self.fp32_groups_gradients_flat[i], 0, partition_id * partition_size, partition_size)) - - # track fp32 gradient updates - self.fp32_groups_has_gradients.append([False] * len(self.bf16_groups[i])) - - # Record padding required for alignment - if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: - padding = self.bf16_groups_flat[i].numel() - length_without_padding - else: - padding = 0 - - self.group_paddings.append(padding) - - # update optimizer param groups to reference fp32 params partition - param_group['params'] = [self.fp32_groups_flat_partition[i]] - - see_memory_usage(f'after initializing group {i}', force=True) - - see_memory_usage('before initialize_optimizer', force=True) - self.initialize_optimizer_states() - see_memory_usage('end initialize_optimizer', force=True) - - # Need optimizer states initialized before linking lp to optimizer state - self._link_all_hp_params() - self._enable_universal_checkpoint() - self._param_slice_mappings = self._create_param_mapping() - - def _enable_universal_checkpoint(self): - for lp_param_group in self.bf16_groups: - enable_universal_checkpoint(param_list=lp_param_group) - - def _create_param_mapping(self): - param_mapping = [] - for i, _ in enumerate(self.optimizer.param_groups): - param_mapping_per_group = OrderedDict() - for lp in self.bf16_groups[i]: - if lp._hp_mapping is not None: - lp_name = self.param_names[lp] - param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() - param_mapping.append(param_mapping_per_group) - - return param_mapping - - def _link_all_hp_params(self): - dp_world_size = dist.get_world_size(group=self.dp_process_group) - for i, _ in enumerate(self.optimizer.param_groups): - # Link bf16 and fp32 params in partition - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - partition_size = self.bf16_groups_flat[i].numel() // dp_world_size - flat_hp_partition = self.fp32_groups_flat_partition[i] - link_hp_params(lp_param_list=self.bf16_groups[i], - flat_hp_partition=flat_hp_partition, - gradient_dict=self.fp32_groups_gradient_dict, - offload_gradient_dict=None, - use_offload=False, - param_group_index=i, - partition_start=partition_id * partition_size, - partition_size=partition_size, - partition_optimizer_state=self.optimizer.state[flat_hp_partition], - dp_group=self.real_dp_process_group[i]) - - def initialize_optimizer_states(self): - """Take an optimizer step with zero-valued gradients to allocate internal - optimizer state. - - This helps prevent memory fragmentation by allocating optimizer state at the - beginning of training instead of after activations have been allocated. - """ - for param_partition, grad_partition in zip(self.fp32_groups_flat_partition, - self.fp32_groups_gradient_flat_partition): - # In case of grad acc dtype different than FP32, need to cast to high precision. - param_partition.grad = grad_partition.to( - param_partition.dtype) if grad_partition.dtype != param_partition.dtype else grad_partition - - self.optimizer.step() - - if self.grad_acc_dtype is not torch.float32: - for param_partition in self.fp32_groups_flat_partition: - param_partition.grad = None - - self.clear_hp_grads() - - def _split_flat_tensor(self, flat_tensor, num_elem_list): - assert sum(num_elem_list) <= flat_tensor.numel() - tensor_list = [] - offset = 0 - for num_elem in num_elem_list: - dense_tensor = torch.narrow(flat_tensor, 0, offset, num_elem) - tensor_list.append(dense_tensor) - offset += num_elem - - return tensor_list - - def _update_storage_to_flattened_tensor(self, tensor_list, flat_tensor): - updated_params = self.unflatten(flat_tensor, tensor_list) - for p, q in zip(tensor_list, updated_params): - p.data = q.data - - def _flatten_dense_tensors_aligned(self, tensor_list, alignment): - return self.flatten(align_dense_tensors(tensor_list, alignment)) - - @torch.no_grad() - def step(self, closure=None): - if closure is not None: - raise NotImplementedError(f'{self.__class__} does not support closure.') - - all_groups_norm = get_global_norm_of_tensors(input_tensors=self.get_grads_for_norm(), - mpu=self.mpu, - norm_type=self.norm_type, - use_graph=self.graph_harvesting) - self._global_grad_norm = all_groups_norm - - assert all_groups_norm > 0. - if self.clip_grad > 0.: - clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(for_clipping=True), - max_norm=self.clip_grad, - global_norm=all_groups_norm, - mpu=self.mpu, - use_graph=self.graph_harvesting) - - self.optimizer.step() - - self.update_lp_params() - - self.clear_hp_grads() - - def backward(self, loss, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs): - """Perform a backward pass and copy the low-precision gradients to the - high-precision copy. - - We copy/accumulate to the high-precision grads now to prevent accumulating in the - bf16 grads after successive backward() calls (i.e., grad accumulation steps > 1) - - The low-precision grads are deallocated during this procedure. - """ - self.clear_lp_grads() - loss.backward(**bwd_kwargs) - - if update_hp_grads: - self.update_hp_grads(clear_lp_grads=clear_lp_grads) - - @torch.no_grad() - def update_hp_grads(self, clear_lp_grads=False): - - def _update_hp_grads_func(clear_lp_grads=False): - for i, group in enumerate(self.bf16_groups): - for j, lp in enumerate(group): - if lp.grad is None: - continue - hp_grad = self.fp32_groups_gradients[i][j] - assert hp_grad is not None, \ - f'high precision param has no gradient, lp param_id = {id(lp)} group_info = [{i}][{j}]' - hp_grad.data.add_(lp.grad.data.to(hp_grad.dtype).view(hp_grad.shape)) - lp._hp_grad = hp_grad - self.fp32_groups_has_gradients[i][j] = True - # clear gradients - if clear_lp_grads: - lp.grad._zero() - - if self.graph_harvesting: - graph_process(False, _update_hp_grads_func, clear_lp_grads) - else: - _update_hp_grads_func(clear_lp_grads) - #cpu op - for i, group in enumerate(self.bf16_groups): - for j, lp in enumerate(group): - if lp.grad is None: - continue - self.fp32_groups_has_gradients[i][j] = True - - @torch.no_grad() - def get_grads_for_reduction(self): - return self.fp32_groups_gradients_flat - - @torch.no_grad() - def get_grads_for_norm(self, for_clipping=False): - grads = [] - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) - for i, group in enumerate(self.bf16_groups): - for j, lp in enumerate(group): - if not for_clipping: - if hasattr(lp, PIPE_REPLICATED) and lp.ds_pipe_replicated: - continue - - if not (tensor_mp_rank == 0 or is_model_parallel_parameter(lp)): - continue - - if not self.fp32_groups_has_gradients[i][j]: - continue - - grads.append(self.fp32_groups_gradients[i][j]) - - return grads - - @torch.no_grad() - def update_lp_params(self): - for i, (bf16_partitions, - fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - bf16_partitions[partition_id].data.copy_(fp32_partition.data) - # print_rank_0(f'update_lp_params {i=} {partition_id=}', force=True) - # if i == 0: - # print_rank_0(f'{fp32_partition[:10]=}', force=True) - - all_gather_dp_groups(groups_flat=self.bf16_groups_flat, - partitioned_param_groups=self.bf16_partitioned_groups, - dp_process_group=self.real_dp_process_group, - start_alignment_factor=self.nccl_start_alignment_factor, - allgather_bucket_size=self.allgather_bucket_size) - - def clear_hp_grads(self): - for flat_gradients in self.fp32_groups_gradients_flat: - flat_gradients.zero_() - - for i, group in enumerate(self.fp32_groups_gradients): - self.fp32_groups_has_gradients[i] = [False] * len(group) - - def clear_lp_grads(self): - for group in self.bf16_groups: - for param in group: - if param.grad is not None: - # Using zero_() fixed memory address for graph replay - param.grad.zero_() - - def state_dict(self): - state_dict = {} - state_dict[CLIP_GRAD] = self.clip_grad - state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() - state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = self.fp32_groups_flat_partition - state_dict[GROUP_PADDINGS] = self.group_paddings - state_dict[PARTITION_COUNT] = self.partition_count - state_dict[DS_VERSION] = version - state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings - - return state_dict - - # Restore base optimizer fp32 weights bfloat16 weights - def _restore_from_bit16_weights(self): - for i, group in enumerate(self.bf16_groups): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - for bf16_partitions, fp32_partition in zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition): - fp32_partition.data.copy_(bf16_partitions[partition_id].data) - - def refresh_fp32_params(self): - self._restore_from_bit16_weights() - - def load_state_dict(self, - state_dict_list, - checkpoint_folder, - load_optimizer_states=True, - load_from_fp32_weights=False, - load_serial=None): - if checkpoint_folder: - self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) - else: - self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) - - def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): - - dp_rank = dist.get_rank(group=self.dp_process_group) - current_rank_sd = state_dict_list[dp_rank] - - ckpt_version = current_rank_sd.get(DS_VERSION, False) - assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" - ckpt_version = pkg_version.parse(ckpt_version) - - self.clip_grad = current_rank_sd.get(CLIP_GRAD, self.clip_grad) - - if load_optimizer_states: - self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) - - if load_from_fp32_weights: - for current, saved in zip(self.fp32_groups_flat_partition, - current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): - src_tensor = _get_padded_tensor(saved, current.numel()) - current.data.copy_(src_tensor.data) - - if load_optimizer_states: - self._link_all_hp_params() - - def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): - self._load_hp_checkpoint_state(checkpoint_folder) - - @property - def param_groups(self): - """Forward the wrapped optimizer's parameters.""" - return self.optimizer.param_groups - - def _load_hp_checkpoint_state(self, checkpoint_dir): - checkpoint_dir = os.path.join(checkpoint_dir, "zero") - tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) - tp_world_size = self.mpu.get_slice_parallel_world_size() - - for i, _ in enumerate(self.optimizer.param_groups): - for lp in self.bf16_groups[i]: - if lp._hp_mapping is not None: - #print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}") - lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank, - tp_world_size) - - -def _get_padded_tensor(src_tensor, size): - if src_tensor.numel() >= size: - return src_tensor - padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) - slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) - slice_tensor.data.copy_(src_tensor.data) - return padded_tensor diff --git a/opensora/adaptor/engine.py b/opensora/adaptor/engine.py deleted file mode 100644 index 5208a9e41..000000000 --- a/opensora/adaptor/engine.py +++ /dev/null @@ -1,3568 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import re -import stat -import torch -import hashlib -from collections import defaultdict, OrderedDict, deque -from shutil import copyfile -import gc - -from torch.nn.modules import Module -from torch.nn.parameter import Parameter -from torch.optim import Optimizer -from torch.optim.lr_scheduler import _LRScheduler -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors - -from typing import Callable, Dict, Union, Iterable - -import deepspeed - -from deepspeed import comm as dist -from deepspeed.runtime.utils import see_memory_usage, DummyOptim -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum -from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer -from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus -from deepspeed.runtime.zero.utils import is_zero_supported_optimizer, ZeRORuntimeException -from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload -from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION - -from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer -from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer -from deepspeed.runtime.bf16_optimizer import BF16_Optimizer - -from deepspeed.runtime.config import DEEPSPEED_OPTIMIZERS, \ - ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, \ - TORCH_ADAM_PARAM, ADAM_W_MODE, ADAM_W_MODE_DEFAULT, ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, \ - MUSGD_OPTIMIZER, LION_OPTIMIZER - -from deepspeed.runtime.dataloader import DeepSpeedDataLoader -from deepspeed.runtime.constants import \ - ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \ - PLD_THETA, PLD_GAMMA, BFLOAT16, FP16, AMP, GRADIENT_ACCUMULATION_STEPS, \ - DATA_PARALLEL_GROUP, GLOBAL_RANK -from deepspeed.runtime.zero.config import ZeroStageEnum -from deepspeed.compression import compression_scheduler -from deepspeed.compression.constants import \ - WEIGHT_QUANTIZE_IN_FORWARD_ENABLED, \ - WEIGHT_QUANTIZATION, SHARED_PARAMETERS, \ - WEIGHT_QUANTIZE_ENABLED, \ - WEIGHT_QUANTIZE_GROUPS, \ - WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE, \ - WEIGHT_QUANTIZE_CHANGE_RATIO, \ - WEIGHT_QUANTIZE_TYPE, \ - WEIGHT_QUANTIZE_ROUNDING, \ - WEIGHT_QUANTIZE_VERBOSE, \ - WEIGHT_QUANTIZE_KERNEL -from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FROZEN_PARAM_FRAGMENTS -from deepspeed.runtime.sparse_tensor import SparseTensor - -from deepspeed.runtime import lr_schedules -from deepspeed.utils import groups -from deepspeed.utils import logger, log_dist, instrument_w_nvtx -from deepspeed.utils.timer import NoopTimer, ThroughputTimer, SynchronizedWallClockTimer, \ - FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, \ - STEP_MICRO_TIMER, \ - FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ - STEP_GLOBAL_TIMER -from deepspeed.utils.debug import debug_extract_module_and_param_names -from deepspeed.monitor.monitor import MonitorMaster -from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop -from deepspeed.runtime.utils import clip_grad_norm_ -from deepspeed.runtime.eigenvalue import Eigenvalue -from deepspeed.runtime.data_pipeline.constants import DATA_SAMPLING, \ - DATA_ROUTING, DATA_SAMPLING_ENABLED, CURRICULUM_LEARNING, \ - CURRICULUM_LEARNING_ENABLED, DATA_SAMPLING_NUM_WORKERS, RANDOM_LTD, \ - RANDOM_LTD_ENABLED, RANDOM_LTD_LAYER_ID, RANDOM_LTD_LAYER_NUM, \ - RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE, RANDOM_LTD_LAYER_TOKEN_LR_ENABLED, \ - RANDOM_LTD_GLOBAL_BATCH_SIZE, RANDOM_LTD_MICRO_BATCH_SIZE, DATA_EFFICIENCY -from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler -from deepspeed.runtime.data_pipeline.data_routing.scheduler import RandomLTDScheduler -from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict -from deepspeed.runtime.data_pipeline.data_routing.basic_layer import RandomLayerTokenDrop - -from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine -from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint - -from deepspeed.runtime.pipe.module import PipelineModule -from deepspeed.runtime.utils import get_ma_status -from deepspeed.ops.adam import FusedAdam -from deepspeed.moe.sharded_moe import TopKGate, MOELayer -from deepspeed.moe.layer import MoE -from deepspeed.moe.utils import is_moe_param -from deepspeed.git_version_info import version - -from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler -from deepspeed.utils.logging import print_json_dist, print_configuration - -from deepspeed.accelerator import get_accelerator - -from deepspeed.runtime.config import DtypeEnum - -from opensora.adaptor.zp_manager import zp_manager - -MEMORY_OPT_ALLREDUCE_SIZE = 500000000 - -DeepSpeedOptimizerCallable = \ - Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer] -DeepSpeedSchedulerCallable = Callable[[Optimizer], _LRScheduler] - -try: - import apex - from apex import amp - APEX_INSTALLED = True -except ImportError: - # Fail silently so we don't spam logs unnecessarily if user isn't using amp - APEX_INSTALLED = False - - -def split_half_float_double_sparse(tensors): - device_type = get_accelerator().device_name() - supported_types = [ - "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type), - "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type), - SparseTensor.type() - ] - - for t in tensors: - assert t.type() in supported_types, f"attempting to reduce an unsupported grad type: {t.type()}" - - buckets = [] - for i, dtype in enumerate(supported_types): - bucket = [t for t in tensors if t.type() == dtype] - if bucket: - buckets.append((dtype, bucket)) - return buckets - - -class EngineTimers(object): - r"""Wallclock timers for DeepSpeedEngine""" - - def __init__(self, enable_micro_timers, enable_global_timers): - self.forward_timers = [] - self.backward_timers = [] - self.backward_inner_timers = [] - self.backward_reduce_timers = [] - self.step_timers = [] - self.global_timers = [] - self.micro_timers = [] - - if enable_micro_timers: - self.forward_timers += [FORWARD_MICRO_TIMER] - self.backward_timers += [BACKWARD_MICRO_TIMER] - self.backward_inner_timers += [BACKWARD_INNER_MICRO_TIMER] - self.backward_reduce_timers += [BACKWARD_REDUCE_MICRO_TIMER] - self.step_timers += [STEP_MICRO_TIMER] - self.micro_timers += [ - FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, - STEP_MICRO_TIMER - ] - - if enable_global_timers: - self.forward_timers += [FORWARD_GLOBAL_TIMER] - self.backward_timers += [BACKWARD_GLOBAL_TIMER] - self.backward_inner_timers += [BACKWARD_INNER_GLOBAL_TIMER] - self.backward_reduce_timers += [BACKWARD_REDUCE_GLOBAL_TIMER] - self.step_timers += [STEP_GLOBAL_TIMER] - self.global_timers += [ - FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, - STEP_GLOBAL_TIMER - ] - - -class DeepSpeedEngine(Module): - r"""DeepSpeed engine for training.""" - - def __init__( - self, - args, - model, - optimizer=None, - model_parameters=None, - training_data=None, - lr_scheduler=None, - mpu=None, - dist_init_required=None, - collate_fn=None, - config=None, - config_class=None, - dont_change_device=False, - ): - super(DeepSpeedEngine, self).__init__() - self.dont_change_device = dont_change_device - self.client_optimizer = optimizer - self.client_lr_scheduler = lr_scheduler - self.training_data = training_data - self.collate_fn = collate_fn - self.mpu = mpu - self.all_to_all_group = None - self.data_parallel_group = None - self.global_steps = 0 - self.global_samples = 0 - self.micro_steps = 0 - self.skipped_steps = 0 - self.gradient_average = True - self.warn_unscaled_loss = True - self.config = config - self._config = config_class - self.loaded_checkpoint_mp_world_size = None - self.loaded_checkpoint_dp_world_size = None - self.loaded_checkpoint_zp_world_size = None - self.enable_backward_allreduce = True - self.progressive_layer_drop = None - self.eigenvalue = None - self.block_eigenvalue = None - self.gas_boundary_ctr = 0 - self.dist_backend = get_accelerator().communication_backend_name() - self.has_moe_layers = False - self.num_experts = [] - self.gate_modules = [] - self.moe_layers = [] - self._step_applied = False - self._global_grad_norm = None - self.use_ds_comm = False # False --> Use torch.dist, True --> Use ds.comm backend. - - self.checkpoint_engine = None - - self._is_gradient_accumulation_boundary = None - self.scale_wrt_gas = None - self.losses = 0.0 - - # for debug purposes - can then debug print: debug_get_module_name(module) - debug_extract_module_and_param_names(model) - - self._do_args_sanity_check(args) - self._configure_with_arguments(args, mpu) - self._do_sanity_check() - see_memory_usage(f"DeepSpeed Engine: After args sanity test", force=self.memory_breakdown()) - if mpu is not None: - if self.elasticity_enabled(): - if not self.is_elastic_model_parallel_supported(): - assert not self.elasticity_enabled(), ("Elasticity is not currently supported" - " with model parallelism.") - - self._set_distributed_vars(args) - - dist.configure(self._config) - - self.monitor = MonitorMaster(self._config.monitor_config) - - see_memory_usage( - f"DeepSpeed Engine: Before configure distributed model", - force=self.memory_breakdown(), - ) - - self.pipeline_parallelism = isinstance(model, PipelineModule) - - # Configure distributed model - self._configure_distributed_model(model) - - # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict - self.param_names = {param: name for name, param in model.named_parameters()} - - self._get_model_parameters() - - see_memory_usage(f"DeepSpeed Engine: After configure distributed model") - - # Configure wall clock timers - self.timers = SynchronizedWallClockTimer() - # Throughput timer - self.tput_timer = ThroughputTimer( - batch_size=self.train_batch_size(), - steps_per_output=self.steps_per_print(), - monitor_memory=False, - ) - - log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0]) - - if self.flops_profiler_enabled(): - self.flops_profiler = FlopsProfiler(self.module, self, self.flops_profiler_recompute_fwd_factor()) - - if training_data: - self.training_dataloader = self.deepspeed_io(training_data) - else: - self.training_dataloader = None - - # Configure optimizer and scheduler - self.optimizer = None - self.basic_optimizer = None - self.lr_scheduler = None - has_optimizer = False - - if optimizer or self.optimizer_name(): - has_optimizer = True - # If no parameters given by init default to module parameters - if model_parameters is None: - model_parameters = self.module.parameters() - - # Convert model parameters from generator to list - if not isinstance(model_parameters, list): - model_parameters = list(model_parameters) - - if has_optimizer: - self._configure_optimizer(optimizer, model_parameters) - self._configure_lr_scheduler(lr_scheduler) - self._report_progress(0) - elif self.zero_optimization(): - # no optim selected but zero is enabled - self.optimizer = self._configure_zero_optimizer(optimizer=None) - elif self.bfloat16_enabled(): - self.optimizer = self._configure_bf16_optimizer(optimizer=None) - - # Hook optimizer for snip_momentum pruning - if hasattr(model, 'pruners'): - from deepspeed.compression.helper import rewrite_optimizer_step - self.optimizer.pruners = model.pruners - rewrite_optimizer_step(self.optimizer) - - # Bookkeeping for sparse support - self.sparse_tensor_module_names = set() - # if self.sparse_gradients_enabled(): - for name, module in self.module.named_modules(): - if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)) and self.sparse_gradients_enabled(): - self.sparse_tensor_module_names.add(name + ".weight") - logger.info("Will convert {} to sparse tensor during training".format(name)) - - self.save_non_zero_checkpoint = False - self.save_zero_checkpoint = False - if not isinstance(self.optimizer, DeepSpeedZeRoOffload): - self._configure_checkpointing(dist_init_required) - - if self.eigenvalue_enabled(): - self.eigenvalue = self._configure_eigenvalue() - - if self.pld_enabled(): - self.progressive_layer_drop = self._configure_progressive_layer_drop() - - if self.curriculum_enabled_legacy(): - self.curriculum_scheduler_legacy = self._configure_curriculum_scheduler_legacy() - - if self.random_ltd_enabled(): - random_ltd_config = self.random_ltd_config() - random_ltd_config[RANDOM_LTD_GLOBAL_BATCH_SIZE] = self.train_batch_size() - random_ltd_config[RANDOM_LTD_MICRO_BATCH_SIZE] = self.train_micro_batch_size_per_gpu() - self.random_ltd_scheduler = self._configure_random_ltd_scheduler(random_ltd_config) - - # Engine timers - - self.engine_timers = EngineTimers(enable_micro_timers=self.wall_clock_breakdown(), - enable_global_timers=self.wall_clock_breakdown() - or self.flops_profiler_enabled()) - - if self.global_rank == 0: - self._config.print("DeepSpeedEngine configuration") - if self.dump_state(): - print_configuration(self, "DeepSpeedEngine") - - # Use torch (un)flatten ops - self.flatten = _flatten_dense_tensors - self.unflatten = _unflatten_dense_tensors - - def destroy(self): - if self.optimizer is not None and hasattr(self.optimizer, 'destroy'): - self.optimizer.destroy() - - def _get_model_parameters(self): - if self.autotuning_profile_model_info(): - self.autotuning_model_info = {} - num_params = 0 - trainable_num_params = 0 - - for p in self.module.parameters(): - # since user code might call deepspeed.zero.Init() before deepspeed.initialize(), need to check the attribute to check if the parameter is partitioned in zero 3 already or not - n = 0 - if hasattr(p, "ds_tensor"): # if the parameter is partitioned in zero 3 - n += p.ds_numel - else: # if the parameter is not partitioned in zero 3 yet - n += p.numel() - num_params += n - if p.requires_grad: - trainable_num_params += n - if self.global_rank == 0: - self.autotuning_model_info["num_params"] = num_params * self.mp_world_size - self.autotuning_model_info["trainable_num_params"] = trainable_num_params * self.mp_world_size - - logger.info(f"model parameter = {num_params}") - - def get_batch_info(self): - """Get all training batch related settings. - Returns: - train_batch_size (int): The effective training batch size. This is the amount of data - samples that leads to one step of model update. - train_micro_batch_size_per_gpu (int): Batch size to be processed by one GPU in one - step (without gradient accumulation). - gradient_accumulation_steps (int): Number of training steps to accumulate gradients - before averaging and applying them. - """ - return ( - self.train_batch_size, - self.train_micro_batch_size_per_gpu, - self.gradient_accumulation_steps, - ) - - def set_train_batch_size(self, train_batch_size): - """Adjust the global batch size by increasing or decreasing the number of - micro-batches (i.e., gradient accumulation steps). The size of each micro-batch - (i.e., ``train_micro_batch_size_per_gpu``) is not changed. - Args: - train_batch_size (int): The new global batch size for training. - Raises: - ValueError: if ``train_batch_size`` is not divisible by the - configured micro-batch size and data parallelism. - """ - if train_batch_size % (self.train_micro_batch_size_per_gpu() * self.dp_world_size) != 0: - #print(f'{train_batch_size=} {self.train_micro_batch_size_per_gpu()=} {self.dp_world_size=}') - raise ValueError(f'Train batch size must be divisible by micro-batch data parallelism') - new_gas = train_batch_size // (self.train_micro_batch_size_per_gpu() * self.dp_world_size) - # overwrite config - self._config.train_batch_size = train_batch_size - self._config.gradient_accumulation_steps = new_gas - - def set_train_micro_batch_size(self, micro_batch_size): - """Adjust the micro batch size(i.e., the micro batch size in every data parallel group), - while keep the gradient accumulation steps the same. - Args: - micro_batch_size (int): The new micro batch size for training. - """ - # overwrite config - new_global_batch_size = micro_batch_size * self._config.gradient_accumulation_steps * self.dp_world_size - self._config.train_batch_size = new_global_batch_size - self._config.train_micro_batch_size_per_gpu = micro_batch_size - - def set_data_post_process_func(self, post_process_func): - if self.training_dataloader is not None: - self.training_dataloader.post_process_func = post_process_func - - def set_custom_curriculum_learning_schedule(self, schedule_func_dict): - if self.training_dataloader is not None and self.curriculum_learning_enabled(): - self.training_dataloader.data_sampler.set_custom_curriculum_learning_schedule(schedule_func_dict) - - def get_global_grad_norm(self) -> float: - """Return the 2-norm of all gradients. If there is model parallelism, - the norm will be global. - The computed norm will be cached and reused until the next step() pass. - .. note:: - In the presence of model parallelism, this is a collective call - and acts as a barrier among ``mpu.get_model_parallel_group()``. - Returns: - float: norm - """ - return self._global_grad_norm - - def __getattr__(self, name): - """ - Pass through attributes defined in the model if they are not overridden by ds-engine. - """ - - _module = {} - if "module" in self.__dict__: - _module = self.__dict__['module'] - if name in dir(self): - return getattr(self, name) - elif name in dir(_module): - return getattr(_module, name) - else: - raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") - - def checkpoint_tag_validation_enabled(self): - return self._config.checkpoint_tag_validation_enabled - - def checkpoint_tag_validation_fail(self): - return self._config.checkpoint_tag_validation_fail - - def elasticity_enabled(self): - return self._config.elasticity_enabled - - def is_elastic_model_parallel_supported(self): - if self.elasticity_enabled(): - # Add code for finding number of GPUs per node automatically - if self._config.num_gpus_per_node % self._config.elastic_model_parallel_size == 0: - return True - else: - return False - - def pld_enabled(self): - return self._config.pld_enabled - - def pld_params(self): - return self._config.pld_params - - def pld_theta(self): - return self.pld_params()[PLD_THETA] - - def pld_gamma(self): - return self.pld_params()[PLD_GAMMA] - - def eigenvalue_enabled(self): - return self._config.eigenvalue_enabled - - def eigenvalue_verbose(self): - return self._config.eigenvalue_verbose - - def eigenvalue_max_iter(self): - return self._config.eigenvalue_max_iter - - def eigenvalue_tol(self): - return self._config.eigenvalue_tol - - def eigenvalue_stability(self): - return self._config.eigenvalue_stability - - def eigenvalue_gas_boundary_resolution(self): - return self._config.eigenvalue_gas_boundary_resolution - - def eigenvalue_layer_name(self): - return self._config.eigenvalue_layer_name - - def eigenvalue_layer_num(self): - return self._config.eigenvalue_layer_num - - def curriculum_enabled_legacy(self): - return self._config.curriculum_enabled_legacy - - def curriculum_params_legacy(self): - return self._config.curriculum_params_legacy - - def data_efficiency_enabled(self): - return self._config.data_efficiency_enabled - - def data_efficiency_config(self): - return self._config.data_efficiency_config - - def data_sampling_enabled(self): - return self._config.data_efficiency_config[DATA_SAMPLING][DATA_SAMPLING_ENABLED] - - def data_sampling_config(self): - return self._config.data_efficiency_config[DATA_SAMPLING] - - def curriculum_learning_enabled(self): - return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED] - - def curriculum_learning_config(self): - return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING] - - def random_ltd_enabled(self): - return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD][RANDOM_LTD_ENABLED] - - def random_ltd_config(self): - return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD] - - def random_ltd_initialize(self): - assert self.random_ltd_enabled() - random_ltd_config = self.random_ltd_config() - random_ltd_queue = deque([x for x in sorted(random_ltd_config[RANDOM_LTD_LAYER_ID])]) - count = 0 - for name, layer in self.module.named_modules(): - if isinstance(layer, RandomLayerTokenDrop): - if len(random_ltd_queue) != 0 and str(random_ltd_queue[0]) in name: ###[1,2,3] - layer.init_config(random_ltd_config, self.random_ltd_scheduler, count) - random_ltd_queue.popleft() - count += 1 - - if random_ltd_config[RANDOM_LTD_LAYER_NUM] != count: - raise ValueError(f'random_ltd_layer_num {random_ltd_config[RANDOM_LTD_LAYER_NUM]} must be \ - equivalent to the len of random_ltd_layer_id {count}') - - if random_ltd_config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: - assert self.client_lr_scheduler is None - raise ValueError(f'not yet support') - #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler) - - def wall_clock_breakdown(self): - return self._config.wall_clock_breakdown - - def flops_profiler_enabled(self): - return self._config.flops_profiler_config.enabled or self.autotuning_enabled() - - def flops_profiler_recompute_fwd_factor(self): - return self._config.flops_profiler_config.recompute_fwd_factor - - def flops_profiler_profile_step(self): - step = self._config.flops_profiler_config.profile_step - if self._config.autotuning_config.enabled: - step = self.autotuning_start_profile_step() - return step - - def flops_profiler_module_depth(self): - return self._config.flops_profiler_config.module_depth - - def flops_profiler_top_modules(self): - return self._config.flops_profiler_config.top_modules - - def flops_profiler_detailed(self): - if self._config.autotuning_config.enabled: - return False - return self._config.flops_profiler_config.detailed - - def flops_profiler_output_file(self): - return self._config.flops_profiler_config.output_file - - def memory_breakdown(self): - return self._config.memory_breakdown - - def autotuning_enabled(self): - return self._config.autotuning_config.enabled - - def autotuning_start_profile_step(self): - return self._config.autotuning_config.start_profile_step - - def autotuning_end_profile_step(self): - return self._config.autotuning_config.end_profile_step - - def autotuning_metric_path(self): - path = self._config.autotuning_config.metric_path - if not path: - path = os.path.join(os.getcwd(), "autotuning_metric.json") - return path - - def autotuning_model_info_path(self): - path = self._config.autotuning_config.model_info_path - if not path: - path = os.path.join(os.getcwd(), "autotuning_model_info.json") - return path - - def autotuning_metric(self): - return self._config.autotuning_config.metric - - def autotuning_profile_model_info(self): - return self.autotuning_enabled( - ) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get( - "profile", False) - - def sparse_gradients_enabled(self): - return self._config.sparse_gradients_enabled - - def train_batch_size(self): - return self._config.train_batch_size - - def train_micro_batch_size_per_gpu(self): - return self._config.train_micro_batch_size_per_gpu - - def optimizer_name(self): - return (self.client_optimizer.__class__.__name__ if self.client_optimizer else self._config.optimizer_name) - - def optimizer_params(self): - return self._config.optimizer_params - - def optimizer_legacy_fusion(self): - return self._config.optimizer_legacy_fusion - - def scheduler_name(self): - return self._config.scheduler_name - - def scheduler_params(self): - return self._config.scheduler_params - - def quantize_training(self): - return ( - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] - [WEIGHT_QUANTIZE_IN_FORWARD_ENABLED], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ENABLED], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_GROUPS], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] - [WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_CHANGE_RATIO], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_TYPE], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ROUNDING], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_VERBOSE], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_KERNEL], - ) - - def zero_optimization(self): - return self._config.zero_enabled - - def zero_allow_untested_optimizer(self): - return self._config.zero_allow_untested_optimizer - - def zero_force_ds_cpu_optimizer(self): - return self._config.zero_force_ds_cpu_optimizer - - def zero_reduce_scatter(self): - return self._config.zero_config.reduce_scatter - - def zero_overlap_comm(self): - return self._config.zero_config.overlap_comm - - def zero_offload_optimizer(self): - return self._config.zero_config.offload_optimizer - - def zero_offload_param(self): - return self._config.zero_config.offload_param - - def zero_use_cpu_optimizer(self): - if self._config.zero_config.offload_optimizer is not None: - return self._config.zero_config.offload_optimizer.device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme] - return False - - def zero_cpu_offload(self): - if self._config.zero_config.offload_optimizer is not None: - return self._config.zero_config.offload_optimizer.device == OffloadDeviceEnum.cpu - return False - - def zero_partial_offload(self): - return getattr(self._config.zero_config.offload_optimizer, "ratio", 1.0) - - def zero_sub_group_size(self): - return self._config.zero_config.sub_group_size - - def zero_optimization_stage(self): - return self._config.zero_optimization_stage - - def mics_shard_size(self): - return self._config.mics_shard_size - - def zero_reduce_bucket_size(self): - return self._config.zero_config.reduce_bucket_size - - def zero_multi_rank_bucket_allreduce(self): - return self._config.zero_config.use_multi_rank_bucket_allreduce - - def zero_allgather_bucket_size(self): - return self._config.zero_config.allgather_bucket_size - - def zero_optimization_partition_gradients(self): - return self.zero_optimization_stage() >= ZeroStageEnum.gradients - - def zero_optimization_partition_weights(self): - return self.zero_optimization_stage() >= ZeroStageEnum.weights - - def is_first_weights_partition_group(self): - ret = True if self.mics_shard_size() < 0 \ - and self.zero_optimization_partition_weights() else False - if self.mics_shard_size() > 0 and self.global_rank < self.mics_shard_size(): - ret = True - return ret - - def zero_contiguous_gradients(self): - return self._config.zero_config.contiguous_gradients - - def zero_load_from_fp32_weights(self): - return self._config.zero_config.load_from_fp32_weights - - def zero_elastic_checkpoint(self): - return self._config.zero_config.elastic_checkpoint - - def zero_max_live_parameters(self): - return self._config.zero_config.max_live_parameters - - def zero_max_reuse_distance(self): - return self._config.zero_config.max_reuse_distance - - def zero_prefetch_bucket_size(self): - return self._config.zero_config.prefetch_bucket_size - - def zero_param_persistence_threshold(self): - return self._config.zero_config.param_persistence_threshold - - def zero_model_persistence_threshold(self): - return self._config.zero_config.model_persistence_threshold - - def zero_gather_16bit_weights_on_model_save(self): - return self._config.zero_config.gather_16bit_weights_on_model_save - - def zero_grad_hooks(self): - return self._config.zero_config.grad_hooks - - def zero_legacy_stage1(self): - return self._config.zero_config.legacy_stage1 - - def zero_ignore_unused_parameters(self): - return self._config.zero_config.ignore_unused_parameters - - def graph_harvesting(self): - return self._config.graph_harvesting - - def fp16_enabled(self): - return self._config.fp16_enabled - - def bfloat16_enabled(self): - return self._config.bfloat16_enabled - - def fp16_master_weights_and_gradients(self): - return self._config.fp16_master_weights_and_gradients - - def amp_enabled(self): - return self._config.amp_enabled - - def amp_params(self): - return self._config.amp_params - - def fp16_auto_cast(self): - return self._config.fp16_auto_cast - - def loss_scale(self): - return self._config.loss_scale - - def gradient_accumulation_steps(self): - return self._config.gradient_accumulation_steps - - def use_node_local_storage(self): - return self._config.use_node_local_storage - - def load_universal_checkpoint(self): - return self._config.load_universal_checkpoint - - @property - def communication_data_type(self): - res = self._config.communication_data_type - if res is not None: - return res - - if self.fp16_enabled(): - return torch.float16 - - if self.bfloat16_enabled(): - return torch.bfloat16 - - return torch.float32 - - @communication_data_type.setter - def communication_data_type(self, value): - self._config.communication_data_type = value - - def postscale_gradients(self): - return not self._config.prescale_gradients - - def gradient_predivide_factor(self): - return self._config.gradient_predivide_factor - - def steps_per_print(self): - return self._config.steps_per_print - - def zero_allgather_partitions(self): - return self._config.zero_config.allgather_partitions - - def zero_round_robin_gradients(self): - return self._config.zero_config.round_robin_gradients - - def zero_hpz_partition_size(self): - return self._config.zero_config.zero_hpz_partition_size - - def zero_quantized_weights(self): - return self._config.zero_config.zero_quantized_weights - - def zero_quantized_nontrainable_weights(self): - return self._config.zero_config.zero_quantized_nontrainable_weights - - def zero_quantized_gradients(self): - return self._config.zero_config.zero_quantized_gradients - - def dump_state(self): - return self._config.dump_state - - def gradient_clipping(self): - return self._config.gradient_clipping - - def dynamic_loss_scale(self): - return self._config.loss_scale == 0 - - def initial_dynamic_scale(self): - return self._config.initial_dynamic_scale - - def dynamic_loss_scale_args(self): - return self._config.dynamic_loss_scale_args - - def swap_tensor_config(self): - return self._config.swap_tensor_config - - def aio_config(self): - return self._config.aio_config - - def get_data_types(self): - model_dtype = torch.float32 - if self.fp16_enabled(): - model_dtype = torch.float16 - elif self.bfloat16_enabled(): - model_dtype = torch.bfloat16 - - if self._config.grad_accum_dtype is None: - if model_dtype == torch.bfloat16 and not self.zero_optimization(): - grad_accum_dtype = torch.float32 - else: - grad_accum_dtype = model_dtype - else: - grad_accum_dtype = DtypeEnum(self._config.grad_accum_dtype).value - - return (model_dtype, grad_accum_dtype) - - def _optimizer_has_ckpt_event_prologue(self): - return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_prologue') - - def _optimizer_has_ckpt_event_epilogue(self): - return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_epilogue') - - def _configure_lr_scheduler(self, client_lr_scheduler): - # First check for scheduler in json configuration - lr_scheduler = self._scheduler_from_config(self.optimizer) - if lr_scheduler: - log_dist(f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}", ranks=[0]) - self.lr_scheduler = lr_scheduler - else: - if isinstance(client_lr_scheduler, Callable): - log_dist('DeepSpeed using client callable to create LR scheduler', ranks=[0]) - self.lr_scheduler = client_lr_scheduler(self.basic_optimizer) - else: - log_dist('DeepSpeed using client LR scheduler', ranks=[0]) - self.lr_scheduler = client_lr_scheduler - - log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) - - def _configure_checkpointing(self, dist_init_required): - self.checkpoint_engine = TorchCheckpointEngine() - - if self._config is not None and self._config.nebula_config.enabled: - try: - from deepspeed.runtime.checkpoint_engine.nebula_checkpoint_engine import \ - NebulaCheckpointEngine - self.checkpoint_engine = NebulaCheckpointEngine(config_params=self._config.nebula_config) - except ImportError as err: - logger.error(f"No torch_nebula was found! Will fall back to torch.save. Details: {err}") - self.checkpoint_engine = TorchCheckpointEngine() - - dp_rank = groups._get_sequence_data_parallel_rank() - - rank = self.local_rank if self.use_node_local_storage() else dp_rank - - # only the first data parallel process needs to store the model checkpoint - # if you want to use node local storage this must be done by rank 0 on each - # node - self.save_non_zero_checkpoint = (rank == 0) or (self.zero_optimization_partition_weights() - and self.is_first_weights_partition_group()) - - if self.zero_optimization() or self.bfloat16_enabled(): - param_rank = dist.get_rank(group=self.optimizer.zp_process_group) - - # Only the first parameter parallel process needs to store the - # optimizer state checkpoints for zero - self.save_zero_checkpoint = param_rank == dp_rank - - def _scheduler_from_config(self, optimizer): - scheduler_name = self.scheduler_name() - if scheduler_name is not None: - if hasattr(lr_schedules, scheduler_name): - scheduler = getattr(lr_schedules, scheduler_name) - else: - assert hasattr(torch.optim.lr_scheduler, - scheduler_name), f"DeepSpeed does not recognize LR scheduler {scheduler_name}" - - scheduler = getattr(torch.optim.lr_scheduler, scheduler_name) - - scheduler_params = self.scheduler_params() - instantiated_scheduler = scheduler(optimizer, **scheduler_params) - return instantiated_scheduler - else: - return None - - def _set_distributed_vars(self, args): - device_rank = args.device_rank if args is not None and hasattr(args, 'device_rank') else self.local_rank - if device_rank >= 0: - get_accelerator().set_device(device_rank) - self.device = torch.device(get_accelerator().device_name(), device_rank) - self.world_size = dist.get_world_size() - self.global_rank = dist.get_rank() - else: - self.world_size = 1 - self.global_rank = 0 - self.device = torch.device(get_accelerator().device_name()) - - # Configure based on command line arguments - def _configure_with_arguments(self, args, mpu): - # After the distributed backend is initialized we are guaranteed the LOCAL_RANK - # environment variable is set. We must align args.local_rank to this value for - # backwards compatibility with scripts relying on [args|self].local_rank containing - # the correct local rank info. _do_args_sanity_check will ensure this is the case. - - if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: - ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK") - local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank) - assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({ompi_local_rank}), " \ - "not sure how to proceed as we're seeing conflicting local rank info." - os.environ['LOCAL_RANK'] = local_rank - - self.local_rank = int(os.environ['LOCAL_RANK']) - if hasattr(args, 'local_rank'): - args.local_rank = self.local_rank - - # Validate command line arguments - def _do_args_sanity_check(self, args): - assert "LOCAL_RANK" in os.environ or "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment " \ - "variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch's launcher. If using a " \ - "different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed." - - if hasattr(args, 'local_rank') and args.local_rank is not None: - assert isinstance(args.local_rank, - int), f"args.local_rank of {args.local_rank} is an unknown type {type(args.local_rank)}" - if args.local_rank >= 0: - env_local_rank = int(os.environ.get("LOCAL_RANK")) - assert ( - env_local_rank == args.local_rank - ), f"Mismatch in local rank setting, args.local_rank={args.local_rank} but env['LOCAL_RANK']={env_local_rank}." - - def _is_supported_optimizer(self, optimizer_name): - return (optimizer_name in DEEPSPEED_OPTIMIZERS or getattr(torch.optim, optimizer_name, None) is not None) - - def _supported_optims(self): - FairseqOptimizer = None - try: - from fairseq.optim.fairseq_optimizer import FairseqOptimizer - except ImportError: - pass - - expected_optim_types = [Optimizer] - if FairseqOptimizer: - # fairseq optims are not torch.optim objects - expected_optim_types.append(FairseqOptimizer) - return expected_optim_types - - # Validate configuration based on command line arguments - def _do_sanity_check(self): - expected_optim_types = self._supported_optims() - expected_optim_types += [type(None), Callable] - assert isinstance(self.client_optimizer, tuple(expected_optim_types)), \ - f'Client Optimizer is of unexpected type {type(self.client_optimizer)}' - - if not self.client_optimizer: - if self.optimizer_name() is not None: - assert self._is_supported_optimizer( - self.optimizer_name()), "{} is not a supported DeepSpeed Optimizer".format(self.optimizer_name()) - - if (self.optimizer_name() == LAMB_OPTIMIZER or self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER): - assert (self.dynamic_loss_scale()), "DeepSpeed {} optimizer requires dynamic loss scaling".format( - self.optimizer_name()) - - # Detect invalid combinations of client optimizer and client scheduler - if isinstance(self.client_lr_scheduler, _LRScheduler): - assert isinstance(self.client_optimizer, Optimizer), \ - f'Client Optimizer (type = {type(self.client_optimizer)} is not instantiated but Client LR Scheduler is instantiated' - - def _broadcast_model(self): - - def is_replicated(p): - if hasattr(p, "ds_status") and p.ds_status is not ZeroParamStatus.AVAILABLE: - return False - return True - - for p in self.module.parameters(): - # Broadcast the model for different parameters - if is_moe_param(p): - if torch.is_tensor(p) and is_replicated(p): - dist.broadcast(p, - groups._get_expert_broadcast_src_rank(p.group_name), - group=self.expert_data_parallel_group[p.group_name]) - else: - if torch.is_tensor(p) and is_replicated(p): - dist.broadcast(p, groups._get_broadcast_src_rank(), group=self.seq_data_parallel_group) - - @staticmethod - def __check_params(model: Module, dtype: torch.dtype) -> None: - return - if not all(param.dtype == dtype for param in model.parameters()) and dist.get_rank() == 0: - raise ValueError(f"{dtype} is enabled but the following parameters have dtype that is " - f"not {dtype}: " - f"{[(n, p.dtype) for n, p in model.named_parameters() if p.dtype != dtype]}") - - def _set_client_model(self, model): - # register client model in _modules so that nn.module methods work correctly - modules = self.__dict__.get('_modules') - modules['module'] = model - # register module attribute in engine but avoid getattr - self.__dict__['module'] = model - - def _configure_distributed_model(self, model): - self._set_client_model(model) - is_zero_init_model = self.zero_optimization_partition_weights() and any( - [hasattr(param, "ds_id") for param in self.module.parameters()]) - - if self.fp16_enabled(): - if is_zero_init_model: - self.__check_params(self.module, torch.half) - self.module.half() - elif self.bfloat16_enabled(): - if is_zero_init_model: - self.__check_params(self.module, torch.bfloat16) - self.module.bfloat16() - else: - self.__check_params(self.module, torch.float) - - # zero.Init() handles device placement of model - if not (self.dont_change_device or is_zero_init_model): - self.module.to(self.device) - - # MoE related initialization - for _, module in self.module.named_modules(): - if isinstance(module, MoE): - self.has_moe_layers = True - self.num_experts.append(module.num_experts) - - if self.has_moe_layers: - for _, module in self.module.named_modules(): - if isinstance(module, TopKGate): - self.gate_modules.append(module) - if self.wall_clock_breakdown(): - module.wall_clock_breakdown = True - if isinstance(module, MOELayer): - self.moe_layers.append(module) - if self.wall_clock_breakdown(): - module.wall_clock_breakdown = True - - # Pass the mpu from here to groups. For subsequent use, just query groups - if self.mpu is not None: - groups.mpu = self.mpu - - # Set deepspeed parallelism spec. for the model including expert parallelism - for _, module in self.module.named_modules(): - if hasattr(module, 'set_deepspeed_parallelism'): - module.set_deepspeed_parallelism(self._config.use_data_before_expert_parallel_) - - # Query the groups module to get information about various parallel groups - self.local_all_to_all_group = None - if self.zero_quantized_gradients(): - log_dist("Using quantized gradients", ranks=[0]) - self.local_all_to_all_group = groups._get_local_all_to_all_group() - self.data_parallel_group = groups._get_data_parallel_group() - self.dp_world_size = groups._get_data_parallel_world_size() - self.zp_world_size = zp_manager.zp_size - self.seq_data_parallel_group = groups._get_sequence_data_parallel_group() - self.seq_dp_world_size = groups._get_sequence_data_parallel_world_size() - self.mp_world_size = groups._get_model_parallel_world_size() - self.expert_parallel_group = groups._get_expert_parallel_group_dict() - self.expert_data_parallel_group = groups._get_expert_data_parallel_group_dict() - self.sequence_parallel_size = groups._get_sequence_parallel_world_size() - if self.sequence_parallel_size > 1: - self.communication_data_type = self._config.seq_parallel_communication_data_type - - if not (self.amp_enabled() or is_zero_init_model): - self._broadcast_model() - - # check if parameters are duplicated in optimizer param_groups - def _check_for_duplicates(self, optimizer): - for name, param in self.module.named_parameters(): - param_id = id(param) - - def ids_list(group): - return [id(param) for param in group] - - occurrence = sum([ - ids_list(group['params']).count(param_id) if param_id in ids_list(group['params']) else 0 - for group in optimizer.param_groups - ]) - assert occurrence <= 1, f"Parameter with name: {name} occurs multiple times in optimizer.param_groups. Make sure it only appears once to prevent undefined behavior." - - def _do_optimizer_sanity_check(self, basic_optimizer): - model_dtype, grad_accum_dtype = self.get_data_types() - zero_enabled = self.zero_optimization() - amp_enabled = self.amp_enabled() - # config based assertions - assert ( - not (amp_enabled and zero_enabled) - ), "Amp and ZeRO are not currently compatible, please use (legacy) fp16 mode which performs similar to amp opt_mode=O2" - if zero_enabled: - if not is_zero_supported_optimizer(basic_optimizer): - assert ( - self.zero_allow_untested_optimizer() - ), 'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.' - - if self.global_rank == 0: - logger.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****") - if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage( - ) == 1 and not self.zero_cpu_offload(): - return BFLOAT16 - return ZERO_OPTIMIZATION - elif amp_enabled: - if model_dtype != grad_accum_dtype: - raise NotImplementedError( - "Model data type and gradient accumulation data type must be equal to use Amp") - if model_dtype == torch.bfloat16 or model_dtype == torch.float16: - raise NotImplementedError("Cannot enable both amp with (legacy) fp16 or bfloat16 mode") - try: - logger.info("Initializing Apex amp from: {}".format(amp.__path__)) - except NameError: - # If apex/amp is available it will be imported above - raise RuntimeError("Unable to import apex/amp, please make sure it is installed") - return AMP - # data type checks - elif model_dtype == grad_accum_dtype: - if model_dtype == torch.bfloat16: - if self.pipeline_parallelism: - logger.warning( - "**** BF16 gradient accumulation is not safe numerically with large number of accumulation steps, proceed with caution *****" - ) - return BFLOAT16 - else: - raise NotImplementedError( - "Bfloat16 wrapper must use a gradient accumulation type of fp32, enable ZeRO to use Bfloat16 gradient accumulation" - ) - if model_dtype == torch.float16: - return FP16 - # else optimizer_wrapper = None - elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32: - return BFLOAT16 - else: - raise NotImplementedError("unsupported mix of model dtype and gradient accumulation type") - - return None - - # Configure optimizer - def _configure_optimizer(self, client_optimizer, model_parameters): - if client_optimizer is None: - basic_optimizer = self._configure_basic_optimizer(model_parameters) - log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0]) - else: - if isinstance(client_optimizer, tuple(self._supported_optims())): - basic_optimizer = client_optimizer - log_dist('Using client Optimizer as basic optimizer', ranks=[0]) - else: - basic_optimizer = client_optimizer(model_parameters) - log_dist('Using client callable to create basic optimizer', ranks=[0]) - - if self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam): - if self.zero_force_ds_cpu_optimizer(): - msg = f'You are using ZeRO-Offload with a client provided optimizer ({type(basic_optimizer)}) which in most cases will yield poor performance. Please either use deepspeed.ops.adam.DeepSpeedCPUAdam or set an optimizer in your ds-config (https://www.deepspeed.ai/docs/config-json/#optimizer-parameters). If you really want to use a custom optimizer w. ZeRO-Offload and understand the performance impacts you can also set <"zero_force_ds_cpu_optimizer": false> in your configuration file.' - raise ZeRORuntimeException(msg) - - basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0] - log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0]) - - self._check_for_duplicates(basic_optimizer) - - self.basic_optimizer = basic_optimizer - log_dist("DeepSpeed Basic Optimizer = {}".format(basic_optimizer.__class__.__name__), ranks=[0]) - - optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) - - if optimizer_wrapper == ZERO_OPTIMIZATION: - self.optimizer = self._configure_zero_optimizer(basic_optimizer) - elif optimizer_wrapper == AMP: - amp_params = self.amp_params() - log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) - model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params) - self._set_client_model(model) - self._broadcast_model() - # TODO: maybe need to broadcast experts differently? - elif optimizer_wrapper == FP16: - self.optimizer = self._configure_fp16_optimizer(basic_optimizer) - elif optimizer_wrapper == BFLOAT16: - self.optimizer = self._configure_bf16_optimizer(basic_optimizer) - else: - self.optimizer = basic_optimizer - - log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer_name()), ranks=[0]) - - self.compression_scheduler = self._configure_compression_scheduler() - self.quantizer = self._configure_quantization() - - def _configure_basic_optimizer(self, model_parameters): - optimizer_parameters = self.optimizer_params() - if optimizer_parameters is None: - optimizer_parameters = {} - # print(optimizer_parameters.keys()) - if "max_grad_norm" in optimizer_parameters.keys(): - raise ValueError( - "'max_grad_norm' is not supported as an optimizer parameter, please switch to using the deepspeed parameter 'gradient_clipping' see: https://www.deepspeed.ai/docs/config-json/#gradient-clipping for more details" - ) - - if self.optimizer_name() in [ADAM_OPTIMIZER, ADAMW_OPTIMIZER]: - torch_adam = optimizer_parameters.pop(TORCH_ADAM_PARAM, False) - adam_w_mode = optimizer_parameters.pop(ADAM_W_MODE, ADAM_W_MODE_DEFAULT) - - # Optimizer name of Adam forces AdamW logic unless adam_w_mode is explicitly set - effective_adam_w_mode = self.optimizer_name() == ADAMW_OPTIMIZER or adam_w_mode - - if torch_adam: - if not effective_adam_w_mode: - optimizer = torch.optim.Adam(model_parameters, **optimizer_parameters) - else: - optimizer = torch.optim.AdamW(model_parameters, **optimizer_parameters) - else: - if self.zero_use_cpu_optimizer(): - from deepspeed.ops.adam import DeepSpeedCPUAdam - optimizer = DeepSpeedCPUAdam(model_parameters, - **optimizer_parameters, - adamw_mode=effective_adam_w_mode) - else: - from deepspeed.ops.adam import FusedAdam - - optimizer = FusedAdam( - model_parameters, - **optimizer_parameters, - adam_w_mode=effective_adam_w_mode, - ) - - elif self.optimizer_name() == ADAGRAD_OPTIMIZER: - if self.zero_use_cpu_optimizer(): - from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad - optimizer = DeepSpeedCPUAdagrad(model_parameters, **optimizer_parameters) - else: - optimizer = torch.optim.Adagrad(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == LAMB_OPTIMIZER: - from deepspeed.ops.lamb import FusedLamb - - optimizer = FusedLamb(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER: - assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO" - from deepspeed.runtime.fp16.onebit.adam import OnebitAdam - - optimizer = OnebitAdam(model_parameters, self, **optimizer_parameters) - if not self.fp16_enabled(): - logger.warning(f"Currently the convergence of 1-bit Adam is only verified under FP16") - elif self.optimizer_name() == ZERO_ONE_ADAM_OPTIMIZER: - assert not self.zero_optimization(), "0/1 Adam is not compatible with ZeRO" - from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam - - optimizer = ZeroOneAdam(model_parameters, self, **optimizer_parameters) - if not self.fp16_enabled(): - logger.warning(f'Currently the convergence of 0/1 Adam is only verified under FP16') - elif self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER: - assert not self.zero_optimization(), "1bit-Lamb is not compatible with ZeRO" - from deepspeed.runtime.fp16.onebit.lamb import OnebitLamb - - optimizer = OnebitLamb(model_parameters, self, **optimizer_parameters) - if not self.fp16_enabled(): - logger.warning(f"Currently the convergence of 1-bit Lamb is only verified under FP16") - elif self.optimizer_name() == LION_OPTIMIZER: - if self.zero_use_cpu_optimizer(): - from deepspeed.ops.lion import DeepSpeedCPULion - optimizer = DeepSpeedCPULion(model_parameters, **optimizer_parameters) - else: - from deepspeed.ops.lion import FusedLion - optimizer = FusedLion(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == MUADAM_OPTIMIZER: - try: - from mup import MuAdam - except ImportError: - logger.error(f"Install mup to use MuAdam optimizer") - optimizer = MuAdam(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == MUADAMW_OPTIMIZER: - try: - from mup import MuAdamW - except ImportError: - logger.error(f"Install mup to use MuAdamW optimizer") - optimizer = MuAdamW(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == MUSGD_OPTIMIZER: - try: - from mup import MuSGD - except ImportError: - logger.error(f"Install mup to use MuSGD optimizer") - optimizer = MuSGD(model_parameters, **optimizer_parameters) - else: - torch_optimizer = getattr(torch.optim, self.optimizer_name()) - optimizer = torch_optimizer(model_parameters, **optimizer_parameters) - return optimizer - - def _configure_compression_scheduler(self): - return compression_scheduler(self.module, self._config.compression_config) - - def _configure_random_ltd_scheduler(self, configs): - return RandomLTDScheduler(configs) - - def _configure_quantization(self): - ( - quantize_weight_in_forward, - quantize_enabled, - q_groups, - q_mixed_fp16, - q_change_ratio, - q_type, - q_rounding, - q_verbose, - use_quantizer_kernel, - ) = self.quantize_training() - if quantize_enabled and not quantize_weight_in_forward: - assert self.fp16_enabled( - ), "MoQ (quantize in optimization step) weight quantization is only supported for FP16" - quantizer = None - if quantize_enabled and not quantize_weight_in_forward: - from deepspeed.runtime.quantize import Quantizer - - quantizer = Quantizer( - q_groups, - q_mixed_fp16, - q_change_ratio, - q_type, - q_rounding, - q_verbose, - self.eigenvalue_enabled(), - use_quantizer_kernel, - self.eigenvalue_layer_num() if self.eigenvalue_enabled() else 0, - ) - return quantizer - - def _configure_fp16_optimizer(self, optimizer): - initial_dynamic_scale = self.initial_dynamic_scale() - dynamic_loss_args = self.dynamic_loss_scale_args() - clip_grad = self.gradient_clipping() - if APEX_INSTALLED: - fused_opts = (apex.optimizers.FusedAdam, FusedAdam) - else: - fused_opts = FusedAdam - if isinstance(optimizer, fused_opts) \ - or self.optimizer_name() in [ONEBIT_ADAM_OPTIMIZER, ZERO_ONE_ADAM_OPTIMIZER]: - if self.dynamic_loss_scale(): - log_dist(f'Creating fp16 optimizer with dynamic loss scale', ranks=[0]) - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - optimizer = FP16_Optimizer( - optimizer, - deepspeed=self, - dynamic_loss_scale=True, - initial_dynamic_scale=initial_dynamic_scale, - dynamic_loss_args=dynamic_loss_args, - mpu=self.mpu, - clip_grad=clip_grad, - fused_adam_legacy=self.optimizer_legacy_fusion(), - timers=timers, - has_moe_layers=self.has_moe_layers, - ) - else: - log_dist(f'Creating fp16 optimizer with static loss scale: {self.loss_scale()}', ranks=[0]) - optimizer = FP16_Optimizer( - optimizer, - deepspeed=self, - static_loss_scale=self.loss_scale(), - mpu=self.mpu, - clip_grad=clip_grad, - fused_adam_legacy=self.optimizer_legacy_fusion(), - has_moe_layers=self.has_moe_layers, - ) - else: - log_dist(f'Creating fp16 unfused optimizer with dynamic loss scale', ranks=[0]) - optimizer = FP16_UnfusedOptimizer( - optimizer, - deepspeed=self, - static_loss_scale=self.loss_scale(), - dynamic_loss_scale=self.dynamic_loss_scale(), - dynamic_loss_args=dynamic_loss_args, - mpu=self.mpu, - clip_grad=clip_grad, - fused_lamb_legacy=self.optimizer_name() == LAMB_OPTIMIZER, - ) - - return optimizer - - def _configure_bf16_optimizer(self, optimizer): - clip_grad = self.gradient_clipping() - - if optimizer is None: - optimizer = DummyOptim(list(self.module.parameters())) - - log_dist('Creating BF16 optimizer', ranks=[0]) - - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - optimizer = BF16_Optimizer(optimizer, - self.param_names, - mpu=self.mpu, - clip_grad=clip_grad, - allgather_bucket_size=self.zero_allgather_bucket_size(), - dp_process_group=self.seq_data_parallel_group, - timers=timers, - grad_acc_dtype=self.get_data_types()[1], - graph_harvesting=self.graph_harvesting()) - - return optimizer - - def _configure_zero_optimizer(self, optimizer): - zero_stage = self.zero_optimization_stage() - - mics_shard_size = self.mics_shard_size() - model_dtype, gradient_accumulation_dtype = self.get_data_types() - - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - - if optimizer is None: - optimizer = DummyOptim(list(self.module.parameters())) - - if self.zero_legacy_stage1(): - raise Exception( - "The deprecated version of ZeRO Stage 1 is not supported in deepspeed >= 0.5.9. Please downgrade to a version less than 0.5.9 if you need to use this deprecated version of ZeRO." - ) - - if zero_stage <= ZeroStageEnum.gradients: - overlap_comm = self.zero_overlap_comm() - contiguous_gradients = self.zero_contiguous_gradients() - round_robin_gradients = self.zero_round_robin_gradients() - assert not isinstance(optimizer, DummyOptim), "zero stage {} requires an optimizer".format(zero_stage) - - log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) - # Overlap and contiguous grads are meaningless in stage 1 and are ignored - if zero_stage == ZeroStageEnum.optimizer_states: - overlap_comm = False - round_robin_gradients = False - # Non-MoE requires contiguous grads to be disabled w. stage 1 - if not self.has_moe_layers: - contiguous_gradients = False - - if isinstance(self.module, PipelineModule): - if overlap_comm: - logger.warning("Pipeline parallelism does not support overlapped communication, will be disabled.") - overlap_comm = False - optimizer = DeepSpeedZeroOptimizer( - optimizer, - self.param_names, - timers=timers, - 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=contiguous_gradients, - reduce_bucket_size=self.zero_reduce_bucket_size(), - use_multi_rank_bucket_allreduce=self.zero_multi_rank_bucket_allreduce(), - allgather_bucket_size=self.zero_allgather_bucket_size(), - dp_process_group=self.seq_data_parallel_group, - expert_parallel_group=self.expert_parallel_group if self.has_moe_layers else None, - expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None, - reduce_scatter=self.zero_reduce_scatter(), - overlap_comm=overlap_comm, - offload_optimizer_config=self.zero_offload_optimizer(), - mpu=self.mpu, - postscale_gradients=self.postscale_gradients(), - gradient_predivide_factor=self.gradient_predivide_factor(), - gradient_accumulation_steps=self.gradient_accumulation_steps(), - ignore_unused_parameters=self.zero_ignore_unused_parameters(), - partition_grads=zero_stage == ZeroStageEnum.gradients, - round_robin_gradients=round_robin_gradients, - has_moe_layers=self.has_moe_layers, - fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(), - gradient_accumulation_dtype=gradient_accumulation_dtype, - communication_data_type=self.communication_data_type, - elastic_checkpoint=self.zero_elastic_checkpoint()) - - elif zero_stage == ZeroStageEnum.weights: - assert not self.has_moe_layers, "MoE not supported with Stage 3" - if isinstance(optimizer, DummyOptim): - log_dist("Creating ZeRO Offload", ranks=[0]) - zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() - if self.zero_hpz_partition_size() > 1 and zero_param_parallel_group is None: - self._set_zero_group_parallelism() - zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() - optimizer = DeepSpeedZeRoOffload( - self.module, - timers=timers, - ds_config=self.config, - overlap_comm=self.zero_overlap_comm(), - 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(), - offload_param_config=self.zero_offload_param(), - mpu=self.mpu, - zero_param_parallel_group=zero_param_parallel_group, - zero_quantized_weights=self.zero_quantized_weights(), - zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), - ) - else: - log_dist( - f'Creating fp16 ZeRO stage {zero_stage} optimizer,' - f' MiCS is enabled {mics_shard_size>0},' - f' Hierarchical params gather {self._config.mics_hierarchial_params_gather}', - ranks=[0]) - if mics_shard_size > 0: - return self._return_mics_optimizer(optimizer, timers) - - log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) - from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 - optimizer = DeepSpeedZeroOptimizer_Stage3( - self.module, - 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, - all2all_process_group=self.local_all_to_all_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(), - offload_ratio=self.zero_partial_offload(), - 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, - zero_hpz_partition_size=self.zero_hpz_partition_size(), - zero_quantized_weights=self.zero_quantized_weights(), - zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), - ) - - else: - raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage)) - - return optimizer - - def _return_mics_optimizer(self, basic_optimizer, timers): - from deepspeed.runtime.zero.mics import MiCS_Optimizer - model_dtype, gradient_accumulation_dtype = self.get_data_types() - optimizer = MiCS_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(), - max_iter=self.eigenvalue_max_iter(), - tol=self.eigenvalue_tol(), - stability=self.eigenvalue_stability(), - gas_boundary_resolution=self.eigenvalue_gas_boundary_resolution(), - layer_name=self.eigenvalue_layer_name(), - layer_num=self.eigenvalue_layer_num(), - ) - - return eigenvalue - - def _configure_progressive_layer_drop(self): - pld = ProgressiveLayerDrop(theta=self.pld_theta(), gamma=self.pld_gamma()) - - return pld - - def _configure_curriculum_scheduler_legacy(self): - scheduler = CurriculumScheduler(self.curriculum_params_legacy()) - return scheduler - - @staticmethod - def is_map_style_dataset(obj): - return hasattr(obj, "__getitem__") and hasattr(obj, "__len__") - - @staticmethod - def is_iterable_style_dataset(obj): - return isinstance(obj, torch.utils.data.IterableDataset) # hasattr(obj, "__iter__") should work as well - - def dataloader_drop_last(self): - return self._config.dataloader_drop_last - - def was_step_applied(self) -> bool: - """Returns True if the latest ``step()`` produced in parameter updates. - Note that a ``False`` return is not an error condition. Steps are frequently - no-ops, such as between gradient accumulation boundaries or when overflows - occur. - Returns: - bool: Whether the latest ``step()`` modified model parameters. - """ - return self._step_applied - - def deepspeed_io(self, - dataset, - batch_size=None, - route=ROUTE_TRAIN, - pin_memory=True, - data_sampler=None, - collate_fn=None, - num_local_io_workers=None): - if not (self.is_map_style_dataset(dataset) or self.is_iterable_style_dataset(dataset)): - raise ValueError("Training data must be a torch Dataset") - - if batch_size is None: - batch_size = self.train_micro_batch_size_per_gpu() - - if collate_fn is None: - collate_fn = self.collate_fn - - # Currently we only use timer in train route - deepspeed_io_timer = None - if route == ROUTE_TRAIN: - deepspeed_io_timer = self.tput_timer - - # If mpu is provided, forward world size and parallel rank to sampler. - data_parallel_world_size = self.dp_world_size - data_parallel_rank = self.global_rank - if self.mpu is not None: - data_parallel_world_size = self.mpu.get_data_parallel_world_size() - data_parallel_rank = self.mpu.get_data_parallel_rank() - - if data_sampler is None and (route == ROUTE_PREDICT or route == ROUTE_EVAL): - data_sampler = torch.utils.data.DistributedSampler( - dataset, - num_replicas=data_parallel_world_size, - rank=data_parallel_rank, - shuffle=False, - ) - - deepspeed_dataloader_config = {} - if self.curriculum_learning_enabled(): - deepspeed_dataloader_config = { - CURRICULUM_LEARNING: self.curriculum_learning_enabled(), - DATA_EFFICIENCY: self.data_efficiency_config(), - DATA_PARALLEL_GROUP: self.data_parallel_group, - GRADIENT_ACCUMULATION_STEPS: self.gradient_accumulation_steps(), - GLOBAL_RANK: self.global_rank, - DATA_SAMPLING_NUM_WORKERS: self.data_sampling_config()[DATA_SAMPLING_NUM_WORKERS] - } - - return DeepSpeedDataLoader(dataset=dataset, - batch_size=batch_size, - pin_memory=pin_memory, - collate_fn=collate_fn, - local_rank=self.local_rank, - tput_timer=deepspeed_io_timer, - num_local_io_workers=num_local_io_workers, - data_sampler=data_sampler, - data_parallel_world_size=data_parallel_world_size, - data_parallel_rank=data_parallel_rank, - dataloader_drop_last=self.dataloader_drop_last(), - deepspeed_dataloader_config=deepspeed_dataloader_config) - - def train(self, mode=True): - r"""""" - - self.warn_unscaled_loss = True - self.module.train(mode) - - def eval(self): - r"""""" - - self.warn_unscaled_loss = True - self.module.train(False) - - def _scale_loss_by_gas(self, prescaled_loss): - if isinstance(prescaled_loss, torch.Tensor): - scaled_loss = prescaled_loss / self.gradient_accumulation_steps() - elif isinstance(prescaled_loss, tuple) or isinstance(prescaled_loss, list): - scaled_loss = [] - for l in prescaled_loss: - if isinstance(l, torch.Tensor): - scaled_loss.append(l / self.gradient_accumulation_steps()) - else: - scaled_loss.append(l) - else: - scaled_loss = prescaled_loss - if self.warn_unscaled_loss: - logger.warning(f"DeepSpeed unable to scale loss because of type: {type(prescaled_loss)}") - self.warn_unscaled_loss = False - - return scaled_loss - - @instrument_w_nvtx - def forward(self, *inputs, **kwargs): - r"""Execute forward propagation - Arguments: - *inputs: Variable length input list - **kwargs: variable length keyword arguments - """ - - if self.autotuning_profile_model_info(): - ma = get_ma_status() - else: - see_memory_usage("Engine before forward", force=self.memory_breakdown()) - - flops_profiler_active = (self.flops_profiler_enabled() - and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) - - # used to check quantization happens at step 0! - if self.global_steps == 0 and hasattr(self, "compression_scheduler"): - self.compression_scheduler.step(step_zero_check=True) - if self.quantizer: - tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( - ) == 2 else self.optimizer.fp16_groups - if self.compression_scheduler.weight_quantization_enabled: - self.quantizer.quantize( - tensor_to_quantize, - (self.optimizer.overflow if self.fp16_enabled() else False), - self.eigenvalue_enabled(), - None, - ) - - if flops_profiler_active: - self.flops_profiler.start_profile(ignore_list=None) - - if self.module.training: - if self.progressive_layer_drop: - kwargs.update(self.progressive_layer_drop.get_state()) - - if self.__class__.__name__ != "PipelineEngine": - # TODO: The above if condition is a HACK since for PipelineEngine - # it's difficult to inject argument in forward pass. - if self.module.training and self.curriculum_enabled_legacy(): - self.curriculum_scheduler_legacy.update_difficulty(self.global_steps + 1) - if self.curriculum_params_legacy()["curriculum_type"] == "seqlen": - kwargs.update({"curriculum_seqlen": self.curriculum_scheduler_legacy.get_current_difficulty()}) - - if self.module.training and self.random_ltd_enabled(): - self.random_ltd_scheduler.update_seq(self.global_steps) - - if self.zero_optimization_partition_weights(): - # Enable automated discovery of external parameters by indicating that - # we are in a forward pass. - for module in self.module.modules(): - module._parameters._in_forward = True - - self._start_timers(self.engine_timers.forward_timers) - - if self.training_dataloader is None: - self.tput_timer.start() - - if self.fp16_auto_cast(): - inputs = self._cast_inputs_half(inputs) - # print(f"RANK[{self.global_rank}] self.fp16_auto_cast() is {self.fp16_auto_cast()}") - - loss = self.module(*inputs, **kwargs) - - # print(f"RANK[{self.global_rank}]'s loss is {loss}") - - if self.zero_optimization_partition_weights(): - # Disable automated discovery of external parameters - for module in self.module.modules(): - module._parameters._in_forward = False - - self._stop_timers(self.engine_timers.forward_timers) - - if flops_profiler_active: - self.flops_profiler.stop_profile() - - if self.autotuning_profile_model_info(): - activation_mem = get_ma_status() - ma - self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem - print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path()) - exit() - else: - see_memory_usage("Engine after forward", force=self.memory_breakdown()) - return loss - - def _cast_inputs_half(self, inputs): - if isinstance(inputs, (list, tuple)): - new_inputs = [] - for v in inputs: - new_inputs.append(self._cast_inputs_half(v)) - return inputs.__class__(new_inputs) - elif isinstance(inputs, dict): - new_inputs = {} - for k, v in inputs.items(): - new_inputs[k] = self._cast_inputs_half(v) - return new_inputs - elif hasattr(inputs, 'half'): - return inputs.half() - else: - return inputs - - def print_forward_breakdown(self, fwd_time): - gate_time = 0.0 - moe_time = 0.0 - falltoall = 0.0 - salltoall = 0.0 - - for gate in self.gate_modules: - #logger.info(f"Individual TopK gate time: {gate.gate_time:.2f} ms") - gate_time += gate.gate_time - - for l in self.moe_layers: - #logger.info(f"MoE layer; total: {l.time_moe:.2f} ms, first alltoall: {l.time_falltoall:.2f}, second alltoall: {l.time_salltoall:.2f}") - moe_time += l.time_moe - falltoall += l.time_falltoall - salltoall += l.time_salltoall - - # TODO: Allreduce/average them across ranks for more accurate timing. - - # if deepspeed.comm.get_rank() == 0: - log_dist( - f"time (ms) | fwd: {fwd_time:.2f} (fwd_moe: {moe_time:.2f}, 1st_a2a: {falltoall:.2f}, 2nd_a2a: {salltoall:.2f}, top_k: {gate_time:.2f})", - ranks=[0]) - - @instrument_w_nvtx - def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): - assert not (self.bfloat16_enabled() and self.pipeline_parallelism), \ - f'allreduce_gradients() is not valid when bfloat+pipeline_parallelism is enabled' - - # Pass (PP) gas boundary flag to optimizer (required for zero) - self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() - # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well - if self.zero_optimization_partition_gradients(): - self.optimizer.overlapping_partition_gradients_reduce_epilogue() - - # Communicate only at gradient accumulation boundaries - elif self.is_gradient_accumulation_boundary(): - if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr( - self.optimizer, 'reduce_gradients'): - self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism) - else: - self.buffered_allreduce_fallback(elements_per_buffer=bucket_size) - - @instrument_w_nvtx - def backward(self, loss, allreduce_gradients=True, release_loss=False, retain_graph=False, scale_wrt_gas=True): - r"""Execute backward pass on the loss - Arguments: - loss: Torch tensor on which to execute backward propagation - allreduce_gradients: is deprecated, ignored, and will soon be removed' - retain_graph: bool, default: false - forward on user defined choice of retain_graph - """ - - see_memory_usage("Engine before backward", force=self.memory_breakdown()) - - if self.scale_wrt_gas is not None: - scale_wrt_gas = self.scale_wrt_gas - - if not allreduce_gradients: - logger.warning(f"Argument `allreduce_gradients` is deprecated, ignored, and will soon be removed") - - # scale loss w.r.t. gradient accumulation if needed - if self.gradient_accumulation_steps() > 1 and scale_wrt_gas: - loss = self._scale_loss_by_gas(loss.float()) - - # Log training loss - self.losses += loss.mean().item() - if self.monitor.enabled: - if self.is_gradient_accumulation_boundary(): - if self.global_rank == 0: - self.summary_events = [( - f"Train/Samples/train_loss", - self.losses, - self.global_samples, - )] - self.monitor.write_events(self.summary_events) - - self._start_timers(self.engine_timers.backward_timers) - - assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ - "must provide optimizer during init in order to use backward" - - self._start_timers(self.engine_timers.backward_inner_timers) - - if self.zero_optimization(): - self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() - self.optimizer.backward(loss, retain_graph=retain_graph) - elif self.amp_enabled(): - # AMP requires delaying unscale when inside gradient accumulation boundaries - # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations - delay_unscale = not self.is_gradient_accumulation_boundary() - with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss: - scaled_loss.backward(retain_graph=retain_graph) - elif self.fp16_enabled(): - if self.eigenvalue_enabled(): - self.optimizer.backward(loss, create_graph=True, retain_graph=True) - else: - self.optimizer.backward(loss, retain_graph=retain_graph) - elif self.bfloat16_enabled(): - self.optimizer.backward(loss) - else: - if self.eigenvalue_enabled(): - loss.backward(create_graph=True, retain_graph=True) - else: - loss.backward(retain_graph=retain_graph) - - self._stop_timers(self.engine_timers.backward_inner_timers) - - self._start_timers(self.engine_timers.backward_reduce_timers) - - if allreduce_gradients and self.enable_backward_allreduce: - # Traditional code path that allreduces the module parameter grads - self.allreduce_gradients() - - self._stop_timers(self.engine_timers.backward_reduce_timers) - - self._stop_timers(self.engine_timers.backward_timers) - - if release_loss: - # loss.data = None - pass - - see_memory_usage("Engine after backward", force=self.memory_breakdown()) - - return loss - - def is_gradient_accumulation_boundary(self): - """ - Query whether the current micro-batch is at the boundary of - gradient accumulation, and thus will trigger gradient reductions and - an optimizer step. - - Returns: - bool: if the current step is a gradient accumulation boundary. - - """ - if self._is_gradient_accumulation_boundary is None: - return (self.micro_steps + 1) % \ - self.gradient_accumulation_steps() == 0 - else: - return self._is_gradient_accumulation_boundary - - def set_gradient_accumulation_boundary(self, is_boundary): - """ - Manually overrides the DeepSpeed engine's gradient accumulation boundary state, this is an optional - feature and should be used with care. The state should be set before to the intended - value before each forward/backward. The final forward/backward should have the - boundary state set to True. This style allows client code to only call engine.step() once after all - the gradient accumulation passes are complete. See example below: - .. code-block:: python - engine.set_gradient_accumulation_boundary(False) - for _ in range(gradient_accumulation_steps - 1): - micro_batch = next(data_loader) - loss = engine(micro_batch) - engine.backward(loss) - engine.set_gradient_accumulation_boundary(True) - micro_batch = next(data_loader) - loss = engine(micro_batch) - engine.backward(loss) - engine.step() - Arguments: - is_boundary (bool): are we at a gradient accumulation boundary or not? - """ - self._is_gradient_accumulation_boundary = is_boundary - self.optimizer.is_gradient_accumulation_boundary = is_boundary - - def zero_grad(self): - """ - Zero parameter grads. - """ - for param_name, param in self.module.named_parameters(): - param.grad = None - - def clip_fp32_gradients(self): - clip_grad_norm_(parameters=self.module.parameters(), max_norm=self.gradient_clipping(), mpu=self.mpu) - - def _take_model_step(self, lr_kwargs, block_eigenvalue={}): - if self.gradient_clipping() > 0.0: - if not (self.fp16_enabled() or self.bfloat16_enabled() or self.amp_enabled() or self.zero_optimization()): - self.clip_fp32_gradients() - elif self.amp_enabled(): - # AMP's recommended way of doing clipping - # https://nvidia.github.io/apex/advanced.html#gradient-clipping - master_params = amp.master_params(self.optimizer) - clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clipping(), mpu=self.mpu) - self.optimizer.step() - - if hasattr(self.optimizer, '_global_grad_norm'): - self._global_grad_norm = self.optimizer._global_grad_norm - - # Quantize the updated parameter if there is no overflow - if self.quantizer: - tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( - ) == 2 else self.optimizer.fp16_groups - if self.compression_scheduler.weight_quantization_enabled: - self.quantizer.quantize( - tensor_to_quantize, - (self.optimizer.overflow if self.fp16_enabled() else False), - self.eigenvalue_enabled(), - block_eigenvalue, - ) - # zero grad in basic optimizer could be unreliable and may not exhibit - # the behavior that we want - if self.bfloat16_enabled(): - # TODO: Temporary until bf16_optimizer and zero_optimizer are integrated - if self.zero_optimization() and hasattr(self.optimizer, "zero_grad"): - self.optimizer.zero_grad() - else: - pass - elif self.zero_optimization() or self.fp16_enabled() or self.amp_enabled(): - self.optimizer.zero_grad() - else: - self.zero_grad() - - report_progress = self.global_rank == 0 if self.global_rank else True - - # Check overflow here since in DS fp16 optimizer, the overflow is updated in above step() function. - overflow = False - if hasattr(self.optimizer, "overflow"): - overflow = self.optimizer.overflow - self._step_applied = not overflow - - if overflow: - self.skipped_steps += 1 - else: - self.compression_scheduler.step() - if self.lr_scheduler is not None: - try: - self.lr_scheduler.step(**(lr_kwargs or {})) - except TypeError: - # XXX Hack to work with Megatron 2.0 and DeepSpeed pipelines. - # We don't currently have a way to specify lr_kwargs from - # pipe_engine.train_batch() - self.lr_scheduler.step(self.train_batch_size()) - - if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0: - self._report_progress(self.global_steps + 1) - - self.losses = 0.0 - self.global_steps += 1 - self.global_samples += self.train_batch_size() - - def step(self, lr_kwargs=None): - r"""Execute the weight update step after forward and backward propagation - on effective_train_batch. - """ - see_memory_usage("Engine before step", force=self.memory_breakdown()) - - # Check early because self.global_steps is incremented at some point here. - # TODO: Delay self.global_steps increment until very end of this function. - flops_profiler_active = self.flops_profiler_enabled( - ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 - - self._start_timers(self.engine_timers.step_timers) - - assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ - "must provide optimizer during init in order to use step" - - report_progress = False - - self._step_applied = False # assume False, will flip to True - - # Update the model when we reach gradient accumulation boundaries - if self.is_gradient_accumulation_boundary(): - self.gas_boundary_ctr += 1 - - if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) - and self.quantizer.any_precision_switch()): - log_dist(f"computing eigenvalue...", ranks=[0]) - self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, - self.optimizer.cur_scale) - - if self.progressive_layer_drop: - self.progressive_layer_drop.update_state(self.global_steps) - - if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() - and self.quantizer.any_precision_switch()): - self._take_model_step(lr_kwargs, self.block_eigenvalue) - else: - self._take_model_step(lr_kwargs) - - report_progress = self.global_rank == 0 if self.global_rank else True - - self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) - - self._stop_timers(self.engine_timers.step_timers) - - # Log learning rate - if self.monitor.enabled: - if self.is_gradient_accumulation_boundary(): - if self.global_rank == 0: - self.summary_events = [(f"Train/Samples/lr", self.get_lr()[0], self.global_samples)] - - if self.fp16_enabled() and hasattr(self.optimizer, "cur_scale"): - self.summary_events.append(( - f"Train/Samples/loss_scale", - self.optimizer.cur_scale, - self.global_samples, - )) - - if (self.eigenvalue_enabled() - and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()): - ev_values = self.block_eigenvalue.values() - for i in range(len(ev_values)): - self.summary_events.append(( - f"Train/Eigenvalues/ModelBlockParam_{i}", - self.ev_values[i][0], - self.global_samples, - )) - self.monitor.write_events(self.summary_events) - - # Check flops profiling - if flops_profiler_active: - if self.autotuning_enabled(): - self.flops = self.flops_profiler.get_total_flops() * 3 - self.fwd_duration = self.flops_profiler.get_total_duration() - else: - self.flops_profiler.print_model_profile( - profile_step=self.global_steps, - module_depth=self.flops_profiler_module_depth(), - top_modules=self.flops_profiler_top_modules(), - detailed=self.flops_profiler_detailed(), - output_file=self.flops_profiler_output_file(), - ) - self.flops_profiler.end_profile() - - if self.autotuning_enabled() and self.global_steps == (self.autotuning_end_profile_step() + 1): - self._autotuning_exit() - - if self.wall_clock_breakdown(): - # Log micro timing and reset - self.timers.log(names=self.engine_timers.micro_timers, memory_breakdown=self.memory_breakdown()) - - if self.wall_clock_breakdown() or self.flops_profiler_enabled(): - # Log global timing and reset - if self.is_gradient_accumulation_boundary(): - if self.monitor.enabled: - self._write_monitor() - - if self.has_moe_layers: - fwd_time = self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False) - self.print_forward_breakdown(fwd_time=fwd_time) - - self.timers.log(self.engine_timers.global_timers) - - self.micro_steps += 1 - see_memory_usage("Engine after step", force=self.memory_breakdown()) - - def _start_timers(self, timer_names): - for name in timer_names: - self.timers(name).start() - - def _stop_timers(self, timer_names): - record = self.is_gradient_accumulation_boundary() and \ - self.flops_profiler_enabled() and \ - (self.global_steps >= self.flops_profiler_profile_step()) - for name in timer_names: - self.timers(name).stop(record=record) - - def _autotuning_exit(self): - if self.global_rank == 0: - msg = self.timers.get_mean([ - FORWARD_GLOBAL_TIMER, - BACKWARD_GLOBAL_TIMER, - STEP_GLOBAL_TIMER, - ], reset=False) - titer = 0.0 - titer += msg[FORWARD_GLOBAL_TIMER] if FORWARD_GLOBAL_TIMER in msg else 0 - titer += msg[BACKWARD_GLOBAL_TIMER] if BACKWARD_GLOBAL_TIMER in msg else 0 - titer += msg[STEP_GLOBAL_TIMER] if STEP_GLOBAL_TIMER in msg else 0 - titer *= self.gradient_accumulation_steps() - msg["latency"] = titer - msg["FLOPS_per_gpu"] = self.flops * 1_000_000 * self.gradient_accumulation_steps() / titer - msg["throughput"] = self.train_batch_size() * 1_000_000 / \ - msg["latency"] - print_json_dist(msg, [0], path=self.autotuning_metric_path()) - log_dist( - f"Wrote metrics to {self.autotuning_metric_path()}, {os.path.abspath(self.autotuning_metric_path())}", - ranks=[0]) - import atexit - atexit.register(print, "Autotuning: done with running current ds config.") - exit() - - def _write_monitor(self): - if self.global_rank == 0: - self.summary_events = [ - ( - f"Train/Samples/elapsed_time_ms_forward", - self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_backward", - self.timers(BACKWARD_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_backward_inner", - self.timers(BACKWARD_INNER_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_backward_allreduce", - self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_step", - self.timers(STEP_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ] - self.monitor.write_events(self.summary_events) - - def _get_optimizer_param(self, param_name): - result = [] - if not self.optimizer: - return result - for group in self.optimizer.param_groups: - if param_name in group: - result.append(group[param_name]) - else: - result.append(0.0) - return result - - def get_lr(self): - return self._get_optimizer_param("lr") - - def get_type(self): - return self._get_optimizer_param("type") - - def get_mom(self): - if self.optimizer_name() in ["SGD", "RMSprop"]: - return self._get_optimizer_param("momentum") - else: - return self._get_optimizer_param("betas") - - def get_pld_theta(self): - if self.progressive_layer_drop: - return self.progressive_layer_drop.get_theta() - else: - return None - - def _report_progress(self, step): - lr = self.get_lr() - mom = self.get_mom() - log_dist(f"step={step}, skipped={self.skipped_steps}, lr={lr}, mom={mom}", ranks=[0]) - - def allreduce_bucket(self, bucket, dp_group): - tensor = self.flatten(bucket) - - tensor_to_allreduce = tensor - - if self.communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(self.communication_data_type) - - if self.postscale_gradients(): - if self.gradient_predivide_factor() != 1.0: - tensor_to_allreduce.mul_(1.0 / self.gradient_predivide_factor()) - - dist.all_reduce(tensor_to_allreduce, group=dp_group) - if self.gradient_average: - if self.gradient_predivide_factor() != dist.get_world_size(group=dp_group): - tensor_to_allreduce.mul_(self.gradient_predivide_factor() / dist.get_world_size(group=dp_group)) - else: - tensor_to_allreduce.mul_(1. / dist.get_world_size(group=dp_group)) - dist.all_reduce(tensor_to_allreduce, group=dp_group) - - if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - tensor.copy_(tensor_to_allreduce) - - return tensor - - def allreduce_and_copy(self, small_bucket, dp_group): - allreduced = self.allreduce_bucket(small_bucket, dp_group) - for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): - buf.copy_(synced) - - def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000): - small_bucket = [] - numel = 0 - for tensor in bucket: - small_bucket.append(tensor) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy(small_bucket, dp_group) - small_bucket = [] - numel = 0 - if len(small_bucket) > 0: - self.allreduce_and_copy(small_bucket, dp_group) - - def _get_gradients_for_reduction(self): - non_expert_grads = [] - expert_grads = {} - if self.has_moe_layers: - for key in self.expert_data_parallel_group.keys(): - expert_grads[key] = [] - - for param_name, param in self.module.named_parameters(): - if not param.requires_grad: - continue - - if param.grad is None: - # In cases where there is an imbalance of empty grads across - # ranks we must create empty grads, this will ensure that every - # rank is reducing the same size. In some cases it may make - # sense in the future to support the ability to average not - # w.r.t. world size but with a different value. - param.grad = torch.zeros(param.size(), dtype=param.dtype, device=param.device) - - grad_data = param.grad.data - if param_name in self.sparse_tensor_module_names or grad_data.is_sparse: - # Call param.grad without data to avoid problem with setting of updated grads - grad_data = SparseTensor(param.grad) - - if is_moe_param(param): - expert_grads[param.group_name].append(grad_data) - else: - non_expert_grads.append(grad_data) - - return non_expert_grads, expert_grads - - def _reduce_non_expert_gradients(self, grads, elements_per_buffer): - split_buckets = split_half_float_double_sparse(grads) - for _, bucket_tuple in enumerate(split_buckets): - bucket_type, bucket = bucket_tuple - - if self.pipeline_parallelism: - dp_group = self.mpu.get_data_parallel_group() - else: - dp_group = groups._get_sequence_data_parallel_group() - - if bucket_type == SparseTensor.type(): - self.sparse_allreduce_no_retain(bucket, dp_group=dp_group) - else: - self.allreduce_no_retain(bucket, dp_group=dp_group, numel_per_bucket=elements_per_buffer) - - def _reduce_expert_gradients(self, expert_grads, elements_per_buffer): - for ep_name, expert_grads_group in expert_grads.items(): - expert_split_buckets = split_half_float_double_sparse(expert_grads_group) - for i, bucket_tuple in enumerate(expert_split_buckets): - bucket_type, bucket = bucket_tuple - if bucket_type == SparseTensor.type(): - self.sparse_allreduce_no_retain(bucket, groups._get_expert_data_parallel_group(ep_name)) - else: - # Separate between diff groups - self.allreduce_no_retain(bucket, - dp_group=groups._get_expert_data_parallel_group(ep_name), - numel_per_bucket=elements_per_buffer) - - def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000): - if grads is None: - non_expert_grads, expert_grads = self._get_gradients_for_reduction() - else: - assert not self.has_moe_layers, "attempting to reduce grads in unsupported way w.r.t. MoE" - non_expert_grads = grads - - self._reduce_non_expert_gradients(non_expert_grads, elements_per_buffer) - - if self.has_moe_layers: - self._reduce_expert_gradients(expert_grads, elements_per_buffer) - - def sparse_allreduce_no_retain(self, bucket, dp_group): - allreduced_sparses = self.sparse_allreduce_bucket(bucket, dp_group) - # Densify sparse tensor and copy back to original location - for tensor in allreduced_sparses: - if tensor.is_sparse: - tensor.orig_dense_tensor.data = tensor.to_coo_tensor() - else: - tensor.orig_dense_tensor.copy_(tensor.to_dense()) - - def sparse_allreduce_bucket(self, bucket, dp_group): - sparse_list = [] - for sparse in bucket: - sparse_list.append(self.sparse_allreduce(sparse, dp_group)) - return sparse_list - - def sparse_allreduce(self, sparse, dp_group): - original_data_type = sparse.values.dtype - if self.communication_data_type != sparse.values.dtype: - if self.communication_data_type in (torch.float16, torch.bfloat16): - indices = sparse.indices.to(torch.int32) - else: - indices = sparse.indices - values = sparse.values.to(self.communication_data_type) - else: - indices = sparse.indices - values = sparse.values - - if self.postscale_gradients(): - if self.gradient_average: - values.mul_(self.gradient_predivide_factor() / - (dist.get_world_size(group=dp_group) / float(self.sequence_parallel_size))) - else: - values.mul_(1. / (dist.get_world_size(group=dp_group) / float(self.sequence_parallel_size))) - - indices_device_list = self.sparse_all_gather(indices, dp_group) - values_device_list = self.sparse_all_gather(values, dp_group) - - sparse.indices = torch.cat(indices_device_list).to(torch.long) - sparse.values = torch.cat(values_device_list).to(original_data_type) - return sparse - - def sparse_all_gather(self, value, dp_group): - my_size = torch.LongTensor([value.size()[0]]).to(self.device) - all_sizes = self.all_gather_scalar(my_size, dp_group) - max_size = torch.cat(all_sizes).max() - fill_size = max_size - my_size - - assert value.dim() in [1, 2] - if value.dim() == 1: - if fill_size > 0: - value = torch.cat([value, value.new_empty(fill_size)]) - tensor_list = [value.new_empty(max_size) for _ in range(dist.get_world_size(group=dp_group))] - else: - if fill_size > 0: - value = torch.cat([value, value.new_empty(fill_size, value.size()[1])]) - tensor_list = [ - value.new_empty(max_size, - value.size()[1]) for _ in range(dist.get_world_size(group=dp_group)) - ] - - dist.all_gather(tensor_list, value, group=dp_group) - tensors = [] - for dev_idx, t in enumerate(tensor_list): - size = all_sizes[dev_idx][0] - tensors.append(t.index_select(0, torch.arange(size, dtype=torch.long, device=self.device))) - - return tensors - - def all_gather_scalar(self, value, dp_group): - tensor_list = [value.new_zeros(value.size()) for _ in range(dist.get_world_size(group=dp_group))] - dist.all_gather(tensor_list, value, group=dp_group) - return tensor_list - - def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclude_frozen_parameters=False): - sd = self.module.state_dict(destination, prefix, keep_vars) - - # Remove frozen parameter weights from state_dict if specified - if exclude_frozen_parameters: - for n, p in self.module.named_parameters(): - if not p.requires_grad and n in sd: - del sd[n] - - if self.random_ltd_enabled(): - sd = remove_random_ltd_state_dict(sd) - return sd - - @staticmethod - def load_moe_state_dict(checkpoint_path, - tag, - state_dict, - old_moe_load, - model=None, - mpu=None, - num_experts=1, - checkpoint_engine=TorchCheckpointEngine()): - if old_moe_load: - expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name()) - - num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size( - groups._get_max_expert_size_name()) - for local_expert_id in range(num_local_experts): - global_expert_id = expp_rank * num_local_experts + local_expert_id - expert_state_dict = checkpoint_engine.load( - DeepSpeedEngine._get_expert_ckpt_name( - checkpoint_path, - -1, # -1 means ignore layer_id - global_expert_id, - tag, - mpu), - map_location=torch.device('cpu')) - - # Updating global -> local expert ids - moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' - for key in list(expert_state_dict.keys()): - local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', - f'{moe_str_prefix}{local_expert_id}') - expert_state_dict[local_key] = expert_state_dict.pop(key) - state_dict.update(expert_state_dict) - - else: - moe_layer_id = 0 - for n_module, module in model.named_modules(): - if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: - group_name = module.expert_group_name - num_local_experts = module.num_local_experts - expp_rank = groups._get_expert_parallel_rank(group_name) - # loop all local_experts - for local_expert_id in range(num_local_experts): - global_expert_id = expp_rank * num_local_experts + local_expert_id - expert_state_dict = checkpoint_engine.load(DeepSpeedEngine._get_expert_ckpt_name( - checkpoint_path, moe_layer_id, global_expert_id, tag, mpu), - map_location=torch.device('cpu')) - # print(expert_state_dict.keys()) - # Updating global -> local expert ids - moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' - for key in list(expert_state_dict.keys()): - local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', - f'{moe_str_prefix}{local_expert_id}') - expert_state_dict[local_key] = expert_state_dict.pop(key) - state_dict.update(expert_state_dict) - moe_layer_id += 1 - - def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): - if fetch_z3_params: - params_to_fetch = [ - p for p in self.module.parameters() - if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE - ] - else: - params_to_fetch = [] - - with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0): - module_state_dict = checkpoint['module'] - if custom_load_fn: - custom_load_fn(src=module_state_dict, dst=self.module) - else: - self.module.load_state_dict( - module_state_dict, # TODO - strict=strict) - - if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: - saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] - for param in self.module.parameters(): - if param.requires_grad: - continue - if param not in self.param_names: - raise ValueError(f"failed to find frozen {param} in named params") - name = self.param_names[param] - if hasattr(param, 'ds_id'): - param.ds_tensor.data.copy_(saved_frozen_params[name].data) - else: - param.data.copy_(saved_frozen_params[name].data) - - def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode): - return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}' - - def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode): - file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode) - zero_ckpt_name = os.path.join( - checkpoints_path, - str(tag), - f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt", - ) - return zero_ckpt_name - - def _get_zero_ckpt_name(self, checkpoints_path, tag): - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - pp_rank = dist.get_rank(group=self.optimizer.zp_process_group) - bf16_mode = self.bfloat16_enabled() - return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode) - - def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None): - if mp_placeholder is not None: - mp_rank_str = mp_placeholder - else: - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - mp_rank_str = f"{mp_rank:02d}" - - if self.zero_optimization_partition_weights(): - filename = "zero_pp_rank_{}".format(dist.get_rank(group=self.optimizer.zp_process_group)) - ckpt_name = os.path.join( - checkpoints_path, - str(tag), - f"{filename}_mp_rank_{mp_rank_str}_model_states.pt", - ) - else: - ckpt_name = os.path.join( - checkpoints_path, - str(tag), - "mp_rank_" + mp_rank_str + "_model_states.pt", - ) - return ckpt_name - - def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank): - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - ckpt_name = os.path.join(checkpoints_path, str(tag), - f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt') - return ckpt_name - - @staticmethod - def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None): - mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank() - if layer_id <= -1: - # Used to support old checkpoint loading - ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), - f'expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') - else: - # Used to support new checkpoint loading - ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), - f'layer_{layer_id}_expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') - return ckpt_name - - def _get_all_ckpt_names(self, checkpoints_path, tag): - # It is required that (checkpoints_path, tag) are consistent among all ranks. - ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*") - import glob - - ckpt_files = glob.glob(ckpt_file_pattern) - ckpt_files.sort() - return ckpt_files - - def load_checkpoint(self, - load_dir, - tag=None, - load_module_strict=True, - load_optimizer_states=True, - load_lr_scheduler_states=True, - load_module_only=False, - custom_load_fn=None): - """ - Load training checkpoint - - Arguments: - load_dir: Required. Directory to load the checkpoint from - tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file - load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. - load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance - load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. - load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. - custom_load_fn: Optional. Custom model load function. - - Returns: - A tuple of ``load_path`` and ``client_state``. - *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. - *``client_state``: State dictionary used for loading required training states in the client code. - - Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right - after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and - ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine - before ``load_checkpoint()``. - - """ - - if tag is None: - latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" - latest_path = os.path.join(load_dir, latest_tag) - if os.path.isfile(latest_path): - with open(latest_path, "r") as fd: - tag = fd.read().strip() - else: - if self.load_universal_checkpoint(): - raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') - else: - logger.warning( - f"Unable to find latest file at {latest_path}, if trying to load latest " - "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." - ) - return None, None - - if self._optimizer_has_ckpt_event_prologue(): - # Prepare for checkpoint load by ensuring all parameters are partitioned - self.optimizer.checkpoint_event_prologue() - - load_path, client_states = self._load_checkpoint(load_dir, - tag, - load_module_strict=load_module_strict, - load_optimizer_states=load_optimizer_states, - load_lr_scheduler_states=load_lr_scheduler_states, - load_module_only=load_module_only, - custom_load_fn=custom_load_fn) - - load_zero_checkpoint = load_optimizer_states and load_path is not None and (self.zero_optimization() - or self.bfloat16_enabled()) - if load_zero_checkpoint: - success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states) - if not success: - self.optimizer._restore_from_bit16_weights() - - if self._optimizer_has_ckpt_event_epilogue(): - self.optimizer.checkpoint_event_epilogue() - - if self.load_universal_checkpoint(): - self.optimizer.update_lp_params() - if load_zero_checkpoint: - self.update_optimizer_step(step=client_states['iteration'] + 1) - - return load_path, client_states - - def _load_checkpoint(self, - load_dir, - tag, - load_module_strict=True, - load_optimizer_states=True, - load_lr_scheduler_states=True, - load_module_only=False, - custom_load_fn=None): - - from deepspeed.runtime.state_dict_factory import SDLoaderFactory - - ckpt_list = self._get_all_ckpt_names(load_dir, tag) - sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine) - - is_pipe_parallel = isinstance(self.module, PipelineModule) - - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel) - - if checkpoint is None: - return None, None - - fetch_z3_params = False - if self.zero_optimization_partition_weights() and not load_optimizer_states: - checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) - fetch_z3_params = True - - if is_pipe_parallel: - # Pipeline parallelism uses this to load its own checkpoint files. - self._curr_ckpt_path = os.path.join(load_dir, tag) - - if self.has_moe_layers: - # print(checkpoint.keys()) - old_moe_load = False - if not isinstance(checkpoint['num_experts'], list): - old_moe_load = True - DeepSpeedEngine.load_moe_state_dict(load_dir, - tag, - state_dict=checkpoint['module'], - old_moe_load=old_moe_load, - model=self.module, - mpu=self.mpu, - num_experts=self.num_experts, - checkpoint_engine=self.checkpoint_engine) - if not self.load_universal_checkpoint(): - self.load_module_state_dict(checkpoint=checkpoint, - strict=load_module_strict, - custom_load_fn=custom_load_fn, - fetch_z3_params=fetch_z3_params) - - self.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size'] - if 'zp_world_size' not in checkpoint: - checkpoint['zp_world_size'] = self.zp_world_size - self.loaded_checkpoint_zp_world_size = checkpoint['zp_world_size'] - - optim_checkpoint = None - if load_module_only: - deepspeed_states = ['module'] - if self.optimizer is not None and self.fp16_enabled(): - self.optimizer.refresh_fp32_params() - else: - has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() - if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state: - if self.has_moe_layers: - largest_group_name = groups._get_max_expert_size_name() - expp_rank = groups._get_expert_parallel_rank(largest_group_name) - optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank) - optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu')) - else: - optim_checkpoint = checkpoint - - if self.fp16_enabled() or self.bfloat16_enabled(): - self.optimizer.load_state_dict(optim_checkpoint['optimizer'], - load_optimizer_states=load_optimizer_states) - else: - optim_checkpoint = checkpoint - - self.optimizer.load_state_dict(optim_checkpoint['optimizer']) - - if load_lr_scheduler_states and self.lr_scheduler is not None: - self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) - - if self.random_ltd_enabled() and self.random_ltd_scheduler is not None and 'random_ltd' in checkpoint: - self.random_ltd_scheduler.load_state_dict(checkpoint['random_ltd']) - - if self.training_dataloader is not None and self.curriculum_learning_enabled( - ) and 'data_sampler' in checkpoint: - self.training_dataloader.data_sampler.load_state_dict(checkpoint['data_sampler']) - - def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters, loaded_parameters): - result = set() - - for name in original_set: - if name in loaded_parameters and name not in loaded_set: - continue # parameter existed in previous model and was not sparse - result.add(name) - - for name in loaded_set: - if name in original_parameters: - result.add(name) # parameter exists in both configs and it was sparse - - return result - - if 'sparse_tensor_module_names' in checkpoint: - sparse_tensor_module_names = checkpoint['sparse_tensor_module_names'] - elif 'csr_tensor_module_names' in checkpoint: - sparse_tensor_module_names = checkpoint['csr_tensor_module_names'] - else: - sparse_tensor_module_names = None - if sparse_tensor_module_names is not None: - if load_module_strict: - self.sparse_tensor_module_names = sparse_tensor_module_names - else: - self.sparse_tensor_module_names = get_sparse_tensor_module_names( - self.sparse_tensor_module_names, sparse_tensor_module_names, - dict(self.module.named_parameters()), checkpoint["module"]) - - self.global_steps = checkpoint['global_steps'] - self.global_samples = checkpoint.get('global_samples', self.global_steps * self.train_batch_size()) - self.skipped_steps = checkpoint['skipped_steps'] - self.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size'] - deepspeed_states = [ - 'module', 'sparse_tensor_module_names', 'skipped_steps', 'global_steps', 'zp_world_size', - 'mp_world_size', 'data_sampler', 'random_ltd', 'dp_world_size', - ] - client_state = {} - - if load_lr_scheduler_states: - deepspeed_states.append('lr_scheduler') - if load_optimizer_states: - deepspeed_states.append('optimizer') - - client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states} - - if optim_checkpoint is not None: - client_state['optimizer'] = optim_checkpoint['optimizer'] - - return load_path, client_state - - def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): - - load_serial = None - # When use loading checkpoint serial, checkpoint loading start from local rank 0, - # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. - if self._config.zero_config.pipeline_loading_checkpoint: - assert self.zero_optimization_stage( - ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" - load_serial = torch.zeros(1).to(self.device) - if dist.get_local_rank() != 0: - dist.recv(tensor=load_serial, src=dist.get_rank() - 1) - if self.load_universal_checkpoint(): - zero_sd_list = None - checkpoint_folder = f'{os.path.join(load_dir, tag)}' - else: - if load_optimizer_states and self.zp_world_size != self.loaded_checkpoint_zp_world_size: - raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ - f"world size of {self.loaded_checkpoint_zp_world_size} but the " \ - f"current world size is {self.zp_world_size}. Automatic adjustment " \ - "of ZeRO's optimizer state partitioning with a new world size is not " \ - "currently supported.") - checkpoint_folder = None - zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) - if zero_sd_list is None: - return False - - self.optimizer.load_state_dict(state_dict_list=zero_sd_list, - load_optimizer_states=load_optimizer_states, - load_from_fp32_weights=self.zero_load_from_fp32_weights(), - checkpoint_folder=checkpoint_folder, - load_serial=load_serial) - - if self.load_universal_checkpoint(): - logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') - else: - logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") - return True - - def update_optimizer_step(self, step): - - def set_step(d): - if isinstance(d['step'], torch.Tensor): - d['step'] = torch.tensor(step, dtype=d['step'].dtype, device=d['step'].device) - else: - d['step'] = step - - optimizer = self.optimizer - base_optimizer = optimizer.optimizer - state = base_optimizer.state - for group in optimizer.param_groups: - if 'step' in group: - set_step(group) - for p in group['params']: - if p in state and len(state[p]) > 0 and 'step' in state[p]: - set_step(state[p]) - - def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): - zero_ckpt_names = [] - for dp_rank in range(dp_world_size): - ckpt_name = self._get_rank_zero_ckpt_name(checkpoints_path=load_dir, - tag=tag, - mp_rank=mp_rank, - dp_rank=dp_rank, - bf16_mode=bf16_mode) - zero_ckpt_names.append(ckpt_name) - - return zero_ckpt_names - - def _get_all_zero_checkpoint_names(self, load_dir, tag, bf16_mode): - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - zero_ckpt_names = self._get_mp_rank_zero_checkpoint_names(load_dir=load_dir, - tag=tag, - mp_rank=mp_rank, - dp_world_size=self.loaded_checkpoint_dp_world_size, - bf16_mode=bf16_mode) - for i, ckpt_name in enumerate(zero_ckpt_names): - if not os.path.exists(ckpt_name): - # transparently handle the old file pattern for optim_states - if "optim_states.pt" in ckpt_name: - ckpt_name_try = ckpt_name.replace("_optim_states.pt", "optim_states.pt") - if os.path.exists(ckpt_name_try): - zero_ckpt_names[i] = ckpt_name_try - continue - - return zero_ckpt_names - - def _get_all_zero_checkpoint_state_dicts(self, zero_ckpt_names): - zero_sd_list = [] - for i, ckpt_name in enumerate(zero_ckpt_names): - _state = None - if ckpt_name is None: - _state = {OPTIMIZER_STATE_DICT: None} - # Fully load state for current rank - elif self.zero_elastic_checkpoint() or dist.get_rank(group=self.optimizer.zp_process_group) == i: - _state = self.checkpoint_engine.load( - ckpt_name, - map_location='cpu', - ) - else: - _state = {OPTIMIZER_STATE_DICT: None} - zero_sd_list.append(_state) - - zero_optimizer_sd = [sd[OPTIMIZER_STATE_DICT] for sd in zero_sd_list] - logger.info(f"successfully read {len(zero_optimizer_sd)} ZeRO state_dicts for rank {self.global_rank}") - return zero_optimizer_sd - - def _get_all_zero_checkpoints(self, load_dir, tag): - for bf16_mode in [self.bfloat16_enabled(), not self.bfloat16_enabled()]: - zero_ckpt_names = self._get_all_zero_checkpoint_names(load_dir, tag, bf16_mode) - if zero_ckpt_names is not None: - # Warn if loading checkpoint of different bit16 type - if bf16_mode is not self.bfloat16_enabled(): - checkpoint_bit16 = BFLOAT16 if bf16_mode else FP16 - engine_bit16 = BFLOAT16 if self.bfloat16_enabled() else FP16 - logger.warn(f'Loading {checkpoint_bit16} zero checkpoints into {engine_bit16} training engine') - return self._get_all_zero_checkpoint_state_dicts(zero_ckpt_names) - - return None - - def _checkpoint_tag_validation(self, tag): - if self.checkpoint_tag_validation_enabled(): - s_hash = hashlib.sha1(tag.encode()) - bhash = torch.ByteTensor([s_hash.digest()]).flatten().to(self.device) - max_bhash = bhash.clone() - min_bhash = bhash.clone() - dist.all_reduce(max_bhash, op=dist.ReduceOp.MAX) - dist.all_reduce(min_bhash, op=dist.ReduceOp.MIN) - valid = all(min_bhash == bhash) and all(max_bhash == bhash) - msg = (f"[rank={dist.get_rank()}] The checkpoint tag name '{tag}' is not consistent across " - "all ranks. Including rank unique information in checkpoint tag could cause issues when " - "restoring with different world sizes.") - if self.checkpoint_tag_validation_fail(): - assert valid, msg - elif not valid: - logger.warning(msg) - - def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False): - """Save training checkpoint - - Arguments: - save_dir: Required. Directory for saving the checkpoint - tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is - used if not provided. Tag name must be the same across all ranks. - client_state: Optional. State dictionary used for saving required training states in the client code. - save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint. - exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. - Important: all processes must call this method and not just the process with rank 0. It is - because each process needs to save its master weights and scheduler+optimizer states. This - method will hang waiting to synchronize with other processes if it's called just for the - process with rank 0. - - """ - if self._optimizer_has_ckpt_event_prologue(): - # Custom preparation for checkpoint save, if applicable - self.optimizer.checkpoint_event_prologue() - - rank = self.local_rank if self.use_node_local_storage() else self.global_rank - - # This is to make sure the checkpoint names are created without collision - # There seems to be issue creating them in parallel - - # Ensure save_dir directory exists - if rank == 0: - self.checkpoint_engine.makedirs(save_dir, exist_ok=True) - dist.barrier() - - if tag is None: - tag = f"global_step{self.global_steps}" - - # Ensure tag is a string - tag = str(tag) - self.checkpoint_engine.create(tag) - - # Ensure checkpoint tag is consistent across ranks - self._checkpoint_tag_validation(tag) - - if self.has_moe_layers: - self.save_non_zero_checkpoint = False - self._create_checkpoint_file(save_dir, tag, False) - self._save_moe_checkpoint(save_dir, - tag, - client_state=client_state, - exclude_frozen_parameters=exclude_frozen_parameters) - - # We distribute the task of saving layer checkpoint files among - # data parallel instances, so all procs should call _save_checkpoint. - # All procs then call module_state_dict(), but only procs of data - # parallel rank 0 save the general model params. - if not self.has_moe_layers: - self._create_checkpoint_file(save_dir, tag, False) - self._save_checkpoint(save_dir, - tag, - client_state=client_state, - exclude_frozen_parameters=exclude_frozen_parameters) - - if self.save_zero_checkpoint: - self._create_zero_checkpoint_files(save_dir, tag) - self._save_zero_checkpoint(save_dir, tag) - - if self._optimizer_has_ckpt_event_epilogue(): - self.optimizer.checkpoint_event_epilogue() - - # Save latest checkpoint tag - self.checkpoint_engine.commit(tag) - if save_latest and rank == 0: - with open(os.path.join(save_dir, 'latest'), 'w') as fd: - fd.write(tag) - - dist.barrier() - - return True - - def _get_non_moe_state_dict(self, full_state_dict): - """ - Get the state dict of the non-moe layers - """ - for key in list(full_state_dict.keys()): - if 'expert' in key and 'moe.gate.wg.weight' not in key: - full_state_dict.pop(key) - - return full_state_dict - - def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): - save_path = self._get_ckpt_name(save_dir, tag) - # A hack to save the checkpointing directory. Pipeline parallelism overrides - # module_state_dict() and uses this path to save the model. module_state_dict() - # then instead just returns None. - - # Using layer_#_export_# to save the model's expert state_dict - moe_layer_id = 0 - for n_module, module in self.module.named_modules(): - if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: - group_name = module.expert_group_name - num_local_experts = module.num_local_experts - expp_rank = groups._get_expert_parallel_rank(group_name) - exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) - # print(expp_rank, exp_dp_rank) - if exp_dp_rank != 0: - moe_layer_id += 1 - continue - - # get all moe parameters - moe_state_dict = {} - for n, p in module.state_dict().items(): - if 'expert' in n and 'moe.gate.wg.weight' not in n: - moe_state_dict[n_module + '.' + n] = p - moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' - # print(moe_state_dict.keys()) # until now, everything is fine. So the bug happens at next few lines - # Reorder the moe name rank, so that each checkpoint only has one expert - experts_state_dict = defaultdict(dict) - for key in list(moe_state_dict.keys()): - m = re.match(f".*{moe_str_prefix}([0-9]+).*", key) - - local_expert_id = None - if not m: - logger.warn(f'No expert found in key {key}.') - else: - local_expert_id = m.group(1) - - global_expert_id = expp_rank * \ - num_local_experts + int(local_expert_id) - expert_key = key.replace(f'{moe_str_prefix}{local_expert_id}', - f'{moe_str_prefix}{global_expert_id}') - # truncating extra tensor (shared) storage - truncated = moe_state_dict.pop(key).clone().detach() - experts_state_dict[str(global_expert_id)][expert_key] = truncated - - # let save the moe parameters - for global_expert_id, expert_state_dict in experts_state_dict.items(): - # save the moe parameters - moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu) - if self.random_ltd_enabled(): - expert_state_dict = remove_random_ltd_state_dict(expert_state_dict) - self.checkpoint_engine.save(expert_state_dict, moe_save_path) - moe_layer_id += 1 - - self._curr_ckpt_path = os.path.join(save_dir, tag) - - largest_group_name = groups._get_max_expert_size_name() - expp_rank = groups._get_expert_parallel_rank(largest_group_name) - exp_dp_rank = groups._get_expert_data_parallel_rank(largest_group_name) - - # In the case of E + D parallelism, only the - # first expert parallel group should save the expert weights - # since each expert parallel group is a copy of the model's experts - if exp_dp_rank != 0: - return - - # Save optimizer states. They are different across each exp parallel rank. - optimizer_state = { - 'optimizer': self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None - } - # TODO: why use BufferedWriter not the path - file_path = self._get_optimizer_ckpt_name(save_dir, tag, expp_rank) - self.checkpoint_engine.save(optimizer_state, file_path) - - # get non-moe parameters - model_state_dict = self._get_non_moe_state_dict( - self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters)) - - if expp_rank == 0: - # TODO: update num experts info,.. in checkpoint - state = { - 'module': - model_state_dict, - 'lr_scheduler': - self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, - 'data_sampler': - self.training_dataloader.data_sampler.state_dict() if - (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, - 'random_ltd': - self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, - 'sparse_tensor_module_names': - self.sparse_tensor_module_names, - 'skipped_steps': - self.skipped_steps, - 'global_steps': - self.global_steps, - 'global_samples': - self.global_samples, - 'zp_world_size': - self.zp_world_size, - 'dp_world_size': - self.dp_world_size, - 'mp_world_size': - self.mp_world_size, - 'num_experts': - self.num_experts - } - state.update(client_state) - logger.info(f'Saving model checkpoint: {save_path}') - self.checkpoint_engine.save(state, save_path) - - def _create_checkpoint_file(self, save_dir, tag, zero_checkpoint): - name_function = (self._get_zero_ckpt_name if zero_checkpoint else self._get_ckpt_name) - try: - checkpoint_name = name_function(save_dir, tag) - path = os.path.dirname(checkpoint_name) - self.checkpoint_engine.makedirs(path, exist_ok=True) - except: - logger.error(f"Failed saving model checkpoint to {save_dir} with tag {tag}") - return False - - return True - - def _create_zero_checkpoint_files(self, save_dir, tag): - success = True - # zero checkpoint files are created sequentially - for rank in range(dist.get_world_size(self.optimizer.zp_process_group)): - if rank == self.global_rank: - success = self._create_checkpoint_file(save_dir, tag, True) - - dist.barrier(group=self.optimizer.zp_process_group) - - return success - - def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): - - save_path = self._get_ckpt_name(save_dir, tag) - - zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() - - save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters - - # A hack to save the checkpointing directory. Pipeline parallelism overrides - # module_state_dict() and uses this path to save the model. module_state_dict() - # then instead just returns None. The module_state_dict() implementation in - # PipelineEngine expects the save path to be set in self._curr_ckpt_path. - self._curr_ckpt_path = os.path.join(save_dir, tag) - module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) - self._curr_ckpt_path = None - - state = dict(module=module, - buffer_names=self._get_buffer_names(), - optimizer=self.optimizer.state_dict() if self.optimizer and not zero_optimizer_state else None, - param_shapes=self._get_zero_param_shapes() if self.optimizer and zero_optimizer_state else None, - frozen_param_shapes=self._get_zero_frozen_param_attributes(self._get_param_shape_func) - if save_frozen_param else None, - shared_params=self._get_shared_params() if self.optimizer and zero_optimizer_state else None, - frozen_param_fragments=self._get_zero_frozen_param_attributes(self._get_param_fragment_func) - if save_frozen_param else None, - lr_scheduler=self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, - data_sampler=self.training_dataloader.data_sampler.state_dict() if - (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, - random_ltd=self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, - sparse_tensor_module_names=self.sparse_tensor_module_names, - skipped_steps=self.skipped_steps, - global_steps=self.global_steps, - global_samples=self.global_samples, - dp_world_size=self.seq_dp_world_size, - mp_world_size=self.mp_world_size, - ds_config=self.config, - ds_version=version) - state.update(client_state) - - if self.save_non_zero_checkpoint: - log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0, 1]) - self.checkpoint_engine.save(state, save_path) - - def _get_buffer_names(self): - buffer_names = [] - - # we save buffer names so that we could extract later the real buffers from the saved - # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they - # are intermixed with param placeholders - - # have to traverse the tree to be able to skip non-persistent buffers - def get_layer_named_buffers(module, prefix=""): - for name, buf in module.named_buffers(recurse=False): - if buf is not None and name not in module._non_persistent_buffers_set: - buffer_names.append(prefix + name) - - for name, child in module.named_children(): - if child is not None: - get_layer_named_buffers(child, prefix + name + ".") - - get_layer_named_buffers(self.module, prefix="") - - return buffer_names - - def _get_param_shape_func(self, param): - return param.ds_shape if hasattr(param, 'ds_id') else param.shape - - def _get_param_fragment_func(self, param): - return param.ds_tensor.detach().cpu() if hasattr(param, 'ds_id') else param.detach().cpu() - - def _get_zero_frozen_param_attributes(self, attr_func): - frozen_param_fragments = OrderedDict() - - for param in self.module.parameters(): - if param.requires_grad: - continue - if param not in self.param_names: - raise ValueError(f"failed to find frozen {param} in named params") - name = self.param_names[param] - frozen_param_fragments[name] = attr_func(param) - - return frozen_param_fragments - - def _get_zero_param_shapes(self): - """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the - optimizer. the names are exactly as in state_dict. The order is absolutely important, since - the saved data is just flattened data with no identifiers and requires reconstruction in the - same order it was saved. - We can't rely on self.module.named_parameters() to get the saved tensors, as some params - will be missing and others unsaved and then it'd be impossible to reconstruct state_dict - from the flattened weights. - optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions. - """ - param_group_shapes = [] - cnt = 0 - numel = 0 - - # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups - - # if we don't use it, we get parameters ordered incorrectly - if hasattr(self.optimizer, "round_robin_bit16_groups"): - bit16_groups = self.optimizer.round_robin_bit16_groups - elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"): - bit16_groups = self.optimizer.bf16_groups - else: - bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage( - ) == 2 else self.optimizer.fp16_groups - - for bit16_group in bit16_groups: - param_shapes = OrderedDict() - for param in bit16_group: - cnt += 1 - numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel() - shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape - if param not in self.param_names: - raise ValueError(f"failed to find optimizer param in named params") - name = self.param_names[param] - param_shapes[name] = shape - - # uncomment to debug zero_to_fp32.py problems - # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})") - param_group_shapes.append(param_shapes) - # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params") - - return param_group_shapes - - def _get_shared_params(self): - """ - Returns a dict of shared params, which can later be used to reconstruct the original state dict, - e.g. in `zero_to_fp32`. Each dict entry is a pair of param names, where the key is the name - of the variable that isn't stored and the value is the actual param holding data. - """ - shared_index = {} - shared_params_by_full_name = {} - - is_zero3_model = (self.zero_optimization_partition_weights() - and any(hasattr(param, "ds_id") for param in self.module.parameters())) - - def get_layer_state_dict(module, prefix=""): - # handle params - for name, param in module.named_parameters(recurse=False): - if param is None or (is_zero3_model and not hasattr(param, "ds_id")): - continue - key = prefix + name - - # When weights are manged by stage 3, we can't rely on param.data_ptr() as it will be reused - # as weights get gathered and reduced, but param.ds_id is unique across all zero weights - # (and shared params will have the same param.ds_id) - param_id = param.ds_id if is_zero3_model else param.data_ptr() - - if param_id in shared_index: - # shared weights - #print(f"`{key}` is shared with `{shared_index[param_id]}`") - shared_params_by_full_name[key] = shared_index[param_id] - else: - shared_index[param_id] = key - - for name, child in module.named_children(): - if child is not None: - get_layer_state_dict(child, prefix + name + ".") - - if dist.get_rank() == 0: - get_layer_state_dict(self.module, prefix="") - - return shared_params_by_full_name - - def _copy_recovery_script(self, save_path): - base_dir = os.path.dirname(os.path.dirname(__file__)) - script = "zero_to_fp32.py" - src = os.path.join(base_dir, "utils", script) - dst = os.path.join(save_path, script) - #logger.info(f"creating recovery script {dst}") - copyfile(src, dst) - self._change_recovery_script_permissions(dst) - - def _change_recovery_script_permissions(self, dst): - # make executable (safeguard for file shares - Azure as example) - try: - os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC) - except (FileNotFoundError, PermissionError) as e: - #this message is used in unit test TestZeRONonDistributed - logger.info( - f'Warning: Could not change permissions for {dst} due to error: {e}. Continuing without changing permissions.' - ) - - def _save_zero_checkpoint(self, save_path, tag): - zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag) - zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version) - self.checkpoint_engine.save(zero_sd, zero_checkpoint_name) - - if self.global_rank == 0: - self._copy_recovery_script(save_path) - ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero' - logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}') - - def _zero3_consolidated_16bit_state_dict(self): - """ - Get a full non-partitioned state_dict with fp16 weights on cpu. - Important: this function must be called on all ranks and not just rank 0. - This is similar to nn.Module.state_dict (modelled after _save_to_state_dict), but: - 1. consolidates the weights from different partitions on gpu0 - 2. works on one layer at a time to require as little gpu0 memory as possible, by - moving the already consolidated weights to cpu - 3. takes care to keep the shared params shared when gradually copying the params to cpu - Returns: - a consolidated fp16 ``state_dict`` on cpu on rank 0, ``None`` on other ranks - """ - if not self.zero_optimization_partition_weights(): - raise ValueError("this function requires ZeRO-3 mode") - - state_dict = OrderedDict() if dist.get_rank() == 0 else None - shared_params = {} - - def get_layer_state_dict(module, prefix=""): - # gather one layer at a time to be memory-efficient - # must use modifier_rank=0 to release GPU memory after each layer gathered - #see_memory_usage("before GatheredParameters", force=True) - with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): - if dist.get_rank() == 0: - # handle params - for name, param in module.named_parameters(recurse=False): - if param is None: - continue - key = prefix + name - # can't rely on param.data_ptr() as it will be reused as weights gets - # gathered and reduced, but param.ds_id is unique across all zero weights - # (and shared params will have the same param.ds_id) - if param.ds_id in shared_params: - # shared weights - #print(f"`{key}` is shared with `{shared_params[param.ds_id]}`") - state_dict[key] = state_dict[shared_params[param.ds_id]] - else: - state_dict[key] = param.detach().cpu() - shared_params[param.ds_id] = key - #print(f"param {param.ds_id} {param.shape} {key} ") - - # now buffers - not sure if need to take care of potentially shared weights here - for name, buf in module.named_buffers(recurse=False): - if (buf is not None and name not in module._non_persistent_buffers_set): - state_dict[prefix + name] = buf.detach().cpu() - #see_memory_usage("after GatheredParameters", force=True) - - for name, child in module.named_children(): - if child is not None: - get_layer_state_dict(child, prefix + name + ".") - - # Prepare for checkpoint save by ensuring all parameters are partitioned - if self._optimizer_has_ckpt_event_prologue(): - self.optimizer.checkpoint_event_prologue() - - see_memory_usage("before get_layer_state_dict", force=False) - get_layer_state_dict(self.module, prefix="") - see_memory_usage("after get_layer_state_dict", force=False) - - if self._optimizer_has_ckpt_event_epilogue(): - self.optimizer.checkpoint_event_epilogue() - - return state_dict - - def save_fp16_model(self, save_dir, save_filename="pytorch_model.bin"): - """has been renamed to save_16bit_model, keeping this around for backwards - compatibility""" - return self.save_16bit_model(save_dir, save_filename) - - def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin"): - """ - Save 16bit model weights - - This method saves the 16bit model weights at the desired destination. - - Arguments: - save_dir: Required. Directory for saving the model - save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin`` - - Returns: - ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if - stage3_gather_16bit_weights_on_model_save is ``False``. - - Important: all processes must call this method and not just the process with rank 0. It is - because the processes need to work in sync to gather the weights. This method will hang - waiting to synchronize with other processes if it's called just for the process with rank 0. - - """ - - path = os.path.join(save_dir, save_filename) - - if self.zero_optimization_partition_weights(): - if self.zero_gather_16bit_weights_on_model_save(): - # consolidation is expensive in time and memory and therefore isn't a default - state_dict = self._zero3_consolidated_16bit_state_dict() - else: - # the model will be bogus if not consolidated so don't confuse the user by saving it - logger.info( - f"Did not save the model {path} because `stage3_gather_16bit_weights_on_model_save` is False") - return False - else: - state_dict = self.module.state_dict() - - tag = f"global_step{self.global_steps}" - tag = str(tag) - self.checkpoint_engine.create(tag) - - if dist.get_rank() == 0: - self.checkpoint_engine.makedirs(save_dir, exist_ok=True) - logger.info(f"Saving model weights to {path}, tag: {tag}") - self.checkpoint_engine.save(state_dict, path) - - self.checkpoint_engine.commit(tag) - - return True - - def empty_partition_cache(self): - """ - Release GPU memory consumed by offloaded model parameters. - """ - if hasattr(self.optimizer, 'empty_partition_cache'): - self.optimizer.empty_partition_cache() - gc.collect() - get_accelerator().empty_cache() diff --git a/opensora/adaptor/stage_1_and_2.py b/opensora/adaptor/stage_1_and_2.py deleted file mode 100644 index 35f9a6d55..000000000 --- a/opensora/adaptor/stage_1_and_2.py +++ /dev/null @@ -1,2522 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import os -import pdb -from deepspeed import comm as dist -from packaging import version as pkg_version -from collections import OrderedDict -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors - -from deepspeed.runtime import ZeROOptimizer -from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler -from deepspeed.runtime.utils import (bwc_tensor_model_parallel_rank, get_global_norm, empty_cache, see_memory_usage, - inf, is_model_parallel_parameter, align_dense_tensors, all_gather_dp_groups) - -from deepspeed.runtime.zero.config import ZeroStageEnum -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum -from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.utils import logger -from deepspeed.moe.utils import is_moe_param -from deepspeed.git_version_info import version - -from deepspeed.runtime.constants import PIPE_REPLICATED -from deepspeed.accelerator import get_accelerator - -from deepspeed.checkpoint.constants import (DS_VERSION, GROUP_PADDINGS, PARTITION_COUNT, LOSS_SCALER, - SINGLE_PARTITION_OF_FP32_GROUPS, BASE_OPTIMIZER_STATE, - BASE_OPTIMIZER_STATE_STEP, CLIP_GRAD, ZERO_STAGE, PARAM_SLICE_MAPPINGS) -from deepspeed.utils import link_hp_params -from deepspeed.checkpoint import enable_universal_checkpoint - -from deepspeed.utils import groups - -from opensora.adaptor.zp_manager import zp_manager - -# Toggle this to true to enable correctness test -# with gradient partitioning and without -pg_correctness_test = False - -OPTIMIZER_ALLGATHER_TIMER = 'optimizer_allgather' -OPTIMIZER_GRADIENTS_TIMER = 'optimizer_gradients' -OPTIMIZER_STEP_TIMER = 'optimizer_step' -OPTIMIZER_TIMERS = [OPTIMIZER_ALLGATHER_TIMER, OPTIMIZER_GRADIENTS_TIMER, OPTIMIZER_STEP_TIMER] - - -def input(msg): - return - - -def split_half_float_double(tensors): - device_type = get_accelerator().device_name() - dtypes = [ - "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type), - "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type) - ] - buckets = [] - for i, dtype in enumerate(dtypes): - bucket = [t for t in tensors if t.type() == dtype] - if bucket: - buckets.append(bucket) - return buckets - - -def isclose(a, b, rtol=1e-09, atol=0.0): - return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) - - -def lcm(x, y): - from fractions import gcd # or can import gcd from `math` in Python 3 - return x * y // gcd(x, y) - - -def get_alignment_padding(tensor_list, alignment): - num_elements = sum([tensor.numel() for tensor in tensor_list]) - remainder = num_elements % alignment - return (alignment - remainder) if remainder else remainder - - -def move_to_cpu(tensor_list): - for tensor in tensor_list: - tensor.data = tensor.data.cpu() - - -def print_rank_msg(msg): - print(f"rank {dist.get_rank()} - {msg}") - - -def _get_padded_tensor(src_tensor, size): - if src_tensor.numel() >= size: - return src_tensor - padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) - slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) - slice_tensor.data.copy_(src_tensor.data) - return padded_tensor - - -def contigous_flatten(tensors): - return _flatten_dense_tensors([tensor.contiguous() for tensor in tensors]) - - -def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, zp_process_group): - for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)): - partition_id = dist.get_rank(group=zp_process_group[group_id]) - dp_world_size = dist.get_world_size(group=zp_process_group[group_id]) - if dp_world_size == 1: - # no groups share optimizer states - # pipeline parallel with bf16 will default call this even if dp size = 1. - continue - input_tensor = partitioned_params[partition_id].contiguous() - # print(f"call all_gather_into_tensor_dp_groups, input size is {input_tensor.size()}, " - # f"output size is {group_flat.size()}") - # - # print(f"groups_flat.size = {groups_flat.numel()}") - # print(f"partitioned_param_groups = {sum([v.numel() for v in partitioned_param_groups])}") - dist.all_gather_into_tensor(group_flat, input_tensor, zp_process_group[group_id]) - - -class DeepSpeedZeroOptimizer(ZeROOptimizer): - """ - DeepSpeedZeroOptimizer designed to reduce the memory footprint - required for training large deep learning models. - - For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models - https://arxiv.org/abs/1910.02054 - - For usage examples, refer to TODO: DeepSpeed Tutorial - - """ - - def __init__(self, - init_optimizer, - param_names, - timers, - static_loss_scale=1.0, - dynamic_loss_scale=False, - dynamic_loss_args=None, - verbose=True, - contiguous_gradients=True, - reduce_bucket_size=500000000, - use_multi_rank_bucket_allreduce=True, - allgather_bucket_size=5000000000, - dp_process_group=None, - expert_parallel_group=None, - expert_data_parallel_group=None, - reduce_scatter=True, - overlap_comm=False, - offload_optimizer_config=None, - mpu=None, - clip_grad=0.0, - gradient_accumulation_dtype=torch.float32, - communication_data_type=torch.float16, - postscale_gradients=True, - gradient_predivide_factor=1.0, - gradient_accumulation_steps=1, - ignore_unused_parameters=True, - partition_grads=True, - round_robin_gradients=False, - has_moe_layers=False, - fp16_master_weights_and_gradients=False, - elastic_checkpoint=False): - - if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: - self.cpu_offload = True - self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory - else: - self.cpu_offload = False - self.cpu_offload_pin_memory = False - - if dist.get_rank() == 0: - logger.info(f"Reduce bucket size {reduce_bucket_size}") - logger.info(f"Allgather bucket size {allgather_bucket_size}") - logger.info(f"CPU Offload: {self.cpu_offload}") - logger.info(f'Round robin gradient partitioning: {round_robin_gradients}') - # The fused optimizer does all the work. We need this layer for two reason: - # 1. maintain same user API from apex.fp16_utils - # 2. keep common stuff here in case we need to add ne552w fused optimizer later - - self.elastic_checkpoint = elastic_checkpoint - self.param_names = param_names - self.mpu = mpu - # differences from apex.fp16_utils: - # - assume all model params in fp16 - # - assume all params requires grad - # - flat by groups, not keeping state. TODO: remove state explicitly? - # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? - if not get_accelerator().is_available(): - raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).") - self.optimizer = init_optimizer - - # Use torch (un)flatten ops - self.flatten = contigous_flatten - self.unflatten = _unflatten_dense_tensors - - # ZeRO stage 1 (False) or 2 (True) - self.partition_gradients = partition_grads - self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1" - - self.timers = timers - - self.reduce_scatter = reduce_scatter - - self.overlap_comm = overlap_comm - - self.deepspeed_adam_offload = self.cpu_offload - - self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu' - - zp_manager.init_group() - self.zp_process_group = zp_manager.zp_group - zp_rank = dist.get_rank(group=self.zp_process_group) - zp_size = dist.get_world_size(group=self.zp_process_group) - print(f"zp rank is {zp_rank}, zp_size={zp_size}") - - self.dp_process_group = dp_process_group - - self.sequence_parallel_size = groups._get_sequence_parallel_world_size() - # expert parallel group - self.ep_process_group = expert_parallel_group - - # data parallel group for experts - self.expert_dp_process_group = expert_data_parallel_group - - # data parallel size for non-experts - dp_size = dist.get_world_size(group=self.dp_process_group) - - # For MoE models this maybe different for different param group - # It will be modified during MoE setup later in the init - self.real_zp_process_group = [self.zp_process_group for i in range(len(self.optimizer.param_groups))] - self.real_dp_process_group = [self.dp_process_group for i in range(len(self.optimizer.param_groups))] - self.partition_count = [zp_manager.zp_size for i in range(len(self.optimizer.param_groups))] - - self.is_gradient_accumulation_boundary = True - - # CPU-Offload requires contiguous gradients - self.contiguous_gradients = contiguous_gradients or self.cpu_offload - - self.has_moe_layers = has_moe_layers - if self.has_moe_layers: - self._configure_moe_settings() - self._global_grad_norm = 0. - - if mpu is None: - self.model_parallel_group = None - self.model_parallel_world_size = 1 - self.model_parallel_rank = 0 - else: - self.model_parallel_group = mpu.get_model_parallel_group() - self.model_parallel_world_size = mpu.get_model_parallel_world_size() - self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu) - - self.overflow = False - self.clip_grad = clip_grad - self.communication_data_type = communication_data_type - self.gradient_predivide_factor = gradient_predivide_factor - self.postscale_gradients = postscale_gradients - self.gradient_accumulation_steps = gradient_accumulation_steps - self.micro_step_id = 0 - self.ignore_unused_parameters = ignore_unused_parameters - self.round_robin_gradients = round_robin_gradients - - self.extra_large_param_to_reduce = None - self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients - - if self.fp16_master_weights_and_gradients: - assert self.cpu_offload and type(self.optimizer) in [DeepSpeedCPUAdam], \ - f"fp16_master_and_gradients requires optimizer to support keeping fp16 master and gradients while keeping the optimizer states in fp32." \ - f"Currently only supported using ZeRO-Offload with DeepSpeedCPUAdam. But current setting is ZeRO-Offload:{self.cpu_offload} and optimizer type {type(self.optimizer)}." \ - f"Either disable fp16_master_weights_and_gradients or enable {self.zero_stage_string} Offload with DeepSpeedCPUAdam." - - if self.reduce_scatter: - valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) - assert self.communication_data_type in valid_reduce_scatter_dtypes, f"{self.zero_stage_string} supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" - assert self.gradient_predivide_factor == 1.0, "gradient_predivide_factor != 1.0 is not yet supported with {self.zero_stage_string} with reduce scatter enabled" - assert self.postscale_gradients, "pre-scale gradients is not yet supported with {self.zero_stage_string} with reduce scatter enabled" - - # param flattened by groups - self.bit16_groups = [] - self.bit16_groups_flat = [] - - # param partitioned by data parallel degree - # this will contain a list of equal sized tensors - # each of which will be updated by a different process - self.parallel_partitioned_bit16_groups = [] - - # a single 32-bit partition of the parallel partitioned parameters - # that this process will update - self.single_partition_of_fp32_groups = [] - - # param partition info - - # These are the parameters in each group that will not be updated by this process directly - self.params_not_in_partition = [] - - # These are the parameters that will be updated by this process directly - self.params_in_partition = [] - - # Offset from the first parameter in the self.params_in_partition - # the parameter boundaries may not align with partition boundaries - # so we need to keep track of the offset - self.first_offset = [] - - # number of elements per partition in each group - self.partition_size = [] - - # align nccl all-gather send buffers to 4-byte boundary - self.nccl_start_alignment_factor = 16 # 4-byte alignment/sizeof(fp16) = 2 - - assert ( - allgather_bucket_size % self.nccl_start_alignment_factor == 0 - ), f"allgather_bucket_size must be a multiple of nccl_start_alignment_factor, {self.nccl_start_alignment_factor} " - - self.all_reduce_print = False - self.dtype = self.optimizer.param_groups[0]['params'][0].dtype - self.gradient_accumulation_dtype = gradient_accumulation_dtype - - if self.dtype != self.gradient_accumulation_dtype: - self.use_separate_grad_accum = True - else: - self.use_separate_grad_accum = False - if self.use_separate_grad_accum and not self.partition_gradients: - self.use_grad_accum_attribute = True - else: - self.use_grad_accum_attribute = False - - self.round_robin_bit16_groups = [] - self.round_robin_bit16_indices = [] - - # Use different parallel to do all_to_all_reduce related things - # padding on each partition for alignment purposes - self.groups_padding = [] - # loop to deal with groups - for i, param_group in enumerate(self.optimizer.param_groups): - partition_id = dist.get_rank(group=self.real_zp_process_group[i]) - - # push this group to list before modify - # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group - trainable_parameters = [] - for param in param_group['params']: - if param.requires_grad: - param.grad_accum = None - trainable_parameters.append(param) - self.bit16_groups.append(trainable_parameters) - - # not sure why apex was cloning the weights before flattening - # removing cloning here - - see_memory_usage(f"Before moving param group {i} to CPU") - # move all the parameters to cpu to free up GPU space for creating flat buffer - move_to_cpu(self.bit16_groups[i]) - empty_cache() - see_memory_usage(f"After moving param group {i} to CPU", force=False) - - # Reorder group parameters for load balancing of gradient partitioning during backward among ranks. - # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks. - # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging - # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m). - if self.round_robin_gradients: - round_robin_tensors, round_robin_indices = self._round_robin_reorder( - self.bit16_groups[i], dist.get_world_size(group=self.real_zp_process_group[i])) - else: - round_robin_tensors = self.bit16_groups[i] - round_robin_indices = list(range(len(self.bit16_groups[i]))) - - self.round_robin_bit16_groups.append(round_robin_tensors) - self.round_robin_bit16_indices.append(round_robin_indices) - - # create flat buffer in CPU and move to GPU - self.bit16_groups_flat.append( - self.flatten_dense_tensors_aligned( - self.round_robin_bit16_groups[i], - self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_zp_process_group[i])).to( - get_accelerator().current_device_name())) - see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False) - - # Record padding required for alignment - if partition_id == dist.get_world_size(group=self.real_zp_process_group[i]) - 1: - padding = self.bit16_groups_flat[i].numel() - sum( - [t.numel() for t in self.round_robin_bit16_groups[i]]) - else: - padding = 0 - self.groups_padding.append(padding) - - if dist.get_rank(group=self.real_zp_process_group[i]) == 0: - see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False) - - # set model bit16 weight to slices of flattened buffer - self._update_model_bit16_weights(i) - - # divide the flat weights into near equal partition equal to the data parallel degree - # each process will compute on a different part of the partition - data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i) - self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) - - # print(f"self.bit16_groups_flat[i].size = {self.bit16_groups_flat[i].numel()}") - # print(f"data_parallel_partitions = {sum([v.numel() for v in data_parallel_partitions])}") - - # verify that data partition start locations are 4-byte aligned - for partitioned_data in data_parallel_partitions: - assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0) - - # A partition of the fp32 master weights that will be updated by this process. - # Note that the params in single_partition_of_fp32_groups is cloned and detached - # from the origin params of the model. - if not fp16_master_weights_and_gradients: - self.single_partition_of_fp32_groups.append(self.parallel_partitioned_bit16_groups[i][partition_id].to( - self.device).clone().float().detach()) - else: - self.single_partition_of_fp32_groups.append(self.parallel_partitioned_bit16_groups[i][partition_id].to( - self.device).clone().half().detach()) - - # Set local optimizer to have flat params of its own partition. - # After this, the local optimizer will only contain its own partition of params. - # In that case, the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params(zero stage1). - self.single_partition_of_fp32_groups[ - i].requires_grad = True # keep this in case internal optimizer uses it - param_group['params'] = [self.single_partition_of_fp32_groups[i]] - - partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_zp_process_group[i]) - params_in_partition, params_not_in_partition, first_offset = self.get_partition_info( - self.round_robin_bit16_groups[i], partition_size, partition_id) - - self.partition_size.append(partition_size) - self.params_in_partition.append(params_in_partition) - self.params_not_in_partition.append(params_not_in_partition) - self.first_offset.append(first_offset) - - self.reduce_bucket_size = int(reduce_bucket_size) - self.use_multi_rank_bucket_allreduce = use_multi_rank_bucket_allreduce - self.allgather_bucket_size = int(allgather_bucket_size) - - self.reduction_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream() - # self.copy_grad_stream = get_accelerator().Stream() - self.callback_queued = False - - self.param_dict = {} - - # map between param_id and bool to specify if a param is in this partition - self.is_param_in_current_partition = {} - - self.grads_in_ipg_bucket = [] - self.params_in_ipg_bucket = [] - self.elements_in_ipg_bucket = 0 - self.params_already_reduced = [] - self._release_ipg_buffers() - self.previous_reduced_grads = None - self.ipg_bucket_has_moe_params = False - - # simplified param id - self.param_id = {} - - # interesting code: unique ids being assigned to individual parameters - largest_param_numel = 0 - count = 0 - for i, params_group in enumerate(self.bit16_groups): - for param in params_group: - unique_id = id(param) - self.param_id[unique_id] = count - self.param_dict[count] = param - self.params_already_reduced.append(False) - if param.numel() > largest_param_numel: - largest_param_numel = param.numel() - count = count + 1 - - for param_group in self.params_in_partition: - for param in param_group: - self.is_param_in_current_partition[self.get_param_id(param)] = True - - for param_group in self.params_not_in_partition: - for param in param_group: - self.is_param_in_current_partition[self.get_param_id(param)] = False - - if self.cpu_offload: - self.accumulated_grads_in_cpu = {} - self.norm_for_param_grads = {} - self.local_overflow = False - self.grad_position = {} - self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel, - device=self.device, - dtype=self.dtype) - if self.cpu_offload_pin_memory: - self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory( - self.temp_grad_buffer_for_cpu_offload) - self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel, - device=get_accelerator().current_device_name(), - dtype=self.dtype) - for i, params_group in enumerate(self.bit16_groups): - self.get_grad_position(i, self.params_in_partition[i], self.first_offset[i], self.partition_size[i]) - - # mapping from parameter to partition that it belongs to - self.param_to_partition_ids = {} - - # stores if a partition has been reduced in this step - self.is_partition_reduced = {} - - # number of grads in partition that still need to be computed - self.remaining_grads_in_partition = {} - - # total number of grads in partition - self.total_grads_in_partition = {} - - # stores if a grad in a partition has been computed or not - self.is_grad_computed = {} - - # stores the offset at which a parameter gradient needs to be inserted in a partition - self.grad_partition_insertion_offset = {} - - # the offset in the gradient at which it must be inserted at the beginning of the partition - self.grad_start_offset = {} - - # will store the averaged gradients required by this partition - self.averaged_gradients = {} - - # For cpu_offload, will store the averaged gradients required by this partition - self.offload_gradient_dict = {} - - # store index of first parameter in each partition - self.first_param_index_in_partition = {} - - # initializes all data structures for implementing gradient partitioning - self.initialize_gradient_partitioning_data_structures() - - # resets the data structure value for the next backward propagation - self.reset_partition_gradient_structures() - - # creates backward hooks for gradient partitioning - if self.partition_gradients or self.overlap_comm: - self.create_reduce_and_remove_grad_hooks() - - self.custom_loss_scaler = False - self.external_loss_scale = None - - # we may have a way of fusing dynamic scale. Do not support for now - self.loss_scaler = CreateLossScaler(dtype=self.dtype, - static_loss_scale=static_loss_scale, - dynamic_scaling=dynamic_loss_scale, - dynamic_loss_args=dynamic_loss_args) - self.dynamic_loss_scale = self.loss_scaler.dynamic - - if self.dtype != torch.float16: - # Only fp16 should use dynamic loss scaling - assert self.loss_scaler.cur_scale == 1.0 - assert not self.dynamic_loss_scale - - see_memory_usage("Before initializing optimizer states", force=True) - self.initialize_optimizer_states() - see_memory_usage("After initializing optimizer states", force=True) - - if dist.get_rank() == 0: - logger.info(f"optimizer state initialized") - - if dist.get_rank(group=self.zp_process_group) == 0: - see_memory_usage(f"After initializing ZeRO optimizer", force=True) - - self._link_all_hp_params() - self._enable_universal_checkpoint() - self._param_slice_mappings = self._create_param_mapping() - - def _enable_universal_checkpoint(self): - for lp_param_group in self.bit16_groups: - enable_universal_checkpoint(param_list=lp_param_group) - - def _create_param_mapping(self): - param_mapping = [] - for i, _ in enumerate(self.optimizer.param_groups): - param_mapping_per_group = OrderedDict() - for lp in self.bit16_groups[i]: - if lp._hp_mapping is not None: - lp_name = self.param_names[lp] - param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() - param_mapping.append(param_mapping_per_group) - - return param_mapping - - def _link_all_hp_params(self): - dp_world_size = dist.get_world_size(group=self.zp_process_group) - if self.cpu_offload: - self._get_offload_gradient_dict() - - for i, _ in enumerate(self.optimizer.param_groups): - # Link bit16 and fp32 params in partition - partition_id = dist.get_rank(group=self.real_zp_process_group[i]) - partition_size = self.bit16_groups_flat[i].numel() // dp_world_size - flat_hp_partition = self.single_partition_of_fp32_groups[i] - link_hp_params(lp_param_list=self.bit16_groups[i], - flat_hp_partition=flat_hp_partition, - gradient_dict=self.averaged_gradients, - offload_gradient_dict=self.offload_gradient_dict, - use_offload=self.cpu_offload, - param_group_index=i, - partition_start=partition_id * partition_size, - partition_size=partition_size, - partition_optimizer_state=self.optimizer.state[flat_hp_partition], - dp_group=self.real_zp_process_group[i]) - - def is_moe_group(self, group): - return 'moe' in group and group['moe'] - - def _configure_moe_settings(self): - # if we're using ZeRO stage 2, ensure contiguous gradients are used - if self.partition_gradients: - assert self.contiguous_gradients, "Contiguous Gradients in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" - # NOTE: To run ZeRO stage 1 with MoE, we need to set self.contiguous_gradients to True or ignore the assertion - if not self.partition_gradients and not self.contiguous_gradients: - logger.warn( - "ZeRO Stage 1 has not been thoroughly tested with MoE. This configuration is still experimental.") - assert self.reduce_scatter, "Reduce Scatter in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" - - assert any( - [self.is_moe_group(group) for group in self.optimizer.param_groups] - ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" - self.is_moe_param_group = [] - for i, group in enumerate(self.optimizer.param_groups): - if self.is_moe_group(group): - assert all([is_moe_param(param) - for param in group['params']]), "All params in MoE group must be MoE params" - self.real_zp_process_group[i] = self.expert_dp_process_group[group['name']] - self.partition_count[i] = dist.get_world_size(group=self.expert_dp_process_group[group['name']]) - self.is_moe_param_group.append(True) - else: - self.is_moe_param_group.append(False) - - assert self.expert_dp_process_group is not None, "Expert data parallel group should be configured with MoE" - assert self.ep_process_group is not None, "Expert parallel group should be configured with MoE" - - def _update_model_bit16_weights(self, group_index): - updated_params = self.unflatten(self.bit16_groups_flat[group_index], - self.round_robin_bit16_groups[group_index]) - for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params): - p.data = q.data - - # set model fp16 weight to slices of reordered flattened buffer - for param_index, param in enumerate(self.bit16_groups[group_index]): - new_index = self.round_robin_bit16_indices[group_index][param_index] - param.data = self.round_robin_bit16_groups[group_index][new_index].data - - def _round_robin_reorder(self, tensor_list, num_partitions): - - # disable round robin if need to debug something - # return tensor_list, list(range(len(tensor_list))) - - partition_tensors = {} - - for i, tensor in enumerate(tensor_list): - j = i % num_partitions - if not j in partition_tensors: - partition_tensors[j] = [] - partition_tensors[j].append((i, tensor)) - - reordered_tensors = [] - reordered_indices = {} - - for partition_index in partition_tensors.keys(): - for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]): - reordered_indices[original_index] = len(reordered_tensors) - reordered_tensors.append(tensor) - - return reordered_tensors, reordered_indices - - def _release_ipg_buffers(self): - if self.contiguous_gradients: - self.ipg_buffer = None - self.grads_in_partition = None - self.grads_in_partition_offset = 0 - - def initialize_optimizer_states(self): - - for i, group in enumerate(self.bit16_groups): - single_grad_partition = torch.zeros(int(self.partition_size[i]), - dtype=self.single_partition_of_fp32_groups[i].dtype, - device=self.device) - self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory( - single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition - - # Initialize the optimizer states with the flattened fp32 partition. - # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers - # which do lazy initialization of the state at the first call to step. - if isinstance(self.optimizer, torch.optim.Adagrad): - self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults) - else: - self.optimizer.step() - - if not self.cpu_offload: - for group in self.single_partition_of_fp32_groups: - group.grad = None # class init - - return - - ######################################################################### - #################### ZeRO Stage 1 - reduce gradients #################### - ######################################################################### - def reduce_gradients(self, pipeline_parallel=False): - world_size = dist.get_world_size(self.zp_process_group) - my_rank = dist.get_rank(self.zp_process_group) - - # with PP we must create ipg buffer, since backward is handled outside zero - if pipeline_parallel and self.contiguous_gradients: - self.ipg_buffer = [] - buf_0 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - self.ipg_buffer.append(buf_0) - self.ipg_index = 0 - - if not self.overlap_comm: - for i, group in enumerate(self.bit16_groups): - for param in group: - grad_reduc = self.get_gradient_for_reduction(param) - if grad_reduc is not None: - self.reduce_ready_partitions_and_remove_grads(param, i) - # reduce any pending grads in either hook/non-hook case - self.overlapping_partition_gradients_reduce_epilogue() - - ######################################################################### - #########################ZeRO Partition Gradients######################## - ######################################################################### - - def get_first_param_index(self, group_id, param_group, partition_id): - for index, param in enumerate(param_group): - param_id = self.get_param_id(param) - if partition_id in self.param_to_partition_ids[group_id][param_id]: - return index - return None - - def initialize_gradient_partitioning_data_structures(self): - - for i, param_group in enumerate(self.round_robin_bit16_groups): - total_partitions = dist.get_world_size(group=self.real_zp_process_group[i]) - - self.param_to_partition_ids[i] = {} - self.is_partition_reduced[i] = {} - self.total_grads_in_partition[i] = {} - self.remaining_grads_in_partition[i] = {} - self.is_grad_computed[i] = {} - self.grad_partition_insertion_offset[i] = {} - self.grad_start_offset[i] = {} - self.first_param_index_in_partition[i] = {} - - for partition_id in range(total_partitions): - self.is_grad_computed[i][partition_id] = {} - self.grad_partition_insertion_offset[i][partition_id] = {} - self.grad_start_offset[i][partition_id] = {} - self.total_grads_in_partition[i][partition_id] = 0 - self.initialize_gradient_partition(i, param_group, partition_id) - self.is_partition_reduced[i][partition_id] = False - self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( - i, param_group, partition_id) - - def independent_gradient_partition_epilogue(self): - self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) - self.reduce_ipg_grads() - self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) - - # if dist.get_rank() == 0: - # logger.info("Params already reduced %s", self.params_already_reduced) - for i in range(len(self.params_already_reduced)): - self.params_already_reduced[i] = False - - if self.overlap_comm: - get_accelerator().synchronize() - # It is safe to clear previously reduced grads of other partitions - self._clear_previous_reduced_grads() - - if self.cpu_offload is False: - for i, _ in enumerate(self.bit16_groups): - - if not i in self.averaged_gradients or self.averaged_gradients[i] is None: - self.averaged_gradients[i] = self.get_flat_partition( - self.params_in_partition[i], - self.first_offset[i], - self.partition_size[i], - dtype=self.gradient_accumulation_dtype, - device=get_accelerator().current_device_name(), - return_tensor_list=True) - else: - avg_new = self.get_flat_partition(self.params_in_partition[i], - self.first_offset[i], - self.partition_size[i], - dtype=self.gradient_accumulation_dtype, - device=get_accelerator().current_device_name(), - return_tensor_list=True) - - for accumulated_grad, new_avg_grad in zip(self.averaged_gradients[i], avg_new): - accumulated_grad.add_(new_avg_grad) - - self._release_ipg_buffers() - - # No need to keep the gradients anymore. - # All gradients required by the step - # are in self.averaged_gradients - self.zero_grad(set_to_none=True) - see_memory_usage(f"End ipg_epilogue") - - # resets all partition to no reduced - # sets remaining grads to the total number of grads in each partition - # set is grad computed to false for all grads in partition - def reset_partition_gradient_structures(self): - for i, _ in enumerate(self.bit16_groups): - total_partitions = dist.get_world_size(group=self.real_zp_process_group[i]) - for partition_id in range(total_partitions): - self.is_partition_reduced[i][partition_id] = False - self.remaining_grads_in_partition[i][partition_id] = self.total_grads_in_partition[i][partition_id] - - for param_id in self.is_grad_computed[i][partition_id]: - self.is_grad_computed[i][partition_id][param_id] = False - - def initialize_gradient_partition(self, i, param_group, partition_id): - - def set_key_value_list(dictionary, key, value): - if key in dictionary: - dictionary[key].append(value) - else: - dictionary[key] = [value] - - def increment_value(dictionary, key): - if key in dictionary: - dictionary[key] += 1 - else: - dictionary[key] = 1 - - partition_size = self.partition_size[i] - - start_index = partition_size * partition_id - end_index = partition_size * (partition_id + 1) - - current_index = 0 - first_offset = 0 - - for param in param_group: - - param_size = param.numel() - param_id = self.get_param_id(param) - - if start_index <= current_index < end_index: - set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) - increment_value(self.total_grads_in_partition[i], partition_id) - - self.is_grad_computed[i][partition_id][param_id] = False - - self.grad_partition_insertion_offset[i][partition_id][param_id] = current_index - start_index - self.grad_start_offset[i][partition_id][param_id] = 0 - - elif current_index < start_index < (current_index + param_size): - assert (first_offset == 0 - ), "This can happen either zero or only once as this must be the first tensor in the partition" - first_offset = start_index - current_index - - set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) - increment_value(self.total_grads_in_partition[i], partition_id) - - self.is_grad_computed[i][partition_id][param_id] = False - - self.grad_partition_insertion_offset[i][partition_id][param_id] = 0 - self.grad_start_offset[i][partition_id][param_id] = first_offset - - current_index = current_index + param_size - - def overlapping_partition_gradients_reduce_epilogue(self): - self.independent_gradient_partition_epilogue() - - def fill_grad_accum_attribute(self): - for group in self.bit16_groups: - for param in group: - if param.grad is not None: - if param.grad_accum is None: - param.grad_accum = param.grad.to(self.gradient_accumulation_dtype) - else: - param.grad_accum.add_( - param.grad.to(self.gradient_accumulation_dtype).view(param.grad_accum.shape)) - param.grad = None - - def get_gradient_for_reduction(self, param): - if self.use_grad_accum_attribute: - return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None - else: - return param.grad - - def get_param_gradient_attribute(self, param): - return param.grad_accum if self.use_grad_accum_attribute else param.grad - - # Clear the tensor the reduction gradient attribute is pointing to - def clear_grad_attribute(self, param): - if self.use_grad_accum_attribute: - param.grad_accum = None - else: - param.grad = None - - def create_reduce_and_remove_grad_hooks(self): - self.grad_accs = [] - for i, param_group in enumerate(self.bit16_groups): - for param in param_group: - if param.requires_grad: - def wrapper(param, i): - param_tmp = param.expand_as(param) - grad_acc = param_tmp.grad_fn.next_functions[0][0] - - def reduce_partition_and_remove_grads(*notneeded): - self.reduce_ready_partitions_and_remove_grads(param, i) - - grad_acc.register_hook(reduce_partition_and_remove_grads) - self.grad_accs.append(grad_acc) - - wrapper(param, i) - - def get_param_id(self, param): - unique_id = id(param) - return self.param_id[unique_id] - - def report_ipg_memory_usage(self, tag, param_elems): - elem_count = self.elements_in_ipg_bucket + param_elems - percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size - see_memory_usage( - f"{tag}: elems in_bucket {self.elements_in_ipg_bucket} param {param_elems} max_percent {percent_of_bucket_size}" - ) - - # create a flat tensor aligned at the alignment boundary - def flatten_dense_tensors_aligned(self, tensor_list, alignment): - return self.flatten(align_dense_tensors(tensor_list, alignment)) - - ############### Independent Partition Gradient ######################## - def reduce_independent_p_g_buckets_and_remove_grads(self, param, i): - - grad_reduc = self.get_gradient_for_reduction(param) - if self.elements_in_ipg_bucket + param.numel() > self.reduce_bucket_size: - self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel()) - self.reduce_ipg_grads() - if self.contiguous_gradients and self.overlap_comm: - # Swap ipg_index between 0 and 1 - self.ipg_index = 1 - self.ipg_index - self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel()) - - param_id = self.get_param_id(param) - assert self.params_already_reduced[param_id] == False, \ - f"The parameter {param_id} has already been reduced. \ - Gradient computed twice for this partition. \ - Multiple gradient reduction is currently not supported" - - if self.contiguous_gradients: - if param.numel() > self.reduce_bucket_size: - self.extra_large_param_to_reduce = param - else: - # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening - new_grad_tensor = self.ipg_buffer[self.ipg_index].narrow(0, self.elements_in_ipg_bucket, param.numel()) - new_grad_tensor.copy_(grad_reduc.view(-1)) - grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) - - self.elements_in_ipg_bucket += param.numel() - - assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient" - - self.grads_in_ipg_bucket.append(grad_reduc) - self.params_in_ipg_bucket.append((i, param, param_id)) - - # make sure the average tensor function knows how to average the gradients - if is_moe_param(param): - self.ipg_bucket_has_moe_params = True - - self.report_ipg_memory_usage("End ipg_remove_grads", 0) - - def print_rank_0(self, message): - if dist.get_rank() == 0: - logger.info(message) - - def gradient_reduction_w_predivide(self, tensor): - - dp_world_size = dist.get_world_size(group=self.dp_process_group) - - tensor_to_allreduce = tensor - - if self.communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(self.communication_data_type) - - if self.postscale_gradients: - if self.gradient_predivide_factor != 1.0: - tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor) - - dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) - - if self.gradient_predivide_factor != dp_world_size: - tensor_to_allreduce.mul_(self.gradient_predivide_factor / - (dp_world_size / float(self.sequence_parallel_size))) - else: - tensor_to_allreduce.div_(dp_world_size / float(self.sequence_parallel_size)) - dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) - - if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - tensor.copy_(tensor_to_allreduce) - - return tensor - - def allreduce_and_copy_with_multiple_ranks(self, - small_bucket, - log=None, - divide=True, - process_group=None, - bucket_ranks=None): - process_group = self.zp_process_group if process_group is None else process_group - allreduced = self.allreduce_bucket(small_bucket, log=log, divide=divide, process_group=process_group) - for buf, synced, bucket_rank in zip(small_bucket, self.unflatten(allreduced, small_bucket), bucket_ranks): - if dist.get_rank(group=process_group) == bucket_rank: - buf.copy_(synced) - - def allreduce_and_scatter(self, bucket, numel_per_bucket=500000000, log=None, divide=True, process_group=None): - small_bucket = [] - small_bucket_ranks = [] - numel = 0 - allreduce_sizes = [] - - for i, bucket_elem in enumerate(bucket): - rank, tensor = bucket_elem - small_bucket.append(tensor) - small_bucket_ranks.append(rank) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy_with_multiple_ranks(small_bucket, - log=None, - divide=divide, - process_group=process_group, - bucket_ranks=small_bucket_ranks) - small_bucket = [] - small_bucket_ranks = [] - numel = 0 - - if len(small_bucket) > 0: - self.allreduce_and_copy_with_multiple_ranks(small_bucket, - log=None, - divide=divide, - process_group=process_group, - bucket_ranks=small_bucket_ranks) - - def average_tensor(self, tensor): - if self.overlap_comm: - stream = self.reduction_stream - if not get_accelerator().is_synchronized_device(): - stream.wait_stream(get_accelerator().current_stream()) - else: - stream = get_accelerator().current_stream() - - with get_accelerator().stream(stream): - if not self.reduce_scatter: - self.gradient_reduction_w_predivide(tensor) - return - - # Accumulate destination ranks and bucket offsets for each gradient slice. - # Note: potential future optimization, record access pattern of parameters - # in backward pass and partition gradients w.r.t. access pattern so that our - # bucket is guaranteed to be contiguous w.r.t. ranks - rank_and_offsets = [] - real_dp_process_group = [] - curr_size = 0 - prev_id, prev_process_group = -1, None - - process_group = self.zp_process_group - # count = 0 - for i, param, param_id in self.params_in_ipg_bucket: - - process_group = self.zp_process_group - grad_reduc = self.get_gradient_for_reduction(param) - # Averages gradients at parameter level if ipg has a moe param - # Otherwise averaging is done at the entire buffer level at the end of the loop - # MoE param have different groups - if self.ipg_bucket_has_moe_params: - process_group = self.expert_dp_process_group[param.group_name] if is_moe_param( - param) else self.zp_process_group - grad_reduc.data.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size)) - - partition_ids = self.param_to_partition_ids[i][param_id] - assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids - ]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}" - partition_size = self.partition_size[i] - # Get all partition ids + their offsets - partition_ids_w_offsets = [] - for partition_id in partition_ids: - offset = self.grad_start_offset[i][partition_id][param_id] - partition_ids_w_offsets.append((partition_id, offset)) - partition_ids_w_offsets.sort(key=lambda t: t[1]) - - # Calculate rank and offsets for grad slices - for idx in range(len(partition_ids_w_offsets)): - partition_id, offset = partition_ids_w_offsets[idx] - - # if dist.get_rank() == 0 and count < 100: - # print(f"Rank {dist.get_rank()} rank offset id {idx} calculated dp size {dist.get_world_size(group=process_group)} real dp size {dist.get_world_size(self.real_dp_process_group[i])} and dst: {partition_id}") - # count += 1 - - # Calculate numel for grad slice depending on partition location - if idx == len(partition_ids_w_offsets) - 1: - # Last partition_id uses its own offset - numel = param.numel() - offset - else: - # Set numel to next partition's offset - numel = partition_ids_w_offsets[idx + 1][1] - offset - - # Merge bucket ranges if they belong to the same rank - if partition_id == prev_id and process_group == prev_process_group: - prev_pid, prev_size, prev_numel = rank_and_offsets[-1] - rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel) - else: - rank_and_offsets.append((partition_id, curr_size, numel)) - real_dp_process_group.append(process_group) - curr_size += numel - prev_id, prev_process_group = partition_id, process_group - - if not self.ipg_bucket_has_moe_params: - tensor.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) - - buckets = {} - for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets): - grad_slice = tensor.narrow(0, int(bucket_offset), int(numel)) - bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else ( - dst, real_dp_process_group[i]) - if bucket_key not in buckets: - buckets[bucket_key] = [] - if self.use_multi_rank_bucket_allreduce: - buckets[bucket_key].append((dst, grad_slice)) - else: - buckets[bucket_key].append(grad_slice) - - for bucket_key in buckets: - if self.use_multi_rank_bucket_allreduce: - self.allreduce_and_scatter(buckets[bucket_key], - numel_per_bucket=self.reduce_bucket_size, - divide=self.ipg_bucket_has_moe_params, - process_group=bucket_key) - else: - dst, process_group = bucket_key - self.allreduce_no_retain(buckets[bucket_key], - numel_per_bucket=self.reduce_bucket_size, - rank=dst, - divide=self.ipg_bucket_has_moe_params, - process_group=process_group) - - ############################################################################## - ############################# CPU Offload Methods############################# - ############################################################################## - def get_grad_position(self, group_id, tensor_list, first_offset, partition_size): - current_offset = 0 - - for i, tensor in enumerate(tensor_list): - param_id = self.get_param_id(tensor) - param_start_offset = 0 - - num_elements = tensor.numel() - - # we need to offset to get to the right element - if i == 0 and first_offset > 0: - tensor_offset = first_offset - num_elements = num_elements - tensor_offset - param_start_offset = first_offset - - # we dont need all elements of the tensor - if num_elements > (partition_size - current_offset): - num_elements = partition_size - current_offset - - self.grad_position[param_id] = [ - int(group_id), int(param_start_offset), - int(current_offset), int(num_elements) - ] - current_offset += num_elements - - def update_overflow_tracker_for_param_grad(self, param): - grad_accum = self.get_param_gradient_attribute(param) - if grad_accum is not None and self._has_inf_or_nan(grad_accum.data): - self.local_overflow = True - - def _get_offload_gradient_dict(self): - for param_group_index, _ in enumerate(self.optimizer.param_groups): - self.offload_gradient_dict[param_group_index] = [] - for lp_param in self.params_in_partition[param_group_index]: - param_id = self.get_param_id(lp_param) - [_, _, dest_offset, num_elements] = self.grad_position[param_id] - dest_tensor = self.single_partition_of_fp32_groups[param_group_index].grad.view(-1).narrow( - 0, dest_offset, num_elements) - self.offload_gradient_dict[param_group_index].append(dest_tensor) - - def async_accumulate_grad_in_cpu_via_gpu(self, param): - param_id = self.get_param_id(param) - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - # copy to a preexisiting buffer to avoid memory allocation penalty - dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel()) - - # buffer for storing gradients for this parameter in CPU - def buffer_to_accumulate_to_in_cpu(): - if not self.fp16_master_weights_and_gradients: - buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device) - return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer - else: - return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) - - # accumulate gradients into param.grad_accum or parts of it that belongs to this partition - def accumulate_gradients(): - grad_accum = self.get_param_gradient_attribute(param) - if not self.fp16_master_weights_and_gradients: - dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True) - grad_accum.data.view(-1).add_(dest_buffer) - else: - dest_buffer.narrow(0, source_offset, - num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1), - non_blocking=True) - grad_accum.data.view(-1).narrow(0, source_offset, - num_elements).add_(dest_buffer.narrow(0, source_offset, num_elements)) - - # move accumulated gradients back to CPU - def copy_gradients_to_cpu(): - grad_accum = self.get_param_gradient_attribute(param) - if not self.fp16_master_weights_and_gradients: - self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True) - else: - self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1).narrow( - 0, source_offset, num_elements), - non_blocking=True) - - if param_id not in self.accumulated_grads_in_cpu: - self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu() - - if self.micro_step_id > 0: - accumulate_gradients() - - # at the boundary we will send 32bit directly - if not self.is_gradient_accumulation_boundary: - copy_gradients_to_cpu() - - def set_norm_for_param_grad(self, param): - param_id = self.get_param_id(param) - grad_accum = self.get_param_gradient_attribute(param) - accumulated_grad = self.accumulated_grads_in_cpu[ - param_id] if self.gradient_accumulation_steps > 1 else grad_accum - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - start = source_offset - accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) - - self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) - - def set_norm_for_param_grad_in_gpu(self, param): - param_id = self.get_param_id(param) - grad_accum = self.get_param_gradient_attribute(param) - if grad_accum is None: - accumulated_grad = param.grad - else: - accumulated_grad = grad_accum - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - start = source_offset - accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) - - self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) - - def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param): - param_id = self.get_param_id(param) - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - dest_tensor = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) - - grad_accum = self.get_param_gradient_attribute(param) - if grad_accum is None: - src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) - else: - src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) - if not self.fp16_master_weights_and_gradients: - src_tensor = src_tensor.float() - - dest_tensor.copy_(src_tensor, non_blocking=True) - param.grad = None # offload only - - def complete_grad_norm_calculation_for_cpu_offload(self, params): - total_norm = 0.0 - norm_type = 2.0 - for p in params: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): - param_id = self.get_param_id(p) - # as some model have trainable parameters but skipped in training, - # their backward hooks in self.create_reduce_and_remove_grad_hooks() will not run, - # so they have no norm_for_param_grads - if param_id in self.norm_for_param_grads: - param_norm = self.norm_for_param_grads[param_id] - total_norm += param_norm.item() ** 2 - else: - # As unused parameters in modules may not be expected sometimes, - # add an explicit error msg when it occurred and an option to - # avoid the error - assert self.ignore_unused_parameters, """ - This assert indicates that your module has parameters that - were not used in producing loss. - You can avoid this assert by - (1) enable ignore_unused_parameters option in zero_optimization config; - (2) making sure all trainable parameters and `forward` function - outputs participate in calculating loss. - """ - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) - - self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) - - total_norm = total_norm_cuda[0].item() ** (1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - ############################################################################################ - def copy_grads_in_partition(self, param): - if self.cpu_offload: - - if self.gradient_accumulation_steps > 1: - self.async_accumulate_grad_in_cpu_via_gpu(param) - - if self.is_gradient_accumulation_boundary: - self.set_norm_for_param_grad_in_gpu(param) - - self.update_overflow_tracker_for_param_grad(param) - - self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param) - - return - # print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}") - if self.grads_in_partition is None: - self.grads_in_partition_offset = 0 - total_size = 0 - for group in self.params_in_partition: - for param_in_partition in group: - total_size += param_in_partition.numel() - - see_memory_usage(f"before copying {total_size} gradients into partition") - self.grads_in_partition = torch.empty(int(total_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - see_memory_usage(f"after copying {total_size} gradients into partition") - - grad_reduc = self.get_gradient_for_reduction(param) - # The allreduce buffer will be rewritten. Copy the gradients in partition to a new buffer - new_grad_tensor = self.grads_in_partition.view(-1).narrow(0, self.grads_in_partition_offset, param.numel()) - new_grad_tensor.copy_(grad_reduc.view(-1)) - grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) - # print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}") - self.grads_in_partition_offset += param.numel() - - def reduce_ipg_grads(self): - if self.contiguous_gradients: - if self.extra_large_param_to_reduce is not None: - assert len(self.params_in_ipg_bucket) == 1, "more than 1 param in ipg bucket, this shouldn't happen" - _, _, param_id = self.params_in_ipg_bucket[0] - assert self.get_param_id(self.extra_large_param_to_reduce - ) == param_id, "param in ipg bucket does not match extra-large param" - extra_large_grad_reduc = self.get_gradient_for_reduction(self.extra_large_param_to_reduce) - self.average_tensor(extra_large_grad_reduc.view(-1)) - self.extra_large_param_to_reduce = None - else: - self.average_tensor(self.ipg_buffer[self.ipg_index]) - else: - self.buffered_reduce_fallback(None, - self.grads_in_ipg_bucket, - elements_per_buffer=self.elements_in_ipg_bucket) - - if self.overlap_comm: - stream = self.reduction_stream - elif self.cpu_offload: - # TODO: copy_grad_stream is disabled because of race with reduce. This hurts perf and should be fixed. - # get_accelerator().synchronize() - # stream = self.copy_grad_stream - stream = get_accelerator().current_stream() - else: - stream = get_accelerator().current_stream() - - with get_accelerator().stream(stream): - for _, param, param_id in self.params_in_ipg_bucket: - - assert self.params_already_reduced[param_id] == False, \ - f"The parameter {param_id} has already been reduced. \ - Gradient computed twice for this partition. \ - Multiple gradient reduction is currently not supported" - - self.params_already_reduced[param_id] = True - if self.partition_gradients: - if not self.is_param_in_current_partition[param_id]: - if self.overlap_comm and self.contiguous_gradients is False: - # Clear grads of other partitions during the next reduction - # to avoid clearing them before the reduction is complete. - if self.previous_reduced_grads is None: - self.previous_reduced_grads = [] - self.previous_reduced_grads.append(param) - else: - self.clear_grad_attribute(param) - elif self.contiguous_gradients: - self.copy_grads_in_partition(param) - else: # zero stage 1 - partition only optimizer state - if self.contiguous_gradients and self.is_param_in_current_partition[param_id]: - self.copy_grads_in_partition(param) - - self.grads_in_ipg_bucket = [] - self.params_in_ipg_bucket = [] - self.ipg_bucket_has_moe_params = False - self.elements_in_ipg_bucket = 0 - ##################################################################### - - def reduce_ready_partitions_and_remove_grads(self, param, i): - if self.partition_gradients or self.is_gradient_accumulation_boundary: - self.reduce_independent_p_g_buckets_and_remove_grads(param, i) - - def zero_reduced_gradients(self, partition_id, i): - - def are_all_related_partitions_reduced(params_id): - for partition_id in self.param_to_partition_ids[i][params_id]: - if not self.is_partition_reduced[i][partition_id]: - return False - return True - - for params_id in self.is_grad_computed[i][partition_id]: - if are_all_related_partitions_reduced(params_id): - self.param_dict[params_id].grad = None # dead code - - def flatten_and_print(self, message, tensors, start=0, n=5): - flatten_tensor = self.flatten(tensors) - - def print_func(): - logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) - - self.sequential_execution(print_func, message) - - def get_grads_to_reduce(self, i, partition_id): - - def get_reducible_portion(key): - grad = self.param_dict[key].grad - total_elements = grad.numel() - start = self.grad_start_offset[i][partition_id][key] - num_elements = min(total_elements - start, - self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) - if not pg_correctness_test: - if num_elements == total_elements: - return grad - else: - return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) - else: - if num_elements == total_elements: - return grad.clone() - else: - return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) - - grads_to_reduce = [] - for key in self.is_grad_computed[i][partition_id]: - grad = get_reducible_portion(key) - grads_to_reduce.append(grad) - return grads_to_reduce - - def sequential_execution(self, function, message, group=None): - if group is None: - group = self.zp_process_group - if dist.get_rank(group=group) == 0: - logger.info(message) - for id in range(dist.get_world_size(group=group)): - if id == dist.get_rank(group=group): - function() - dist.barrier(group=group) - - def set_none_gradients_to_zero(self, i, partition_id): - for param_id in self.is_grad_computed[i][partition_id]: - param = self.param_dict[param_id] - if param.grad is None: - param.grad = torch.zero_like(param) - - ######################Reduction Related Methods############################## - def allreduce_bucket(self, bucket, rank=None, log=None, divide=True, process_group=None): - rank = None - tensor = self.flatten(bucket) - - process_group = self.zp_process_group if process_group is None else process_group - - tensor_to_allreduce = tensor - - if pg_correctness_test or self.sequence_parallel_size > 1: - communication_data_type = torch.float32 - else: - communication_data_type = self.communication_data_type - - if communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(communication_data_type) - - if divide: - tensor_to_allreduce.div_( - dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) - - tensor_to_allreduce = tensor_to_allreduce.contiguous() - if rank is None: - # "All Reducing" - dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) - else: - global_rank = dist.get_global_rank(self.dp_process_group, rank) - dist.reduce(tensor_to_allreduce, global_rank, group=self.dp_process_group) - - if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - if rank is None or rank == dist.get_rank(group=process_group): - tensor.copy_(tensor_to_allreduce) - - return tensor - - def _clear_previous_reduced_grads(self): - if self.previous_reduced_grads is not None: - for param in self.previous_reduced_grads: - self.clear_grad_attribute(param) - self.previous_reduced_grads = None - - # if rank is specified do a reduction instead of an allreduce - def allreduce_and_copy(self, small_bucket, rank=None, log=None, divide=True, process_group=None): - process_group = self.zp_process_group if process_group is None else process_group - if self.overlap_comm: - get_accelerator().synchronize() - # It is safe to clear the previously reduced grads of other partitions - self._clear_previous_reduced_grads() - stream = self.reduction_stream - else: - stream = get_accelerator().current_stream() - - with get_accelerator().stream(stream): - allreduced = self.allreduce_bucket( - small_bucket, - rank=rank, - log=log, - divide=divide, - process_group=process_group, - ) - if rank is None or rank == dist.get_rank(group=self.zp_process_group): - for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): - buf.copy_(synced) - - def allreduce_no_retain( - self, - bucket, - numel_per_bucket=500000000, - rank=None, - log=None, - divide=True, - process_group=None, - ): - small_bucket = [] - numel = 0 - for tensor in bucket: - small_bucket.append(tensor) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy(small_bucket, rank=rank, log=None, divide=divide, process_group=process_group) - small_bucket = [] - numel = 0 - - if len(small_bucket) > 0: - self.allreduce_and_copy(small_bucket, rank=rank, log=log, divide=divide, process_group=process_group) - - # allows using reduction of gradients instead of using all_reduce - - def buffered_reduce_fallback(self, rank, grads, elements_per_buffer=500000000, log=None): - split_buckets = split_half_float_double(grads) - - for i, bucket in enumerate(split_buckets): - self.allreduce_no_retain(bucket, numel_per_bucket=elements_per_buffer, rank=rank, log=log) - - ############################################################################# - ############################################################################# - ############################################################################# - - # views the tensor as multiple partitions and returns - # those partitions - def get_data_parallel_partitions(self, tensor, group_id): - partitions = [] - - dp = dist.get_world_size(group=self.real_zp_process_group[group_id]) - # dp_id = dist.get_rank(group=self.real_dp_process_group[group_id]) - - total_num_elements = tensor.numel() - - base_size = total_num_elements // dp - remaining = total_num_elements % dp - - start = 0 - for id in range(dp): - partition_size = base_size - if id < remaining: - partition_size = partition_size + 1 - partitions.append(tensor.narrow(0, start, partition_size)) - start = start + partition_size - return partitions - - def get_partition_info(self, tensor_list, partition_size, partition_id): - params_in_partition = [] - params_not_in_partition = [] - - start_index = partition_size * partition_id - end_index = partition_size * (partition_id + 1) - - current_index = 0 - first_offset = 0 - - for tensor in tensor_list: - - tensor_size = tensor.numel() - - if start_index <= current_index < end_index: - params_in_partition.append(tensor) - - elif current_index < start_index < (current_index + tensor_size): - params_in_partition.append(tensor) - - assert (first_offset == 0 - ), "This can happen either zero or only once as this must be the first tensor in the partition" - first_offset = start_index - current_index - - else: - params_not_in_partition.append(tensor) - - current_index = current_index + tensor_size - - return params_in_partition, params_not_in_partition, first_offset - - def zero_grad(self, set_to_none=True): - """ - Zero FP16 parameter grads. - """ - # FP32 grad should never exist. - # For speed, set model fp16 grad to None by default - # zero all pointers to grad tensors - for group in self.bit16_groups: - for p in group: - if set_to_none: - p.grad = None # epilogue and in step - p.grad_accum = None - else: - if p.grad is not None: - p.grad.detach_() - p.grad.zero_() - - def _model_parallel_all_reduce(self, tensor, op): - """ Perform all reduce within model parallel group, if any. - """ - if self.model_parallel_group is None or self.model_parallel_world_size == 1: - pass - else: - dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) - - def get_grad_norm_direct(self, gradients, params, norm_type=2): - """Clips gradient norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - max_norm (float or int): max norm of the gradients - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - """ - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(g.data.abs().max() for g in gradients) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.dp_process_group) - - # Take max across all GPUs. - self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX) - total_norm = total_norm_cuda[0].item() - else: - total_norm = 0.0 - # if dist.get_rank() == 0: - # logger.info(f"Total Norm beginning {total_norm}") - for g, p in zip(gradients, params): - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): - param_norm = g.data.double().norm(2) - total_norm += param_norm.item() ** 2 - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) - - self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) - - total_norm = total_norm_cuda[0].item() ** (1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - # creates a flat fused tensor from the tensor list starting at the first_offset - # in the first tensor of the list. If there are not enough elements in the tensor - # list then the flat tensor will be padded with zeros - def get_flat_partition(self, tensor_list, first_offset, partition_size, dtype, device, return_tensor_list=False): - flat_tensor_list = [] - current_size = 0 - - for i, tensor in enumerate(tensor_list): - grad_accum = self.get_param_gradient_attribute(tensor) - if grad_accum is None: - grad_accum = torch.zeros_like(tensor, dtype=dtype) - - tensor = grad_accum - num_elements = tensor.numel() - tensor_offset = 0 - - # we need to offset to get to the right element - if i == 0 and first_offset > 0: - tensor_offset = first_offset - num_elements = num_elements - tensor_offset - - # we dont need all elements of the tensor - if num_elements > (partition_size - current_size): - num_elements = partition_size - current_size - - # we need a narrow view of the tensor based on the tensor offset and number of elements that - # we need from this tensor - if tensor_offset > 0 or num_elements < tensor.numel(): - flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) - else: - flat_tensor_list.append(tensor) - - current_size = current_size + num_elements - - # this means its the last partition and does not align with the dp boundary. We need to pad before flattening - if current_size < partition_size: - flat_tensor_list.append(torch.zeros(int(partition_size - current_size), dtype=dtype, device=device)) - - if return_tensor_list: - return flat_tensor_list - - return self.flatten(flat_tensor_list) - - def free_grad_in_param_list(self, param_list): - for p in param_list: - p.grad = None # in step - p.grad_accum = None - - def reset_cpu_buffers(self): - self.norm_for_param_grads = {} - self.local_overflow = False - - def set_lr(self, lr): - """Set the learning rate.""" - for param_group in self.optimizer.param_groups: - param_group["lr"] = lr - - def get_lr(self): - """Return the current learning rate.""" - return self.optimizer.param_groups[0]["lr"] - - def override_loss_scale(self, loss_scale): - if loss_scale != self.external_loss_scale: - logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') - self.custom_loss_scaler = True - self.external_loss_scale = loss_scale - - def scaled_global_norm(self, norm_type=2): - assert norm_type == 2, "only L2 norm supported" - norm_groups = [] - for i, group in enumerate(self.bit16_groups): - partition_id = dist.get_rank(group=self.real_zp_process_group[i]) - if self.cpu_offload: - norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(self.params_in_partition[i])) - single_grad_partition = self.single_partition_of_fp32_groups[i].grad - else: - norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.params_in_partition[i])) - - if self.has_moe_layers: - self._average_expert_grad_norms(norm_groups) - - # note that the get_global_norm function only supports l2 norm - return get_global_norm(norm_list=norm_groups) - - def get_bit16_param_group(self, group_no): - bit16_partitions = self.parallel_partitioned_bit16_groups[group_no] - partition_id = dist.get_rank(group=self.real_zp_process_group[group_no]) - return [bit16_partitions[dist.get_rank(group=self.real_zp_process_group[group_no])]] - - def _optimizer_step(self, group_no): - original_param_groups = self.optimizer.param_groups - self.optimizer.param_groups = [original_param_groups[group_no]] - # Disabling this as the C++ side copy & synchronize is not working correctly - # from deepspeed.ops.adam import DeepSpeedCPUAdam - # if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half: - # self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)]) - # else: - # self.optimizer.step() - self.optimizer.step() - self.optimizer.param_groups = original_param_groups - - def step(self, closure=None): - """ - Not supporting closure. - """ - self.micro_step_id = -1 - - see_memory_usage(f"In step before checking overflow") - - # First compute norm for all group so we know if there is overflow - if self.dtype == torch.float16: - self.check_overflow() - - prev_scale = self.loss_scale - self._update_scale(self.overflow) - if self.overflow: - see_memory_usage('After overflow before clearing gradients') - self.zero_grad(set_to_none=True) - if self.cpu_offload: - self.reset_cpu_buffers() - else: - self.averaged_gradients = {} - - see_memory_usage('After overflow after clearing gradients') - - for timer in OPTIMIZER_TIMERS: - self.timers(timer).start() - self.timers(timer).stop() - return - - # Step 1:- Calculate gradient norm using bit-16 grads - see_memory_usage('Before norm calculation') - scaled_global_grad_norm = self.scaled_global_norm() - self._global_grad_norm = scaled_global_grad_norm / prev_scale - see_memory_usage('After norm before optimizer') - - # Step 2:- run optimizer and upscaling simultaneously - for i, group in enumerate(self.bit16_groups): - self.timers(OPTIMIZER_GRADIENTS_TIMER).start() - partition_id = dist.get_rank(group=self.real_zp_process_group[i]) - if self.cpu_offload: - single_grad_partition = self.single_partition_of_fp32_groups[i].grad - self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) - - self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() - self.timers(OPTIMIZER_STEP_TIMER).start() - self._optimizer_step(i) - - # Disabled, this is not currently working - # from deepspeed.ops.adam import DeepSpeedCPUAdam - # if not (type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half): - # bit16_partitions = self.parallel_partitioned_bit16_groups[i] - # fp32_partition = self.single_partition_of_fp32_groups[i] - # bit16_partitions[partition_id].data.copy_(fp32_partition.data) - bit16_partitions = self.parallel_partitioned_bit16_groups[i] - fp32_partition = self.single_partition_of_fp32_groups[i] - bit16_partitions[partition_id].data.copy_(fp32_partition.data) - - self.timers(OPTIMIZER_STEP_TIMER).stop() - else: - # free gradients for all the parameters that are not updated by this process(ZeRO stage2) - self.free_grad_in_param_list(self.params_not_in_partition[i]) - - # create a flat gradients for parameters updated by this process - # If we are last partition, ensure we have same size grads and partition size, if not pad with zero tensors - if partition_id == dist.get_world_size(group=self.real_zp_process_group[i]) - 1: - single_grad_partition = self.flatten_dense_tensors_aligned( - self.averaged_gradients[i], - int(self.partition_size[i])).to(self.single_partition_of_fp32_groups[i].dtype) - else: - single_grad_partition = self.flatten(self.averaged_gradients[i]).to( - self.single_partition_of_fp32_groups[i].dtype) - assert single_grad_partition.numel() == self.partition_size[i], \ - "averaged gradients have different number of elements that partition size {} {} {} {}".format( - single_grad_partition.numel(), self.partition_size[i], i, partition_id) - - self.single_partition_of_fp32_groups[i].grad = single_grad_partition - # release all the gradient since we have already created a necessary copy in dp_grad_partition(ZeRO stage2) - self.free_grad_in_param_list(self.params_in_partition[i]) - - self.averaged_gradients[i] = None - - self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) - - self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() - - # Step 3:- run the optimizer if no offloading - self.timers(OPTIMIZER_STEP_TIMER).start() - self._optimizer_step(i) - # Step 4:- get rid of the fp32 gradients. Not needed anymore - self.single_partition_of_fp32_groups[i].grad = None - del single_grad_partition - bit16_partitions = self.parallel_partitioned_bit16_groups[i] - fp32_partition = self.single_partition_of_fp32_groups[i] - bit16_partitions[partition_id].data.copy_(fp32_partition.data) - self.timers(OPTIMIZER_STEP_TIMER).stop() - - see_memory_usage('After optimizer before all-gather') - if self.cpu_offload: - self.reset_cpu_buffers() - - self.timers(OPTIMIZER_ALLGATHER_TIMER).start() - - # if dist.get_rank(group=self.dp_process_group) == 0: - # pdb.set_trace() # 或者使用其他调试工具 - - # Gather the updated weights from everyone. - # Then all partitions of the model parameters are updated and ready for next round forward. - all_gather_into_tensor_dp_groups(groups_flat=self.bit16_groups_flat, - partitioned_param_groups=self.parallel_partitioned_bit16_groups, - zp_process_group=self.real_zp_process_group) - self.timers(OPTIMIZER_ALLGATHER_TIMER).stop() - - # TODO: we probably don't need this? just to be safe - for i in range(len(self.bit16_groups)): - self._update_model_bit16_weights(i) - - self.timers.log(OPTIMIZER_TIMERS) - see_memory_usage('After zero_optimizer step') - - return - - @torch.no_grad() - def update_lp_params(self): - for i, (bit16_partitions, fp32_partition) in enumerate( - zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): - partition_id = dist.get_rank(group=self.real_zp_process_group[i]) - bit16_partitions[partition_id].data.copy_(fp32_partition.data) - # print_rank_0(f'update_lp_params {i=} {partition_id=}', force=True) - # if i == 0: - # print_rank_0(f'{fp32_partition[:10]=}', force=True) - all_gather_into_tensor_dp_groups(groups_flat=self.bit16_groups_flat, - partitioned_param_groups=self.parallel_partitioned_bit16_groups, - zp_process_group=self.real_zp_process_group) - - def _average_expert_grad_norms(self, norm_groups): - for i, norm in enumerate(norm_groups): - if self.is_moe_param_group[i]: - scaled_norm = norm * 1.0 / float(dist.get_world_size(group=self.dp_process_group)) - scaled_norm_tensor = torch.tensor(scaled_norm, - device=get_accelerator().device_name(), - dtype=torch.float) - dist.all_reduce(scaled_norm_tensor, group=self.dp_process_group) - norm_groups[i] = scaled_norm_tensor.item() - - def unscale_and_clip_grads(self, grad_groups_flat, total_norm): - # compute combined scale factor for this group - combined_scale = self.loss_scale - if self.clip_grad > 0.: - # norm is in fact norm*scale - clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad - if clip > 1: - combined_scale = clip * self.loss_scale - - for grad in grad_groups_flat: - if isinstance(grad, list): - sub_partitions = grad - for g in sub_partitions: - g.data.mul_(1. / combined_scale) - else: - grad.data.mul_(1. / combined_scale) - - def _check_overflow(self, partition_gradients=True): - self.overflow = self.has_overflow(partition_gradients) - - # `params` is a list / generator of torch.Variable - def has_overflow_serial(self, params, is_grad_list=False): - for p in params: - if p.grad is not None and self._has_inf_or_nan(p.grad.data): - return True - - return False - - def has_overflow_partitioned_grads_serial(self): - for i in range(len(self.bit16_groups)): - for j, grad in enumerate(self.averaged_gradients[i]): - if grad is not None and self._has_inf_or_nan(grad.data, j): - return True - return False - - def has_overflow(self, partition_gradients=True): - if partition_gradients: - overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial() - overflow_gpu = get_accelerator().ByteTensor([overflow]) - '''This will capture overflow across all data parallel and expert parallel process - Since expert parallel process are a subset of data parallel process''' - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) - - else: - params = [] - for group in self.bit16_groups: - for param in group: - params.append(param) - - overflow = self.has_overflow_serial(params, is_grad_list=partition_gradients) - overflow_gpu = get_accelerator().ByteTensor([overflow]) - - # Since each model parallel GPU carries only part of the model, - # make sure overflow flag is synced across all the model parallel GPUs - self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) - - overflow = overflow_gpu[0].item() - return bool(overflow) - - # `x` is a torch.Tensor - @staticmethod - def _has_inf_or_nan(x, j=None): - try: - # if x is half, the .float() incurs an additional deep copy, but it's necessary if - # Pytorch's .sum() creates a one-element tensor of the same type as x - # (which is true for some recent version of pytorch). - cpu_sum = float(x.float().sum()) - # More efficient version that can be used if .sum() returns a Python scalar - # cpu_sum = float(x.sum()) - except RuntimeError as instance: - # We want to check if inst is actually an overflow exception. - # RuntimeError could come from a different error. - # If so, we still want the exception to propagate. - if "value cannot be converted" not in instance.args[0]: - raise - return True - else: - if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: - return True - return False - - def backward(self, loss, retain_graph=False): - """ - :attr:`backward` performs the following steps: - - 1. fp32_loss = loss.float() - 2. scaled_loss = fp32_loss*loss_scale - 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves - """ - self.micro_step_id += 1 - - if self.contiguous_gradients: - self.ipg_buffer = [] - buf_0 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - self.ipg_buffer.append(buf_0) - - # Use double buffers to avoid data access conflict when overlap_comm is enabled. - if self.overlap_comm: - buf_1 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - self.ipg_buffer.append(buf_1) - self.ipg_index = 0 - - if self.custom_loss_scaler: - scaled_loss = self.external_loss_scale * loss - scaled_loss.backward() - else: - self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) - - # Only for Stage 1, Mode 2 - if self.use_grad_accum_attribute: - self.fill_grad_accum_attribute() - - def check_overflow(self, partition_gradients=True): - self._check_overflow(partition_gradients) - - def _update_scale(self, has_overflow=False): - self.loss_scaler.update_scale(has_overflow) - - # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" - def _get_state(self): - return self.optimizer.state - - def _set_state(self, value): - self.optimizer.state = value - - state = property(_get_state, _set_state) - - # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" - # (for example, to adjust the learning rate) - def _get_param_groups(self): - return self.optimizer.param_groups - - def _set_param_groups(self, value): - self.optimizer.param_groups = value - - param_groups = property(_get_param_groups, _set_param_groups) - - # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" - def _get_loss_scale(self): - if self.custom_loss_scaler: - return self.external_loss_scale - else: - return self.loss_scaler.cur_scale - - def _set_loss_scale(self, value): - self.loss_scaler.cur_scale = value - - loss_scale = property(_get_loss_scale, _set_loss_scale) - cur_scale = property(_get_loss_scale, _set_loss_scale) - - # Return group tensor after removing paddings that are added for alignment to DP world size. - # This method works on the assumption that each group contains a single flattened tensor. - def _get_groups_without_padding(self, groups_with_padding): - groups_without_padding = [] - for i, group in enumerate(groups_with_padding): - lean_length = group.numel() - self.groups_padding[i] - groups_without_padding.append(group[:lean_length]) - - return groups_without_padding - - # Return optimizer state after removing paddings that are added for alignment. - def _get_state_without_padding(self, state_with_padding, padding): - lean_state = {} - for key, value in state_with_padding.items(): - if torch.is_tensor(value): - lean_length = value.numel() - padding - lean_state[key] = value[:lean_length] - else: - lean_state[key] = value - - return lean_state - - # Return base optimizer states. - # This method assumes that each param group contains a single flattened tensor. - def _get_base_optimizer_state(self): - optimizer_groups_state = [] - for i, group in enumerate(self.optimizer.param_groups): - p = group['params'][0] - lean_optimizer_state = self._get_state_without_padding(self.optimizer.state[p], self.groups_padding[i]) - optimizer_groups_state.append(lean_optimizer_state) - - return optimizer_groups_state - - def state_dict(self): - """ - Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. - This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict - of the contained Pytorch optimizer. - Example:: - checkpoint = {} - checkpoint['model'] = model.state_dict() - checkpoint['optimizer'] = optimizer.state_dict() - torch.save(checkpoint, "saved.pth") - """ - state_dict = {} - state_dict[LOSS_SCALER] = self.loss_scaler - state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale - state_dict['overflow'] = self.overflow - state_dict[CLIP_GRAD] = self.clip_grad - - if self.elastic_checkpoint: - state_dict[BASE_OPTIMIZER_STATE] = self._get_base_optimizer_state() - - if "step" in self.optimizer.param_groups[0]: - # Assuming "step" is the only item that changes through training iterations - assert all(group["step"] == self.optimizer.param_groups[0]["step"] - for group in self.optimizer.param_groups), "All param groups must have the same step value" - state_dict[BASE_OPTIMIZER_STATE_STEP] = self.optimizer.param_groups[0]["step"] - else: - state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() - - # Remove paddings for DP alignment to enable loading for other alignment values - fp32_groups_without_padding = self._get_groups_without_padding(self.single_partition_of_fp32_groups) - state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding - - state_dict[ - ZERO_STAGE] = ZeroStageEnum.gradients if self.partition_gradients else ZeroStageEnum.optimizer_states - state_dict[GROUP_PADDINGS] = self.groups_padding - state_dict[PARTITION_COUNT] = self.partition_count - - state_dict[DS_VERSION] = version - state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings - - return state_dict - - # Restore base optimizer fp32 weights from elastic checkpoint by: - # 1) Merging fp32 weights from checkpoints of all partitions - # 2) Extracting fp32 weights for current partition from merged weights - # 3) Using extracted weights to update base optimizer weights directly. - def _restore_from_elastic_fp32_weights(self, all_state_dict): - merged_single_partition_of_fp32_groups = [] - - for i in range(len(self.single_partition_of_fp32_groups)): - partition_id = dist.get_rank(group=self.real_zp_process_group[i]) - merged_partitions = [sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict] - if self.is_moe_group(self.optimizer.param_groups[i]): - ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) - merged_partitions = [merged_partitions[i] for i in ranks] - flat_merged_partitions = self.flatten_dense_tensors_aligned( - merged_partitions, - self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_zp_process_group[i])) - dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, i) - merged_single_partition_of_fp32_groups.append(dp_partitions[partition_id]) - - for current, saved in zip(self.single_partition_of_fp32_groups, merged_single_partition_of_fp32_groups): - current.data.copy_(saved.data) - - # Restore base optimizer fp32 weights from ZeRO fp16 or bfloat16 weights - def _restore_from_bit16_weights(self): - for group_id, (bit16_partitions, fp32_partition) in enumerate( - zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): - partition_id = dist.get_rank(group=self.real_zp_process_group[group_id]) - fp32_partition.data.copy_(bit16_partitions[partition_id].data) - - # Refresh the fp32 master params from the fp16 or bfloat16 copies. - def refresh_fp32_params(self): - self._restore_from_bit16_weights() - - # Extract optimizer state for current partition from merged states of all partitions - def _partition_base_optimizer_state(self, state_key, all_partition_states, group_id): - partition_id = dist.get_rank(group=self.real_zp_process_group[group_id]) - alignment = dist.get_world_size(group=self.real_zp_process_group[group_id]) - if torch.is_tensor(all_partition_states[0]): - flat_merged_partitions = self.flatten_dense_tensors_aligned(all_partition_states, alignment) - dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, group_id) - return dp_partitions[partition_id] - else: - # Assume non-tensor states are not partitioned and equal across ranks, so return first one - return all_partition_states[0] - - def _restore_base_optimizer_state(self, base_optimizer_group_states): - if type(base_optimizer_group_states) == dict: - base_optimizer_group_states = base_optimizer_group_states['state'] - for i, group in enumerate(self.optimizer.param_groups): - p = group['params'][0] - for key, saved in base_optimizer_group_states[i].items(): - if torch.is_tensor(self.optimizer.state[p][key]): - dst_tensor = self.optimizer.state[p][key] - src_tensor = _get_padded_tensor(saved, dst_tensor.numel()) - self.optimizer.state[p][key].data.copy_(src_tensor.data) - else: - self.optimizer.state[p][key] = saved - - def get_ep_ranks(self, rank=0, group_name=None): - from deepspeed.utils import groups - expert_parallel_size_ = groups._get_expert_parallel_world_size(group_name) - world_size = groups._get_data_parallel_world_size() - rank = groups._get_expert_parallel_rank(group_name) - ranks = range(rank, world_size, expert_parallel_size_) - return list(ranks) - - # Restore base optimizer state from elastic checkpoint by - # 1) Merging optimizer state from checkpoints of all partitions - # 2) Extracting optimizer state for current partition from the merged state - # 3) Using the extracted value to directly update the base optimizer. - def _restore_elastic_base_optimizer_state(self, all_state_dict): - base_optimizer_group_states = [] - for i in range(len(self.optimizer.param_groups)): - partition_states = {} - all_partition_group_states = [sd[BASE_OPTIMIZER_STATE][i] for sd in all_state_dict] - - if self.is_moe_group(self.optimizer.param_groups[i]): - ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) - all_partition_group_states = [all_partition_group_states[i] for i in ranks] - - for key in all_partition_group_states[0].keys(): - all_partition_states = [all_states[key] for all_states in all_partition_group_states] - partition_states[key] = self._partition_base_optimizer_state(key, all_partition_states, i) - base_optimizer_group_states.append(partition_states) - - self._restore_base_optimizer_state(base_optimizer_group_states) - - # Restore step - if BASE_OPTIMIZER_STATE_STEP in all_state_dict[0]: - assert all(sd[BASE_OPTIMIZER_STATE_STEP] == all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] - for sd in all_state_dict), "State dicts of all partitions must have the same step value" - loaded_param_groups_step = all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] - for param_group in self.optimizer.param_groups: - param_group['step'] = loaded_param_groups_step - - def load_state_dict(self, - state_dict_list, - load_optimizer_states=True, - load_from_fp32_weights=False, - checkpoint_folder=None, - load_serial=None): - if checkpoint_folder: - self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) - else: - self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) - - def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): - self._load_hp_checkpoint_state(checkpoint_folder) - - @property - def param_groups(self): - """Forward the wrapped optimizer's parameters.""" - return self.optimizer.param_groups - - def _load_hp_checkpoint_state(self, checkpoint_dir): - checkpoint_dir = os.path.join(checkpoint_dir, "zero") - optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") - assert os.path.isfile( - optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' - optim_sd = torch.load(optim_state_path) - self._load_global_state(optim_sd) - - tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) - tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \ - else self.mpu.get_tensor_model_parallel_world_size() - - for i, _ in enumerate(self.optimizer.param_groups): - for lp in self.bit16_groups[i]: - if lp._hp_mapping is not None: - # print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}") - lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank, - tp_world_size) - - def _load_global_state(self, sd): - self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) - self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) - self.overflow = sd.get('overflow', self.overflow) - self.clip_grad = sd.get(CLIP_GRAD, self.clip_grad) - - ckpt_version = sd.get(DS_VERSION, False) - assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" - ckpt_version = pkg_version.parse(ckpt_version) - - # zero stage 1 mode - if not self.partition_gradients: - required_version = pkg_version.parse("0.3.17") - error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \ - "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \ - "please use an older version of DeepSpeed (<= 0.5.8) and set 'legacy_stage1': true in your zero config json." - assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}" - - def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): - r"""Loading ZeRO checkpoint - - Arguments: - state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. - Note that the number of saved partitions may differ from number of loading partitions to support - changing GPU count, specifically DP world size, between saving and loading checkpoints. - load_optimizer_states: Boolean indicating whether or not to load base optimizer states - load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 - copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). - """ - """ - Loads a state_dict created by an earlier call to state_dict(). - If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, - whose parameters in turn came from ``model``, it is expected that the user - will call ``model.load_state_dict()`` before - ``fp16_optimizer_instance.load_state_dict()`` is called. - Example:: - model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) - ... - checkpoint = torch.load("saved.pth") - model.load_state_dict(checkpoint['model']) - optimizer.load_state_dict(checkpoint['optimizer']) - """ - - # I think it should actually be ok to reload the optimizer before the model. - dp_rank = dist.get_rank(group=self.zp_process_group) - current_rank_sd = state_dict_list[dp_rank] - self._load_global_state(current_rank_sd) - - ckpt_is_rigid = isinstance(current_rank_sd[BASE_OPTIMIZER_STATE], dict) - - # padding is always at the last rank/partition - # if DP=1024 and param-group elems=16 -> padding will be 1024-16 across all but one rank - # scenario-1 (shrink): saving w. 4 gpus -> loading w. 2 gpus - # scenario-2 (expand): saving w. 2 gpus -> loading w. 4 gpus - # if load_optimizer_states: - # if new_dp_size: - # self.strip_padding() - # self.add_padding_w_new_dp_size() - # self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) - - if load_optimizer_states: - if ckpt_is_rigid: - # loading rigid ckpt into either rigid or elastic exec - self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) - else: - if self.elastic_checkpoint: - # loading elastic into elastic exec - self._restore_elastic_base_optimizer_state(state_dict_list) - else: - # loading an elastic checkpoint into rigid exec - self._restore_base_optimizer_state(current_rank_sd[BASE_OPTIMIZER_STATE]) - - # At this point, the optimizer's references to the model's fp32 parameters are up to date. - # The optimizer's hyperparameters and internal buffers are also up to date. - # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still - # out of date. There are two options. - # 1: Refresh the master params from the model's fp16 params. - # This requires less storage but incurs precision loss. - # 2: Save and restore the fp32 master copies separately. - # We choose option 1 if changing DP degree and option 2 otherwise. - # - # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device - # of their associated parameters, because it's possible those buffers might not exist yet in - # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been - # constructed in the same way as the one whose state_dict we are loading, the same master params - # are guaranteed to exist, so we can just copy_() from the saved master params. - - if load_from_fp32_weights: - # option 2 from above - if self.elastic_checkpoint and not ckpt_is_rigid: - self._restore_from_elastic_fp32_weights(state_dict_list) - else: - # For non-elastic checkpoint, simply copying from saved weights of current rank is sufficient. - for current, saved in zip(self.single_partition_of_fp32_groups, - current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): - src_tensor = _get_padded_tensor(saved, current.numel()) - current.data.copy_(src_tensor.data) - else: - # option 1 from above - self._restore_from_bit16_weights() - - if load_optimizer_states: - self._link_all_hp_params() - - -def _handle_overflow(cpu_sum, x, i): - import math - rank = dist.get_rank() - if rank == 0: - t_i = -1 - for v_i, v in enumerate(x.data.contiguous().view(-1)): - if not math.isfinite(float(v)): - t_i = v_i - break - logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") - - -def estimate_zero2_model_states_mem_needs(total_params, - num_gpus_per_node=1, - num_nodes=1, - cpu_offload=True, - additional_buffer_factor=1.5): - total_gpus = num_nodes * num_gpus_per_node - - if cpu_offload: - gpu_mem = 2 * total_params - cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor - else: - gpu_mem = 4 * total_params + int(16 * total_params / total_gpus) - cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor - - return int(cpu_mem), int(gpu_mem) - - -def model_to_params(model): - # shared params calculated only once - total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) - return total_params - - -def estimate_zero2_model_states_mem_needs_all_live(model, - num_gpus_per_node=1, - num_nodes=1, - additional_buffer_factor=1.5): - """ - Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients - for a given ``model`` and hardware setup. - - If you have an actual model object, use this function and everything will be derived - automatically. - - If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass - the ``total_params`` explicitly. - - Args: - - ``model``: ``nn.Module`` object - - ``num_gpus_per_node``: how many gpus per node (defaults to 1) - - ``num_nodes``: how many nodes (defaults to 1), - - ``additional_buffer_factor``: estimation factor (defaults to 1.5): - - """ - - total_params = model_to_params(model) - - estimate_zero2_model_states_mem_needs_all_cold(total_params=total_params, - num_gpus_per_node=num_gpus_per_node, - num_nodes=num_nodes, - additional_buffer_factor=additional_buffer_factor) - - -def estimate_zero2_model_states_mem_needs_all_cold(total_params, - num_gpus_per_node=1, - num_nodes=1, - additional_buffer_factor=1.5): - """ - Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients - for a given ``model`` and hardware setup. - - If it's a hypothetical model, use this function where you have to pass - the ``total_params`` and ``largest_layer_params`` explicitly. - - If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything - will be derived automatically. - - Args: - - ``total_params``: total model params - - ``num_gpus_per_node``: how many gpus per node (defaults to 1) - - ``num_nodes``: how many nodes (defaults to 1), - - ``additional_buffer_factor``: estimation factor (defaults to 1.5): - - """ - - def format_options(cpu_offload): - enabled = [] - device = f'{OffloadDeviceEnum.cpu:4}' if cpu_offload else "none" - enabled.append(f"offload_optimizer={device}") - return ", ".join(enabled) - - nodes_str = "nodes" if num_nodes > 1 else "node" - gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" - print("Estimated memory needed for params, optim states and gradients for a:\n" - f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" - f"SW: Model with {int(total_params / 1e6)}M total params.") - print(" per CPU | per GPU | Options") - for cpu_offload in [True, False]: - cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params=total_params, - num_gpus_per_node=num_gpus_per_node, - num_nodes=num_nodes, - cpu_offload=cpu_offload, - additional_buffer_factor=additional_buffer_factor) - - options_str = format_options(cpu_offload=cpu_offload) - print(f" {cpu_mem / 2 ** 30:7.2f}GB | {gpu_mem / 2 ** 30:6.2f}GB | {options_str}") diff --git a/opensora/adaptor/utils.py b/opensora/adaptor/utils.py deleted file mode 100644 index c32ff6efc..000000000 --- a/opensora/adaptor/utils.py +++ /dev/null @@ -1,1042 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/Megatron - -Helper functions and classes from multiple sources. -""" - -from collections.abc import Iterable -from deepspeed.moe.utils import is_moe_param -import os -import psutil -import gc -from math import sqrt -from packaging import version as pkg_version - -import torch -from deepspeed import comm as dist - -try: - from torch._six import inf -except ModuleNotFoundError: - from torch import inf - -from deepspeed.utils import groups, logger -from deepspeed.runtime.constants import PIPE_REPLICATED -from numpy import prod -from deepspeed.accelerator import get_accelerator - -from deepspeed.module_inject.policy import transpose -from torch.nn import functional as F - -torch_memory_reserved = get_accelerator().memory_reserved -torch_max_memory_reserved = get_accelerator().max_memory_reserved - - -class DummyOptim(): - """ - Dummy optimizer presents model parameters as a param group, this is - primarily used to allow ZeRO-3 without an optimizer - """ - - def __init__(self, params): - self.param_groups = [] - self.param_groups.append({'params': params}) - - -graph_cache = {} - - -def graph_process(replay_first_step, func, *args, **kwargs): - # `func` should only contain operations on the GPU - # Please ensure that the memory address of the data required by 'func' remains constant - if func.__name__ not in graph_cache: - cuda_stream = get_accelerator().Stream() - cuda_stream.wait_stream(get_accelerator().current_stream()) - with get_accelerator().stream(cuda_stream): - func(*args, **kwargs) - get_accelerator().current_stream().wait_stream(cuda_stream) - graph_cache[func.__name__] = get_accelerator().create_graph() - with get_accelerator().capture_to_graph(graph_cache[func.__name__]): - func(*args, **kwargs) - if replay_first_step: - get_accelerator().replay_graph(graph_cache[func.__name__]) - else: - get_accelerator().replay_graph(graph_cache[func.__name__]) - - -def noop_decorator(func): - return func - - -class noop_context(object): - - def __init__(self): - pass - - def __enter__(self): - pass - - def __exit__(self, exc_type, exc_val, exc_tb): - pass - - -def ensure_directory_exists(filename): - """Create the directory path to ``filename`` if it does not already exist. - - Args: - filename (str): A file path. - """ - dirname = os.path.dirname(filename) - os.makedirs(dirname, exist_ok=True) - - -def set_random_seed(seed): - """Set the random seed for common PRNGs used during training: random, numpy, and torch. - - Args: - seed (int): the seed to use - """ - import numpy - import random - random.seed(seed) - numpy.random.seed(seed) - torch.manual_seed(seed) - - -def is_model_parallel_parameter(p) -> bool: - if hasattr(p, 'model_parallel') and p.model_parallel: - return True - - if hasattr(p, 'tensor_model_parallel') and p.tensor_model_parallel: - return True - - return False - - -def bwc_tensor_model_parallel_rank(mpu=None): - """Backwards-compatible way of querying the tensor model parallel rank from - an ``mpu`` object. - - *Tensor* model parallelism means that tensors are physically split across - processes. This contrasts with *pipeline* model parallelism, in which the - layers are partitioned but tensors left intact. - - The API for tensor model parallelism has changed across versions and this - helper provides a best-effort implementation across versions of ``mpu`` - objects. The preferred mechanism is - ``mpu.get_tensor_model_parallel_rank()``. - - This should "just work" with both Megatron-LM and DeepSpeed's pipeline - parallelism. - - Args: - mpu (model parallel unit, optional): The tensor model parallel rank. - If ``mpu=None``, returns 0. Defaults to ``None``. - - Returns: - int: the rank - """ - if mpu is None: - # No model parallelism in easy :) - return 0 - - if hasattr(mpu, 'get_tensor_model_parallel_rank'): - # New Megatron and DeepSpeed convention (post pipeline-parallelism release) - return mpu.get_tensor_model_parallel_rank() - elif hasattr(mpu, 'get_slice_parallel_rank'): - # Some DeepSpeed + pipeline parallelism versions - return mpu.get_slice_parallel_rank() - else: - # Deprecated Megatron and DeepSpeed convention - return mpu.get_model_parallel_rank() - - -def copy_to_device(item, device, criterion_func): - """ - Return a copy of tensor on specified device. - Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. - Parameters: - item: tensor to copy or (possibly nested) container of tensors to copy. - device: target device - criterion_func: Function to restrict copy operation to items meet criterion - - Returns: - None - """ - if criterion_func(item): - return item.to(device) - elif isinstance(item, list): - return [copy_to_device(v, device, criterion_func) for v in item] - elif isinstance(item, tuple): - return tuple([copy_to_device(v, device, criterion_func) for v in item]) - elif isinstance(item, dict): - return {k: copy_to_device(v, device, criterion_func) for k, v in item.items()} - else: - return item - - -def move_to_device(item, device, criterion_func): - """ - Move tensor on to specified device by changing the storage. - Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. - Parameters: - item: tensor to move or (possibly nested) container of tensors to move. - device: target device - criterion_func: Function to restrict move operation to items meet criterion - - Returns: - None - """ - if criterion_func(item): - device_copy = item.to(device) - item.data = device_copy.data - return item - elif isinstance(item, list): - return [move_to_device(v, device, criterion_func) for v in item] - elif isinstance(item, tuple): - return tuple([move_to_device(v, device, criterion_func) for v in item]) - elif isinstance(item, dict): - return {k: move_to_device(v, device, criterion_func) for k, v in item.items()} - else: - return item - - -class CheckOverflow(object): - '''Checks for overflow in gradient across parallel process''' - - def __init__(self, param_groups=None, mpu=None, zero_reduce_scatter=False, deepspeed=None): - self.mpu = mpu - self.params = [] if param_groups else None - self.zero_reduce_scatter = zero_reduce_scatter - self.deepspeed = deepspeed - self.has_moe_params = False - if param_groups: - for group in param_groups: - for param in group: - self.params.append(param) - if is_moe_param(param): - self.has_moe_params = True - - def check_using_norm(self, norm_group, reduce_overflow=True): - # TODO: I don't think reduce_overflow is needed if mpu is None - overflow = -1 in norm_group - overflow_gpu = get_accelerator().FloatTensor([overflow]) - if self.has_moe_params: - # In this case, we need to do an all_reduce across - # the expert_parallel_group, so that if there was - # an overflow due to expert weights, we detect it - - # Only need to check groups.get_largest_expert_parallel_group() - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) - if self.mpu is not None: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) - elif reduce_overflow: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX) - dist.barrier() - overflow = overflow_gpu[0].item() - return bool(overflow) - - def check(self, param_groups=None): - params = [] - has_moe_params = False - if param_groups is None: - params = self.params - has_moe_params = self.has_moe_params - else: - assert param_groups is not None, \ - "self.params and param_groups both cannot be none" - - for group in param_groups: - for param in group: - params.append(param) - if is_moe_param(param): - has_moe_params = True - - return self.has_overflow(params, has_moe_params=has_moe_params) - - # `params` is a list / generator of torch.Variable - def has_overflow_serial(self, params): - for i, p in enumerate(params): - if p.grad is not None and self._has_inf_or_nan(p.grad.data, i): - return True - return False - - def has_overflow(self, params, has_moe_params=None): - if has_moe_params is None: - has_moe_params = self.has_moe_params - overflow = self.has_overflow_serial(params) - # Since each model parallel GPU carries only part of the model, - # make sure overflow flag is synced across all the model parallel GPUs - overflow_gpu = get_accelerator().ByteTensor([overflow]) - # deepspeed.comm.all_reduce(overflow_gpu, - # op=deepspeed.comm.ReduceOp.MAX, - # group=mpu.get_model_parallel_group()) - if has_moe_params: - # All reduce this across expert_parallel_group, so that if an expert - # overflows, we detect it here - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) - if self.zero_reduce_scatter: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) - elif self.mpu is not None: - if self.deepspeed is not None: - using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - if (using_pipeline and self.deepspeed.pipeline_enable_backward_allreduce is False) or ( - not using_pipeline and self.deepspeed.enable_backward_allreduce is False): - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_data_parallel_group()) - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) - elif self.deepspeed is not None and self.deepspeed.enable_backward_allreduce is False: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) - - overflow = overflow_gpu[0].item() - return bool(overflow) - - # `x` is a torch.Tensor - @staticmethod - def _has_inf_or_nan(x, i): - try: - # if x is half, the .float() incurs an additional deep copy, but it's necessary if - # Pytorch's .sum() creates a one-element tensor of the same type as x - # (which is true for some recent version of pytorch). - cpu_sum = float(x.float().sum()) - # More efficient version that can be used if .sum() returns a Python scalar - # cpu_sum = float(x.sum()) - except RuntimeError as instance: - # We want to check if inst is actually an overflow exception. - # RuntimeError could come from a different error. - # If so, we still want the exception to propagate. - if "value cannot be converted" not in instance.args[0]: - raise - return True - else: - if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: - return True - return False - - -def _handle_overflow(cpu_sum, x, i): - import math - rank = dist.get_rank() - if rank == 0: - t_i = -1 - for v_i, v in enumerate(x.data.contiguous().view(-1)): - if not math.isfinite(float(v)): - t_i = v_i - break - logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") - - -def get_global_norm(norm_list): - """ Compute total from a list of norms - """ - total_norm = 0.0 - for norm in norm_list: - total_norm += norm**2.0 - # logger.info(f'norm_list = {norm_list} global = {sqrt(total_norm)}') - return sqrt(total_norm) - - -def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None): - """Clips gradient norm of an iterable of parameters. - - This has been adapted from Nvidia megatron. We add norm averaging - to consider MoE params when calculating norm as they will result - in different norms across different ranks. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - max_norm (float or int): max norm of the gradients - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) - max_norm = float(max_norm) - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(p.grad.data.abs().max() for p in parameters) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - # Take max across all GPUs. - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item() - else: - total_norm = 0 - for p in parameters: - if mpu is not None: - if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p): - param_norm = p.grad.data.norm(norm_type) - total_norm += param_norm.item()**norm_type - else: - param_norm = p.grad.data.float().norm(norm_type) - total_norm += param_norm.item()**norm_type - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item()**(1. / norm_type) - - # Need to average total_norm across different GPUs due to the presence of moe params - pg = groups._get_data_parallel_group() - scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg)) - - scaled_norm_tensor = get_accelerator().FloatTensor([float(scaled_norm)]) - dist.all_reduce(scaled_norm_tensor, group=pg) - total_norm = scaled_norm_tensor.item() - - clip_coef = max_norm / (total_norm + 1e-6) - if clip_coef < 1: - for p in parameters: - p.grad.data.mul_(clip_coef) - return total_norm - - -def get_grad_norm(parameters, norm_type=2, mpu=None): - """Get grad norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. Taken from Nvidia Megatron. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) - - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(p.grad.data.abs().max() for p in parameters) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - # Take max across all GPUs. - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item() - else: - total_norm = 0. - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) - for p in parameters: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - # Filter to avoid over-counting replicated tensors from tensor - # model parallelism - if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): - continue - - param_norm = p.grad.data.float().norm(norm_type) - total_norm += param_norm.item()**norm_type - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item()**(1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - -def get_grad_zeros(parameters, mpu=None): - """Compute the number of grads with zero values. - - This is adapted from get_grad_norm - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - - Returns: - Total number of params with zero values (viewed as a single vector). - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) - - total_zeros = 0. - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) - for p in parameters: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - # Filter to avoid over-counting replicated tensors from tensor - # model parallelism - if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): - continue - - count_zeros = p.grad.numel() - torch.count_nonzero(p.grad) - total_zeros += count_zeros.item() - - # Sum across all model parallel GPUs. - total_zeros_cuda = get_accelerator().FloatTensor([float(total_zeros)]) - if mpu is not None: - dist.all_reduce(total_zeros_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_zeros = total_zeros_cuda[0].item() - - return total_zeros - - -def get_weight_norm(parameters, norm_type=2, mpu=None): - """Get norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. Taken from Nvidia Megatron. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - -1 if the norm value is NaN or Inf. - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(p.data.abs().max() for p in parameters) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - # Take max across all GPUs. - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item() - else: - total_norm = 0. - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) - for p in parameters: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - # Filter to avoid over-counting replicated tensors from tensor - # model parallelism - if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): - continue - - param_norm = p.data.float().norm(norm_type) - total_norm += param_norm**norm_type - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item()**(1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - -def prefix_sum_inc(weights): - """ Compute an inclusive prefix sum. - - Example: - >>> prefix_sum_inc([3,4,5]) - [3, 7, 12] - """ - weights_ = [w for w in weights] - for x in range(1, len(weights_)): - weights_[x] += weights_[x - 1] - return weights_ - - -def partition_uniform(num_items, num_parts): - import numpy - parts = [0] * (num_parts + 1) - # First check for the trivial edge case - if num_items <= num_parts: - for p in range(num_parts + 1): - parts[p] = min(p, num_items) - return parts - - chunksize = num_items // num_parts - residual = num_items - (chunksize * num_parts) - - parts = numpy.arange(0, (num_parts + 1) * chunksize, chunksize) - - for i in range(residual): - parts[i + 1:] += 1 - parts = parts.tolist() - - return parts - - -def partition_balanced(weights, num_parts): - """ - use dynamic programming solve `The Linear Partition Problem`. - see https://www8.cs.umu.se/kurser/TDBAfl/VT06/algorithms/BOOK/BOOK2/NODE45.HTM - """ - import numpy as np - n = len(weights) - m = num_parts - - if n <= m: - return partition_uniform(n, m) - - dp_max = np.full((n + 1, m + 1), np.inf) - dp_min = np.full((n + 1, m + 1), np.inf) - dp_cost = np.full((n + 1, m + 1), np.inf) - position = np.zeros((n + 1, m + 1), dtype=int) - prefix_sum = np.zeros((n + 1)) - prefix_sum[1:] = np.cumsum(weights) - - dp_max[0, 0] = 0 - dp_cost[0, 0] = 0 - for i in range(1, n + 1): - for j in range(1, min(i, m) + 1): - for k in range(i): - max_sum = max(dp_max[k, j - 1], prefix_sum[i] - prefix_sum[k]) - min_sum = min(dp_min[k, j - 1], prefix_sum[i] - prefix_sum[k]) - cost = max_sum - min_sum - if dp_cost[i, j] >= cost: - dp_cost[i, j] = cost - dp_max[i, j] = max_sum - dp_min[i, j] = min_sum - position[i, j] = k - - parts = [n] - for i in reversed(range(1, m + 1)): - parts.append(position[parts[-1], i]) - parts.reverse() - - return parts - - -class PartitionedTensor: - - def __init__(self, tensor, group, partition_meta=None): - super().__init__() - - self.group = group - self.num_parts = dist.get_world_size(group=self.group) - self.rank = dist.get_rank(group=self.group) - self.orig_size = list(tensor.size()) - self.orig_device = tensor.device - self.local_data, self.partition = self._partition_tensor(tensor) - self.even_split = tensor.numel() % self.num_parts == 0 - - @classmethod - def from_meta(cls, meta, local_part, group, device=get_accelerator().device_name()): - assert meta.dtype == torch.long - dummy = torch.ones(dist.get_world_size(group=group)) - part_obj = cls(tensor=dummy, group=group) - - meta = meta.tolist() - - # [N, list0, ..., listN-1] - part_obj.orig_size = meta[1:(1 + meta[0])] - meta = meta[1 + meta[0]:] - - part_obj.orig_device = device - part_obj.local_data = local_part.detach() - - part_obj.group = group - - # Partition is encoded like the rowptr of a CSR matrix: - # [num_parts, rank, 0, part_1, ..., part_num_parts] - # TODO: support shuffle between different partition granularities - assert part_obj.num_parts == meta[0] - assert part_obj.rank == meta[1] - part_obj.partition = meta[2:] # length num_parts+1 - - return part_obj - - def _partition_tensor(self, tensor): - partition = partition_uniform(num_items=tensor.numel(), num_parts=self.num_parts) - start = partition[self.rank] - length = partition[self.rank + 1] - start - tensor_part = tensor.detach().contiguous().view(-1).narrow(0, start=start, length=length).clone() - - return tensor_part, partition - - def full(self, device=None): - if device is None: - device = self.orig_device - - # Allocate the full tensor as a flat buffer. - full_numel = prod(self.full_size()) - flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device) - if self.even_split: - # Collect the full tensor - dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group) - else: - for part_id in range(self.num_parts): - part_size = self.partition[part_id + 1] - self.partition[part_id] - buf = flat_tensor.narrow(0, start=self.partition[part_id], length=part_size) - if part_id == self.rank: - buf.copy_(self.local_data) - dist.broadcast(buf, part_id, self.group) - return flat_tensor.view(self.full_size()).clone().detach() - - def to_meta(self): - """Returns a torch.LongTensor that encodes partitioning information. - - Can be used along with ``data()`` to serialize a ``PartitionedTensor`` for - communication. - - Returns: - torch.LongTensor: a tensor encoding the meta-information for the partitioning - """ - meta = [] - meta.append(len(self.orig_size)) - meta += list(self.orig_size) - meta.append(self.num_parts) - meta.append(self.rank) - meta += self.partition - return torch.LongTensor(data=meta).to(self.orig_device) - - def data(self): - return self.local_data - - def local_size(self): - return self.local_data.size() - - def full_size(self): - return self.orig_size - - -mem_alloced = 0 -mem_cached = 0 - - -def memory_status(msg, print_rank=-1, reset_max=False): - global mem_alloced, mem_cached - - rank = dist.get_rank() - if print_rank != -1 and rank != print_rank: - return - - get_accelerator().synchronize() - - if reset_max: - get_accelerator().reset_max_memory_cached() - get_accelerator().reset_max_memory_allocated() - - new_alloced = get_accelerator().memory_allocated() - new_cached = get_accelerator().memory_cached() - - delta_alloced = new_alloced - mem_alloced - delta_cached = new_cached - mem_cached - - mem_cached = new_cached - mem_alloced = new_alloced - - max_alloced = get_accelerator().max_memory_allocated() - max_cached = get_accelerator().max_memory_cached() - - # convert to GB for printing - new_alloced /= 1024**3 - new_cached /= 1024**3 - delta_alloced /= 1024**3 - delta_cached /= 1024**3 - max_alloced /= 1024**3 - max_cached /= 1024**3 - - print( - f'RANK={rank} MEMSTATS', msg, f'device={get_accelerator().current_device_name()} ' - f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' - f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') - - -def get_ma_status(): - if dist.is_initialized() and not dist.get_rank() == 0: - return 0 - return get_accelerator().memory_allocated() - - -def empty_cache(): - get_accelerator().empty_cache() - get_accelerator().reset_peak_memory_stats() - - -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 call_to_str(base, *args, **kwargs): - """Construct a string representation of a call. - - Args: - base (str): name of the call - args (tuple, optional): args to ``base`` - kwargs (dict, optional): kwargs supplied to ``base`` - - Returns: - str: A string representation of base(*args, **kwargs) - """ - name = f'{base}(' - if args: - name += ', '.join(repr(arg) for arg in args) - if kwargs: - name += ', ' - if kwargs: - name += ', '.join(f'{key}={repr(arg)}' for key, arg in kwargs.items()) - name += ')' - return name - - -def get_only_unique_item(items): - item_set = set(items) - if len(item_set) != 1: - raise RuntimeError(f"expected there to be only one unique element in {items}") - unique_item, = item_set - - return unique_item - - -def clip_gradients(parameters, max_norm=1.0, global_grad_norm=None, mpu=None, eps=1e-6): - """Clip the gradient of a list of parameters. - Args: - parameters: List of parameters whose .grad will be clipped. - global_grad_norm (float, optional): Precomputed gradient norm. Defaults to None. - mpu (optional): model parallelism unit. Defaults to None. - eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6 - Returns: - float: the global gradient norm - """ - if global_grad_norm is None: - global_grad_norm = get_grad_norm(parameters, mpu=mpu) - clip_coef = max_norm / (global_grad_norm + eps) - if clip_coef < 1: - for p in parameters: - p.grad.detach().mul_(clip_coef) - return global_grad_norm - - -def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, use_graph=False): - """Get norm of an iterable of tensors. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Taken from Nvidia Megatron. - - Arguments: - input_tensors (Iterable[Tensor]): an iterable of Tensors will have norm computed - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the tensors (viewed as a single vector). - """ - assert isinstance(input_tensors, Iterable), f'expected Iterable type not {type(input_tensors)}' - assert all([torch.is_tensor(t) for t in input_tensors]), f'expected list of only tensors' - - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(t.data.abs().max() for t in input_tensors) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item() - else: - if use_graph: - if 'norm_tensors_compute_buffer' not in graph_cache: - graph_cache['norm_tensors_compute_buffer'] = [t.data.float().norm(norm_type) for t in input_tensors] - compute_buffer = graph_cache['norm_tensors_compute_buffer'] - - def _norm_tensors(tensor_list, _compute_buffer, _norm_type): - for i, t in enumerate(tensor_list): - _compute_buffer[i].data.copy_(t.data.float().norm(_norm_type)**_norm_type) - if i != 0: - _compute_buffer[0].data.add_(_compute_buffer[i].data) - - graph_process(False, _norm_tensors, input_tensors, compute_buffer, norm_type) - - total_norm = compute_buffer[0] - else: - total_norm = sum([t.data.float().norm(norm_type).item()**norm_type for t in input_tensors]) - - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]).detach() - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item()**(1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - -def clip_tensors_by_global_norm(input_tensors, max_norm=1.0, global_norm=None, mpu=None, eps=1e-6, use_graph=False): - """Clip list of tensors by global norm. - Args: - input_tensors: List of tensors to be clipped - global_norm (float, optional): Precomputed norm. Defaults to None. - mpu (optional): model parallelism unit. Defaults to None. - eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6 - Returns: - float: the global norm - """ - if global_norm is None: - global_norm = get_global_norm_of_tensors(input_tensors, mpu=mpu, use_graph=use_graph) - clip_coef = max_norm / (global_norm + eps) - if clip_coef < 1: - if use_graph: - - def clip_tensors(_tensor_list, _clip_coef_tensor): - for t in _tensor_list: - t.detach().mul_(_clip_coef_tensor) - - if 'clip_coef_tensor' not in graph_cache: - # Alloc memory - graph_cache['clip_coef_tensor'] = torch.tensor(clip_coef, - dtype=torch.float32).to(get_accelerator().device_name()) - clip_coef_tensor = graph_cache['clip_coef_tensor'] - clip_coef_tensor.copy_(torch.tensor(clip_coef, dtype=torch.float32)) - graph_process(False, clip_tensors, input_tensors, clip_coef_tensor) - - else: - for t in input_tensors: - t.detach().mul_(clip_coef) - return global_norm - - -def align_dense_tensors(tensor_list, alignment): - num_elements = sum(t.numel() for t in tensor_list) - remaining = num_elements % alignment - - if remaining: - elements_to_add = alignment - remaining - pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) - padded_tensor_list = tensor_list + [pad_tensor] - else: - padded_tensor_list = tensor_list - - return padded_tensor_list - - -def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, zp_process_group, dp_process_group=None): - - for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)): - partition_id = dist.get_rank(group=zp_process_group[group_id]) - dp_world_size = dist.get_world_size(group=dp_process_group) - if dp_world_size == 1: - # no groups share optimizer states - # pipeline parallel with bf16 will default call this even if dp size = 1. - continue - # print("call contiguous for all_gather_into_tensor_dp_groups") - dist.all_gather_into_tensor(group_flat, partitioned_params[partition_id].contiguous(), dp_process_group) - - -def all_gather_dp_groups(groups_flat, partitioned_param_groups, zp_process_group, start_alignment_factor, - allgather_bucket_size, dp_process_group=None): - # if dist.has_all_gather_into_tensor(): - return all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, zp_process_group, dp_process_group) - - # for group_id, partitioned_params in enumerate(partitioned_param_groups): - # # Sequential AllGather Best of both worlds - # partition_id = dist.get_rank(group=dp_process_group[group_id]) - # dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) - # - # if dp_world_size == 1: - # # no groups share optimizer states - # # pipeline parallel with bf16 will default call this even if dp size = 1. - # continue - # num_shards = max(1, partitioned_params[partition_id].numel() * dp_world_size // allgather_bucket_size) - # - # shard_size = partitioned_params[partition_id].numel() // num_shards - # - # # Enforce nccl/rccl alignment of start location of each shard - # shard_size = shard_size - (shard_size % start_alignment_factor) - # - # num_elements = shard_size - # - # assert shard_size * num_shards <= partitioned_params[partition_id].numel() - # - # for shard_id in range(num_shards): - # - # if shard_id == (num_shards - 1): - # num_elements = partitioned_params[partition_id].numel() - shard_id * shard_size - # - # shard_list = [] - # for dp_id in range(dp_world_size): - # curr_shard = partitioned_params[dp_id].narrow(0, shard_id * shard_size, num_elements).detach() - # shard_list.append(curr_shard) - # dist.all_gather(shard_list, shard_list[partition_id].contiguous(), dp_process_group[group_id]) - - -class TLinear(torch.nn.Linear): - - def __init__(self, orig_layer, name=""): - self.name = name - super().__init__(orig_layer.weight.shape[1], orig_layer.weight.shape[0], bias=(orig_layer.bias is not None)) - self.weight.data = transpose(orig_layer.weight.data) - self.bias = orig_layer.bias - self._fwd_func = self._fwd_bias_add if self.bias is not None else self._fwd - - def _fwd(self, input): - return F.linear(input, self.weight) - - def _fwd_bias_add(self, input): - return F.linear(input, self.weight, bias=self.bias) - - def forward(self, input): - return self._fwd_func(input) - - -def get_inactive_params(param_list): - from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus - return [param for param in param_list if (hasattr(param, 'ds_id') and \ - param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] - - -def required_torch_version(min_version=None, max_version=None): - assert min_version or max_version, "Must provide a min_version or max_version argument" - - torch_version = pkg_version.parse(torch.__version__) - - if min_version and pkg_version.parse(str(min_version)) > torch_version: - return False - - if max_version and pkg_version.parse(str(max_version)) < torch_version: - return False - - return True diff --git a/opensora/adaptor/zp_manager.py b/opensora/adaptor/zp_manager.py deleted file mode 100644 index e05bfdbcd..000000000 --- a/opensora/adaptor/zp_manager.py +++ /dev/null @@ -1,31 +0,0 @@ -import torch -import os -import torch.distributed as dist - - -class ZPManager(object): - def __init__(self, zp_size=8): - self.rank = int(os.getenv('RANK', '0')) - self.world_size = int(os.getenv("WORLD_SIZE", '1')) - self.zp_size = zp_size - self.zp_group = None - self.zp_rank = None - self.is_initialized = False - - def init_group(self): - if self.is_initialized: - return - - self.is_initialized = True - - """Initialize the sequence parallel group.""" - num_zp_groups: int = self.world_size // self.zp_size - for i in range(num_zp_groups): - ranks = range(i * self.zp_size, (i + 1) * self.zp_size) - group = dist.new_group(ranks) - if self.rank in ranks: - self.zp_group = group - self.zp_rank = self.rank % self.zp_size - - -zp_manager = ZPManager() diff --git a/opensora/sample/sample_t2v_sp.py b/opensora/sample/sample_t2v_sp.py index 5a593320d..a37127299 100644 --- a/opensora/sample/sample_t2v_sp.py +++ b/opensora/sample/sample_t2v_sp.py @@ -17,7 +17,6 @@ from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderM import os, sys -from opensora.adaptor.modules import replace_with_fp32_forwards from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper from opensora.models.diffusion.udit.modeling_udit import UDiTT2V from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh index b811e1bd6..c04e2b22c 100644 --- a/t2v_debug_multi.sh +++ b/t2v_debug_multi.sh @@ -36,7 +36,6 @@ accelerate launch \ --use_image_num 0 \ --tile_overlap_factor 0.125 \ --snr_gamma 5.0 \ - --use_ema \ --ema_start_step 0 \ --cfg 0.1 \ --noise_offset 0.02 \