diff --git a/README.md b/README.md index c602f90144bf..4b752f021787 100755 --- a/README.md +++ b/README.md @@ -96,6 +96,12 @@ If you would like to pre-install any of the DeepSpeed extensions/ops (instead of JIT compiling) or install pre-compiled ops via PyPI please see our [advanced installation instructions](https://www.deepspeed.ai/tutorials/advanced-install/). +On Windows you can build wheel with following steps, currently only inference mode is supported. +1. Install pytorch, such as pytorch 1.8 + cuda 11.1 +2. Install visual cpp build tools, such as VS2019 C++ x64/x86 build tools +3. Launch cmd console with Administrator privilege for creating required symlink folders +4. Run `python setup.py bdist_wheel` to build wheel in `dist` folder + # Features Below we provide a brief feature list, see our detailed [feature overview](https://www.deepspeed.ai/features/) for descriptions and usage. diff --git a/csrc/adam/cpu_adam.cpp b/csrc/adam/cpu_adam.cpp index 6726b895f12c..ef756263b6d6 100644 --- a/csrc/adam/cpu_adam.cpp +++ b/csrc/adam/cpu_adam.cpp @@ -322,42 +322,37 @@ int create_adam_optimizer(int optimizer_id, float betta2 = 0.999, float eps = 1e-8, float weight_decay = 0, - bool adamw_mode = true) + bool adamw_mode = true, + bool should_log = false) { auto opt = std::make_shared(alpha, betta1, betta2, eps, weight_decay, adamw_mode); s_optimizers[optimizer_id] = opt; + + if (should_log) { + std::string avx_type = ""; #if defined(__AVX512__) - std::cout << "Adam Optimizer #" << optimizer_id - << " is created with AVX512 arithmetic capability." << std::endl; - printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n", - alpha, - betta1, - betta2, - weight_decay, - (int)adamw_mode); + avx_type = "AVX512"; #else #if defined(__AVX256__) - std::cout << "Adam Optimizer #" << optimizer_id - << " is created with AVX2 arithmetic capability." << std::endl; - printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n", - alpha, - betta1, - betta2, - weight_decay, - (int)adamw_mode); + avx_type = "AVX2"; #else - std::cout << "Adam Optimizer #" << optimizer_id - << " is created with scalar arithmetic capability." << std::endl; - printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n", - alpha, - betta1, - betta2, - weight_decay, - (int)adamw_mode); + avx_type = "scalar"; #endif #endif + + printf("Adam Optimizer #%d is created with %s arithmetic capability.\n", + optimizer_id, + avx_type.c_str()); + printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n", + alpha, + betta1, + betta2, + weight_decay, + (int)adamw_mode); + } + return 0; } diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index 42ec654a7101..08cf1fa9bacf 100644 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -189,12 +189,16 @@ def _pre_forward_hook(self, module, *inputs, **kwargs): if torch.is_tensor(input): input = input.to(torch.cuda.current_device()) if self.mp_world_size > 1: + if not input.is_contiguous(): + input = input.contiguous() dist.broadcast(input, 0) for k in kwargs: if torch.is_tensor(kwargs[k]): kwargs[k] = kwargs[k].to(torch.cuda.current_device()) if self.mp_world_size > 1: + if not kwargs[k].is_contiguous(): + kwargs[k] = kwargs[k].contiguous() dist.broadcast(kwargs[k], 0) def forward(self, *inputs, **kwargs): @@ -210,12 +214,16 @@ def forward(self, *inputs, **kwargs): if torch.is_tensor(input): input = input.to(torch.cuda.current_device()) if self.mp_world_size > 1: + if not input.is_contiguous(): + input = input.contiguous() dist.broadcast(input, 0) for k in kwargs: if torch.is_tensor(kwargs[k]): kwargs[k] = kwargs[k].to(torch.cuda.current_device()) if self.mp_world_size > 1: + if not kwargs[k].is_contiguous(): + kwargs[k] = kwargs[k].contiguous() dist.broadcast(kwargs[k], 0) return self.model_orig_fwd(*inputs, **kwargs) diff --git a/deepspeed/module_inject/replace_policy.py b/deepspeed/module_inject/replace_policy.py index 325bb37efd74..3758ffd9b522 100755 --- a/deepspeed/module_inject/replace_policy.py +++ b/deepspeed/module_inject/replace_policy.py @@ -19,7 +19,7 @@ def attention(self): def get_hidden_heads(self): """ - retun hidden_size and number of heads + return hidden_size and number of heads """ raise NotImplementedError diff --git a/deepspeed/ops/adam/cpu_adam.py b/deepspeed/ops/adam/cpu_adam.py index 35eeedb86b5d..835742a35604 100755 --- a/deepspeed/ops/adam/cpu_adam.py +++ b/deepspeed/ops/adam/cpu_adam.py @@ -7,6 +7,7 @@ import time from pathlib import Path from ..op_builder import CPUAdamBuilder +from deepspeed.utils.logging import should_log_le class DeepSpeedCPUAdam(torch.optim.Optimizer): @@ -83,7 +84,8 @@ def __init__(self, betas[1], eps, weight_decay, - adamw_mode) + adamw_mode, + should_log_le("info")) def __del__(self): # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index f516e4de5203..efe95f91bac5 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -24,7 +24,7 @@ from deepspeed.runtime.config import DeepSpeedConfig from deepspeed.utils import logger -from deepspeed.runtime.utils import move_to_device, see_memory_usage +from deepspeed.runtime.utils import move_to_device, see_memory_usage, bwc_tensor_model_parallel_rank from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers # DeepSpeed Checkpointing Enabled or Disabled @@ -213,9 +213,12 @@ def model_parallel_cuda_manual_seed(seed): model parallel regions. """ global mpu + + tp_rank = bwc_tensor_model_parallel_rank(mpu) + # 2718 is just for fun and any POSITIVE value will work. offset = seed + 2718 - model_parallel_seed = offset + mpu.get_model_parallel_rank() + model_parallel_seed = offset + tp_rank # Data parallel gets the original sedd. data_parallel_seed = seed @@ -225,7 +228,7 @@ def model_parallel_cuda_manual_seed(seed): 'model parallel rank {}, and data parallel rank {} with ' 'model parallel seed: {} and data parallel seed: {}'.format( torch.distributed.get_rank(), - mpu.get_model_parallel_rank(), + tp_rank, mpu.get_data_parallel_rank(), model_parallel_seed, data_parallel_seed), @@ -384,9 +387,14 @@ def save_args_for_backward(*all_args): global data_offsets, size_offsets if mp_rank is None: if mpu is not None: - mp_rank = mpu.get_model_parallel_rank() - mp_size = mpu.get_model_parallel_world_size() - mp_group = mpu.get_model_parallel_group() + if hasattr(mpu, 'get_tensor_model_parallel_rank'): + mp_rank = mpu.get_tensor_model_parallel_rank() + mp_size = mpu.get_tensor_model_parallel_world_size() + mp_group = mpu.get_tensor_model_parallel_group() + else: + mp_rank = mpu.get_model_parallel_rank() + mp_size = mpu.get_model_parallel_world_size() + mp_group = mpu.get_model_parallel_group() else: mp_rank = 0 mp_size = 1 @@ -416,7 +424,7 @@ def save_args_for_backward(*all_args): inputs = [] for i, item in enumerate(args[:-1]): - if not torch.is_tensor(item): + if not torch.is_tensor(item) or mp_size > item.numel(): inputs.append(item) continue diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 28ad2e977f17..a782581dde8d 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -36,8 +36,10 @@ ZERO_OPTIMIZATION_OPTIMIZER_STATES, ZERO_OPTIMIZATION_GRADIENTS, ZERO_OPTIMIZATION_WEIGHTS from deepspeed.runtime.csr_tensor import CSRTensor import deepspeed.runtime.lr_schedules as lr_schedules +from deepspeed.runtime.utils import get_grad_norm from deepspeed.utils import logger, log_dist, init_distributed from deepspeed.utils.timer import ThroughputTimer, SynchronizedWallClockTimer +from deepspeed.utils.debug import debug_extract_module_and_param_names from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop from deepspeed.runtime.eigenvalue import Eigenvalue @@ -121,6 +123,11 @@ def __init__(self, self.block_eigenvalue = None self.gas_boundary_ctr = 0 self.dist_backend = "nccl" + self._step_applied = False + self._global_grad_norm = None + + # for debug purposes - can then debug print: debug_get_module_name(module) + debug_extract_module_and_param_names(model) # Set config using config_params for backwards compat if self.config is None and config_params is not None: @@ -230,6 +237,40 @@ def get_batch_info(self): """ 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 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 checkpoint_tag_validation_enabled(self): return self._config.checkpoint_tag_validation_enabled @@ -475,6 +516,9 @@ def steps_per_print(self): 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 dump_state(self): return self._config.dump_state @@ -562,6 +606,14 @@ def _configure_with_arguments(self, args, mpu): # environment variable is set. We must align args.local_rank to this value for # backwards compatability 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 conficting 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 @@ -581,8 +633,10 @@ def _do_args_sanity_check(self, args): assert args.deepspeed_config is None, "Not sure how to proceed, we were given both a deepscale_config and deepspeed_config" args.deepspeed_config = args.deepscale_config - assert "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.distributed launcher. If using a different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed." + 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.distributed 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 != 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: @@ -768,6 +822,7 @@ def _configure_basic_optimizer(self, model_parameters): 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(): @@ -775,6 +830,7 @@ def _configure_basic_optimizer(self, model_parameters): f'Currently the convergence of 1-bit 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(): @@ -888,6 +944,15 @@ def _configure_zero_optimizer(self, optimizer): gradient_predivide=self.gradient_predivide) elif zero_stage <= ZERO_OPTIMIZATION_GRADIENTS: overlap_comm = self.zero_overlap_comm() + contiguous_gradients = self.zero_contiguous_gradients() + round_robin_gradients = self.zero_round_robin_gradients() + + # Overlap and contiguous grads are meaningless in stage 1 and are ignored + if zero_stage == ZERO_OPTIMIZATION_OPTIMIZER_STATES: + overlap_comm = False + contiguous_gradients = False + round_robin_gradients = False + if isinstance(self.module, PipelineModule): if overlap_comm: logger.warning( @@ -902,7 +967,7 @@ def _configure_zero_optimizer(self, optimizer): 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(), + contiguous_gradients=contiguous_gradients, reduce_bucket_size=self.zero_reduce_bucket_size(), allgather_bucket_size=self.zero_allgather_bucket_size(), dp_process_group=self.data_parallel_group, @@ -914,9 +979,10 @@ def _configure_zero_optimizer(self, optimizer): 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 == ZERO_OPTIMIZATION_GRADIENTS) + partition_grads=zero_stage == ZERO_OPTIMIZATION_GRADIENTS, + round_robin_gradients=round_robin_gradients) elif zero_stage == ZERO_OPTIMIZATION_WEIGHTS: - print("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None + logger.info("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3 optimizer = FP16_DeepSpeedZeroOptimizer_Stage3( self.module, @@ -976,6 +1042,18 @@ def is_iterable_style_dataset(obj): torch.utils.data.IterableDataset ) # hasattr(obj, "__iter__") should work as well + 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, @@ -1113,6 +1191,10 @@ def forward(self, *inputs, **kwargs): return loss def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): + # 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() @@ -1248,6 +1330,9 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}): 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 no overflow if self.quantizer: self.quantizer.quantize( @@ -1270,12 +1355,19 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}): overflow = False if hasattr(self.optimizer, 'overflow'): overflow = self.optimizer.overflow + self._step_applied = not overflow if overflow: self.skipped_steps += 1 else: if self.lr_scheduler is not None: - self.lr_scheduler.step(**(lr_kwargs or {})) + 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(increment=self.train_batch_size()) if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0: self._report_progress(self.global_steps + 1) @@ -1295,6 +1387,8 @@ def step(self, lr_kwargs=None): "init in order to use step" report_progress = self.global_rank == 0 if self.global_rank else True + 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 @@ -1658,9 +1752,12 @@ def load_checkpoint(self, load_lr_scheduler_states=load_lr_scheduler_states) if self.zero_optimization() and load_path is not None: - self._load_zero_checkpoint(load_dir, - tag, - load_optimizer_states=load_optimizer_states) + success = self._load_zero_checkpoint( + load_dir, + tag, + load_optimizer_states=load_optimizer_states) + if not success: + self.optimizer._restore_from_fp16_weights() return load_path, client_states @@ -1730,7 +1827,7 @@ def _load_checkpoint(self, def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) if zero_sd_list is None: - return + return False self.optimizer.load_state_dict( state_dict_list=zero_sd_list, @@ -1739,6 +1836,7 @@ def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): print( f'loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}' ) + return True def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size): zero_ckpt_names = [] @@ -1903,6 +2001,7 @@ def _save_checkpoint(self, save_dir, tag, client_state={}): self._curr_ckpt_path = os.path.join(save_dir, tag) state = dict(module=self.module_state_dict(), + buffer_names=self._get_buffer_names(), optimizer=self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None, lr_scheduler=self.lr_scheduler.state_dict() @@ -1922,6 +2021,27 @@ def _save_checkpoint(self, save_dir, tag, client_state={}): torch.save(state, save_path) self._curr_save_path = None + 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_shapes(self): param_shapes = OrderedDict() for name, param in self.module.named_parameters(): @@ -1935,7 +2055,7 @@ def _copy_recovery_script(self, save_path): 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}") + #logger.info(f"creating recovery script {dst}") copyfile(src, dst) # make executable os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC) @@ -1948,7 +2068,7 @@ def _save_zero_checkpoint(self, save_path, tag): ds_version=version) torch.save(zero_sd, zero_checkpoint_name) self._copy_recovery_script(save_path) - logger.info('zero checkpoint saved {}'.format(zero_checkpoint_name)) + #logger.info('zero checkpoint saved {}'.format(zero_checkpoint_name)) def _zero3_consolidated_fp16_state_dict(self): """ diff --git a/deepspeed/runtime/fp16/fused_optimizer.py b/deepspeed/runtime/fp16/fused_optimizer.py index fba0d6b1fd59..72dd2c161845 100755 --- a/deepspeed/runtime/fp16/fused_optimizer.py +++ b/deepspeed/runtime/fp16/fused_optimizer.py @@ -9,7 +9,7 @@ import math from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from deepspeed.runtime.utils import get_grad_norm, CheckOverflow, get_weight_norm +from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE from deepspeed.utils import logger, log_dist @@ -45,6 +45,8 @@ def __init__(self, self.fp16_groups_flat = [] self.fp32_groups_flat = [] + self._global_grad_norm = 0. + # loop to deal with groups for i, param_group in enumerate(self.optimizer.param_groups): # push this group to list before modify @@ -161,8 +163,11 @@ def step_fused_adam(self, closure=None): "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) return self.overflow + + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + combined_scale = self.unscale_and_clip_grads(grads_groups_flat, - norm_groups, + self._global_grad_norm, apply_scale=False) # norm is in fact norm*cur_scale self.optimizer.step(grads=[[g] for g in grads_groups_flat], @@ -251,8 +256,10 @@ def step(self, closure=None): all_groups_norm = get_grad_norm(self.fp32_groups_flat, mpu=self.mpu) self.stop_timers([COMPUTE_NORM]) + self._global_grad_norm = get_global_norm(norm_list=[all_groups_norm]) + self.start_timers([UNSCALE_AND_CLIP]) - self.unscale_and_clip_grads(grads_groups_flat, [all_groups_norm]) + self.unscale_and_clip_grads(grads_groups_flat, self._global_grad_norm) self.stop_timers([UNSCALE_AND_CLIP]) self.start_timers([BASIC_STEP]) @@ -277,12 +284,7 @@ def step(self, closure=None): return self.overflow - def unscale_and_clip_grads(self, grad_groups_flat, norm_groups, apply_scale=True): - total_norm = 0.0 - for norm in norm_groups: - total_norm += norm**2.0 - total_norm = math.sqrt(total_norm) - + def unscale_and_clip_grads(self, grad_groups_flat, total_norm, apply_scale=True): # compute combined scale factor for this group combined_scale = self.cur_scale if self.clip_grad > 0.: diff --git a/deepspeed/runtime/fp16/unfused_optimizer.py b/deepspeed/runtime/fp16/unfused_optimizer.py index c30df0bef1d0..08a4a7c41cb2 100755 --- a/deepspeed/runtime/fp16/unfused_optimizer.py +++ b/deepspeed/runtime/fp16/unfused_optimizer.py @@ -9,7 +9,7 @@ from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors import math -from deepspeed.runtime.utils import get_grad_norm, CheckOverflow, get_weight_norm +from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE from deepspeed.utils import logger @@ -32,6 +32,7 @@ def __init__(self, fused_lamb_legacy=False): self.fused_lamb_legacy = fused_lamb_legacy + self._global_grad_norm = 0. if torch.distributed.get_rank() == 0: logger.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ') @@ -148,7 +149,9 @@ def step_fused_lamb(self, closure=None): self.cur_scale)) return self.overflow - combined_scale = self.unscale_and_clip_grads(norm_groups, apply_scale=False) + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + combined_scale = self.unscale_and_clip_grads(self._global_grad_norm, + apply_scale=False) self.optimizer.step(grads=grads_groups, output_params=self.fp16_groups, scale=combined_scale) @@ -197,7 +200,8 @@ def step(self, closure=None): else: fp32_param.grad = fp16_param.grad.to(fp32_param.dtype) - self.unscale_and_clip_grads(norm_groups) + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + self.unscale_and_clip_grads(self._global_grad_norm) self.optimizer.step() @@ -212,12 +216,7 @@ def step(self, closure=None): return self.overflow - def unscale_and_clip_grads(self, norm_groups, apply_scale=True): - total_norm = 0.0 - for norm in norm_groups: - total_norm += norm**2.0 - total_norm = math.sqrt(total_norm) - + def unscale_and_clip_grads(self, total_norm, apply_scale=True): # compute combined scale factor for this group combined_scale = self.cur_scale if self.clip_grad > 0.: diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 196cbe8c6217..cadb5b82e36f 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -110,8 +110,9 @@ def __init__(self, *super_args, **super_kwargs): self.is_model_parallel = self.grid.model_parallel_size > 1 # Partition input/output buffers + # XXX temporarily disable while I revert some partition hacks. self.is_pipe_partitioned = self.is_model_parallel - self.is_grad_partitioned = False + self.is_grad_partitioned = self.is_model_parallel model_parameters = filter(lambda p: p.requires_grad, self.module.parameters()) num_params = sum([p.numel() for p in model_parameters]) @@ -393,6 +394,19 @@ def eval_batch(self, data_iter, compute_loss=True, reduce_output='avg'): return eval_output + 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. + """ + super().set_train_batch_size(train_batch_size) + self.micro_batches = self.gradient_accumulation_steps() + def is_first_stage(self): """True if this process is in the first stage in the pipeline.""" return self.stage_id == 0 @@ -553,12 +567,18 @@ def _exec_forward_pass(self, buffer_id): local_part=inputs[1], group=self.grid.get_slice_parallel_group()) + inputs = part_input.full() + inputs.requires_grad = True + part_input = None + self.pipe_buffers['inputs'][buffer_id] = inputs + ''' inputs = tuple([part_input.full(), inputs[2]]) inputs[0].requires_grad = True # skip mask #inputs[1].requires_grad = True part_input = None self.pipe_buffers['inputs'][buffer_id] = inputs + ''' # Zero out the gradients each time we use the tensor because only the data in # tensor changes across batches @@ -568,13 +588,14 @@ def _exec_forward_pass(self, buffer_id): # Partition the outputs if we are not the last stage if self.is_pipe_partitioned and not self.is_last_stage(): - part = PartitionedTensor(tensor=outputs[0], + assert torch.is_tensor(outputs) + part = PartitionedTensor(tensor=outputs, group=self.grid.get_slice_parallel_group()) # Clear the large output data, but save the computation graph - outputs[0].data = torch.zeros(1) - self.pipe_buffers['output_tensors'][buffer_id] = outputs[0] + outputs.data = torch.zeros(1) + self.pipe_buffers['output_tensors'][buffer_id] = outputs # Inject the partitioned tensor into the output before sending - outputs = tuple([part.to_meta(), part.data(), outputs[1]]) + outputs = tuple([part.to_meta(), part.data()]) part = None self.pipe_buffers['outputs'][buffer_id] = outputs @@ -632,15 +653,11 @@ def _exec_backward_pass(self, buffer_id): local_part=outputs[1], group=self.grid.get_slice_parallel_group()) self.pipe_buffers['output_tensors'][buffer_id].data = part_output.full() - outputs = tuple( - [self.pipe_buffers['output_tensors'][buffer_id], - outputs[2]]) + outputs = self.pipe_buffers['output_tensors'][buffer_id] else: # Already restored from partition - self.pipe_buffers['output_tensors'][buffer_id].data = outputs[0] - outputs = tuple( - [self.pipe_buffers['output_tensors'][buffer_id], - outputs[1]]) + self.pipe_buffers['output_tensors'][buffer_id].data = outputs + outputs = self.pipe_buffers['output_tensors'][buffer_id] grad_tensors = self.grad_layer if self.is_grad_partitioned: @@ -649,7 +666,7 @@ def _exec_backward_pass(self, buffer_id): meta=self.grad_layer[0], local_part=self.grad_layer[1], group=self.grid.get_slice_parallel_group()) - grad_tensors = tuple([part_grad.full(), self.grad_layer[2]]) + grad_tensors = part_grad.full() part_grad = None #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') @@ -872,13 +889,10 @@ def _exec_send_grads(self, buffer_id): # Partition the gradient if self.is_grad_partitioned: - part = PartitionedTensor(tensor=inputs[0].grad, + assert torch.is_tensor(inputs) + part = PartitionedTensor(tensor=inputs.grad, group=self.grid.get_slice_parallel_group()) - # Clear the large output data, but save the computation graph - # Inject the partitoned tensor into the output before sending - - # XXX Hack - inputs = tuple([part.to_meta(), part.data(), inputs[1]]) + inputs = tuple([part.to_meta(), part.data()]) # XXX Terrible hack # Drop the attention mask from the input buffer here. It does not have @@ -899,8 +913,6 @@ def _exec_send_grads(self, buffer_id): # First two sends are partitioned gradient p2p.send(inputs[0], self.prev_stage) p2p.send(inputs[1], self.prev_stage) - # XXX hack hack hack - #p2p.send(inputs[2].grad, self.prev_stage) else: for idx, buffer in enumerate(inputs): # Skip tensors that will not produce a grad @@ -974,7 +986,7 @@ def _exec_recv_grads(self, buffer_id): local_part=outputs[1], group=self.grid.get_slice_parallel_group()) outputs[0].data = part_output.full() - outputs = tuple([outputs[0], outputs[2]]) + outputs = outputs[0] # save for backward self.pipe_buffers['outputs'][buffer_id] = outputs @@ -984,7 +996,7 @@ def _exec_recv_grads(self, buffer_id): s = list(outputs.size()) self.grad_layer = self._allocate_buffer(s, num_buffers=1)[0] else: - sizes = [list(t.size()) for t in outputs if t.is_floating_point()] + sizes = [list(t.size()) for t in outputs] # if t.is_floating_point()] self.grad_layer = self._allocate_buffers(sizes, num_buffers=1)[0] if isinstance(self.grad_layer, torch.Tensor): diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py index c4e111e47315..e61dd1c72878 100644 --- a/deepspeed/runtime/pipe/module.py +++ b/deepspeed/runtime/pipe/module.py @@ -253,7 +253,7 @@ def _build(self): # All pipeline parameters should be considered as model parallel in the context # of our FP16 optimizer for p in self.parameters(): - p.model_parallel = True + p.ds_pipe_replicated = False def _count_layer_params(self): """Count the trainable parameters in individual layers. @@ -472,7 +472,7 @@ def _index_tied_modules(self): # Only count the tied module once in the eyes of the FP16 optimizer if self.global_rank != tied_ranks[0]: for p in self.tied_modules[key].parameters(): - p.model_parallel = False + p.ds_pipe_replicated = True ''' if len(tied_comms) > 0: print(f'RANK={self.global_rank} tied_comms={tied_comms}') @@ -559,7 +559,18 @@ def save_state_dict(self, save_dir): model_ckpt_path = self.ckpt_layer_path(save_dir, idx) if not hasattr(layer, 'state_dict'): continue - torch.save(layer.state_dict(), model_ckpt_path) + # We pass cloned tensors to torch.save() to avoid checkpoint bloat which occurs because torch.save() + # saves the underlying storage rather than the slice of the storage corresponding to individual tensors. + # This is a problem in DeepSpeed because we often allocate tensors using slices of large flattened buffers. + # Tensor cloning helps to avoid this problem because the storage of cloned tensors are closer to the true size. + # It is expected that the garbage collector will reclaim the cloned tensor storage to avoid memory bloat. + # See https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing + orig_state_dict = layer.state_dict() + final_state_dict = type(orig_state_dict)( + {k: v.clone() + for k, + v in orig_state_dict.items()}) + torch.save(final_state_dict, model_ckpt_path) def load_state_dir(self, load_dir, strict=True): for idx, layer in enumerate(self.forward_funcs): @@ -577,15 +588,18 @@ def load_state_dir(self, load_dir, strict=True): layer.load_state_dict(checkpoint) - if self._grid.data_parallel_id == 0: - logger.info( - f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}' - ) + # if self._grid.data_parallel_id == 0: + # logger.info( + # f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}' + # ) self._synchronize_tied_weights() def _is_checkpointable(self, funcs): - if self.__class__.__name__ == 'GPT2ModelPipe': + # This is an unfortunate hack related to torch and deepspeed activation checkpoint implementations. + # Some layers like torch.nn.Embedding will not receive grads if checkpointed, which breaks things. + # I presume it's related to the discrete inputs that cannot require_grad? Need to revisit. + if self.__class__.__name__ in ('GPTModelPipe', 'GPT2ModelPipe'): return all('ParallelTransformerLayerPipe' in f.__class__.__name__ for f in funcs) diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py index 24c0f250a4b9..e129d3d5b177 100644 --- a/deepspeed/runtime/pipe/p2p.py +++ b/deepspeed/runtime/pipe/p2p.py @@ -2,11 +2,27 @@ Copyright 2019 The Microsoft DeepSpeed Team ''' +import pickle +import typing + +import torch import torch.distributed as dist +# To query whether we have send/recv support +from packaging.version import Version +from deepspeed.git_version_info import torch_info + _groups = None _grid = None +_async = [] + + +def can_send_recv() -> bool: + torch_version = Version(torch_info['version']) + sendrecv_min = Version('1.8') + return torch_version >= sendrecv_min + #initializes adjacent process groups #run this only after torch.distributed.init_process_group() has been called @@ -16,7 +32,8 @@ def init_process_groups(grid): assert _grid.pipe_parallel_size > 1, "There is no pipeline parallelism" - _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups] + if not can_send_recv(): + _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups] def _is_valid_send_recv(src_stage, dest_stage): @@ -30,40 +47,117 @@ def _is_valid_send_recv(src_stage, dest_stage): def send(tensor, dest_stage, async_op=False): global _groups - - async_op = False + assert async_op == False, "Doesnt support async_op true" src_stage = _grid.get_stage_id() _is_valid_send_recv(src_stage, dest_stage) - group = _get_send_recv_group(src_stage, dest_stage) - src_rank = _grid.stage_to_global(stage_id=src_stage) + dest_rank = _grid.stage_to_global(stage_id=dest_stage) + if async_op: + global _async + op = dist.isend(tensor, dest_rank) + _async.append(op) + else: - return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) + if can_send_recv(): + return dist.send(tensor, dest_rank) + else: + group = _get_send_recv_group(src_stage, dest_stage) + src_rank = _grid.stage_to_global(stage_id=src_stage) + return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) def recv(tensor, src_stage, async_op=False): - global _groups - - async_op = False + assert async_op == False, "Doesnt support async_op true" dest_stage = _grid.get_stage_id() _is_valid_send_recv(src_stage, dest_stage) - group = _get_send_recv_group(src_stage, dest_stage) src_rank = _grid.stage_to_global(stage_id=src_stage) - return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) - - -def barrier(stage_id): - global _groups, _grid - group_id = _grid.stage_to_global(stage_id=stage_id) - if (dist.get_rank() >= 0): - print("Barrier Group ID", group_id) - print("Barrier Group", _grid.p2p_groups[group_id]) - dist.barrier(group=_groups[group_id]) - if (dist.get_rank() >= 0): - print("Exiting Barrier ", group_id) + if async_op: + global _async + op = dist.irecv(tensor, src_rank) + _async.append(op) + else: + if can_send_recv(): + return dist.recv(tensor, src_rank) + else: + group = _get_send_recv_group(src_stage, dest_stage) + return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) + + +def wait(): + global _async + for op in _async: + op.wait() + _async = [] + + torch.cuda.synchronize() + + +def send_obj(msg: typing.Any, dest: int): + """Send an arbitrary python object to ``dest``. + + Note: ``msg`` must be pickleable. + + WARN: This incurs a CPU -> GPU transfer and should be used sparingly + for performance reasons. + + Args: + msg (typing.Any): The object to send. + dest (int): Destination rank. + """ + # serialize the message + msg = pickle.dumps(msg) + # construct a tensor to send + msg = torch.ByteTensor(torch.ByteStorage.from_buffer(msg)).cuda() + + # Send meta and message + length_tensor = torch.tensor([len(msg)], dtype=torch.long).cuda() + dist.send(length_tensor, dst=dest) + dist.send(msg, dst=dest) + + +def recv_obj(sender: int) -> typing.Any: + """Receive an arbitrary python object from ``sender``. + + WARN: This incur a CPU <-> GPU transfers and should be used sparingly + for performance reasons. + + Args: + sender (int): The rank sending the message. + """ + # Get message meta + length = torch.tensor([0], dtype=torch.long).cuda() + dist.recv(length, src=sender) + + # Receive and deserialize + msg = torch.empty(length.item(), dtype=torch.uint8).cuda() + dist.recv(msg, src=sender) + + msg = pickle.loads(msg.cpu().numpy().tobytes()) + + def _to(x): + """Recursively move to the current device.""" + if torch.is_tensor(x): + return x.cuda() + if isinstance(x, (tuple, list)): + ret = [_to(x_) for x_ in x] + if isinstance(x, tuple): + ret = tuple(ret) + return ret + # handle kwargs + if isinstance(x, dict): + ret = dict() + for key, val in x.items(): + ret[_to(key)] = _to(val) + return ret + + # Anything else is a no-op + return x + + msg = _to(msg) + return msg def _get_send_recv_group(src_stage, dest_stage): diff --git a/deepspeed/runtime/state_dict_factory.py b/deepspeed/runtime/state_dict_factory.py index f6be1dbbe57b..f5562e141c3f 100755 --- a/deepspeed/runtime/state_dict_factory.py +++ b/deepspeed/runtime/state_dict_factory.py @@ -51,10 +51,6 @@ def load(self, self.module_key = module_key num_ckpt = len(self.ckpt_list) idx = mp_rank * num_ckpt // mp_world_size - - logger.info( - f'mp_world_size: {mp_world_size}, mp_rank: {mp_rank}, module_key: {module_key}' - ) """ We have multiple cases to handle here for both training and inference: 1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None a. if no mp_size/pp_size resizing occurs, for both training & inference, loading @@ -82,7 +78,7 @@ def load(self, merge_count = 1 if num_ckpt == mp_world_size: assert os.path.exists(load_path) - logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}') + #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}') sd = torch.load(load_path, map_location=lambda storage, loc: storage) if quantize: @@ -162,7 +158,7 @@ def set_module(self, sd, module): return sd def check_ckpt_list(self): - logger.info(f'checkpoint file list: {self.ckpt_list}') + #logger.info(f'checkpoint file list: {self.ckpt_list}') assert len(self.ckpt_list) > 0 sd = torch.load(self.ckpt_list[0], map_location=lambda storage, loc: storage) diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py index c792d2a6d0db..da9727284232 100755 --- a/deepspeed/runtime/utils.py +++ b/deepspeed/runtime/utils.py @@ -9,7 +9,7 @@ import os import psutil import gc -from math import ceil +from math import ceil, sqrt from math import floor from bisect import bisect_left, bisect_right @@ -49,6 +49,48 @@ def set_random_seed(seed): torch.manual_seed(seed) +def is_model_parallel_parameter(p) -> bool: + return hasattr(p, 'model_parallel') and p.model_parallel + + +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 move_to_device(item, device): """ Move tensor onto device. Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. @@ -198,6 +240,15 @@ def _handle_overflow(cpu_sum, x, i): ) +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 + return sqrt(total_norm) + + def get_grad_norm(parameters, norm_type=2, mpu=None): """Clips gradient norm of an iterable of parameters. @@ -231,15 +282,19 @@ def get_grad_norm(parameters, norm_type=2, mpu=None): 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: - if mpu is not None: - if (mpu.get_model_parallel_rank() == 0 - ) or is_model_parallel_parameter(p): - param_norm = p.grad.data.float().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 + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, 'ds_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 = torch.cuda.FloatTensor([float(total_norm)]) @@ -256,6 +311,48 @@ def get_grad_norm(parameters, norm_type=2, mpu=None): 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, 'ds_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 = torch.cuda.FloatTensor([float(total_zeros)]) + if mpu is not None: + torch.distributed.all_reduce(total_zeros_cuda, + op=torch.distributed.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): """Clips gradient norm of an iterable of parameters. @@ -288,24 +385,19 @@ def get_weight_norm(parameters, norm_type=2, mpu=None): 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: - if mpu is not None: - if (mpu.get_model_parallel_rank() == 0 - ) or is_model_parallel_parameter(p): - try: - param_norm = float(torch.norm(p, norm_type, dtype=torch.float32)) - except TypeError as err: - param_norm = float(torch.norm(p.float(), norm_type)) - - #param_norm = p.data.float().norm(norm_type) - total_norm += param_norm**norm_type - else: - try: - param_norm = float(torch.norm(p, norm_type, dtype=torch.float32)) - except TypeError as err: - param_norm = float(torch.norm(p.float(), norm_type)) - #param_norm = p.data.float().norm(norm_type) - total_norm += param_norm**norm_type + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, 'ds_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 = torch.cuda.FloatTensor([float(total_norm)]) diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py index 377ad94549a7..a48dd4e620b4 100755 --- a/deepspeed/runtime/zero/config.py +++ b/deepspeed/runtime/zero/config.py @@ -39,6 +39,7 @@ def __init__(self, param_dict): self.gather_fp16_weights_on_model_save = None self.ignore_unused_parameters = None + self.round_robin_gradients = None if ZERO_OPTIMIZATION in param_dict.keys(): zero_config_dict = param_dict[ZERO_OPTIMIZATION] @@ -184,3 +185,8 @@ def _initialize(self, zero_config_dict): self.legacy_stage1 = get_scalar_param(zero_config_dict, ZERO_OPTIMIZATION_LEGACY_STAGE1, ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT) + + self.round_robin_gradients = get_scalar_param( + zero_config_dict, + ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS, + ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS_DEFAULT) diff --git a/deepspeed/runtime/zero/constants.py b/deepspeed/runtime/zero/constants.py index eaeb2a95ccd9..e0a26d53609f 100755 --- a/deepspeed/runtime/zero/constants.py +++ b/deepspeed/runtime/zero/constants.py @@ -30,7 +30,8 @@ "sub_group_size" : 1000000000000, "offload_param": {...}, "offload_optimizer": {...}, - "ignore_unused_parameters": [true|false] + "ignore_unused_parameters": [true|false], + "round_robin_gradients": [true|false] } } ''' @@ -91,7 +92,7 @@ ZERO_OPTIMIZATION_OFFLOAD_OPTIMIZER_DEFAULT = None ZERO_OPTIMIZATION_SUB_GROUP_SIZE = 'sub_group_size' -ZERO_OPTIMIZATION_SUB_GROUP_SIZE_DEFAULT = 1000000000000 +ZERO_OPTIMIZATION_SUB_GROUP_SIZE_DEFAULT = 1000000000 #maximum number of parameters per GPU before releasing them ZERO_OPTIMIZATION_MAX_LIVE_PARAMETERS = 'stage3_max_live_parameters' @@ -124,6 +125,10 @@ ZERO_OPTIMIZATION_LEGACY_STAGE1 = "legacy_stage1" ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT = False +# Stage 2 - partition gradients in a round robin fashsion to load-balance reduction and offload copying +ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS = 'round_robin_gradients' +ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS_DEFAULT = False + #yapf: disable ZERO_OPTIMIZATION_DEFAULT = { ZERO_OPTIMIZATION_STAGE: @@ -161,5 +166,7 @@ ZERO_OPTIMIZATION_IGNORE_UNUSED_PARAMETERS: ZERO_OPTIMIZATION_IGNORE_UNUSED_PARAMETERS_DEFAULT, ZERO_OPTIMIZATION_LEGACY_STAGE1: - ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT + ZERO_OPTIMIZATION_LEGACY_STAGE1_DEFAULT, + ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS: + ZERO_OPTIMIZATION_ROUND_ROBIN_GRADIENTS_DEFAULT } diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index 04a82d177611..7dc9076eed3e 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -18,6 +18,7 @@ from ..utils import see_memory_usage from deepspeed.utils import log_dist, init_distributed +from deepspeed.utils.debug import debug_param2name_id_shape, debug_module2name, debug_param2name, debug_param2name_id_shape_status, printflock, log_rank_file from ..swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper, PartitionedParamStatus from ..config import DeepSpeedConfig @@ -27,8 +28,14 @@ def print_rank_0(message, debug=False, force=False): - if torch.distributed.get_rank() == 0 and (debug or force): + rank = torch.distributed.get_rank() + if rank == 0 and (debug or force): print(message) + # other variations + # - print for all ranks w/o interleaving + # printflock(f"[{rank}] {message}") + # - print to log file per rank + # log_rank_file(rank, message) def is_zero_param(parameter): @@ -208,11 +215,16 @@ def recurse(cl): # Inserts _post_init_method at the end of init method # for all sub classes of torch.nn.Module class InsertPostInitMethodToModuleSubClasses(object): - def __init__(self, enabled=True, mem_efficient_linear=True, config=None, dtype=None): + def __init__(self, + enabled=True, + mem_efficient_linear=True, + ds_config=None, + dtype=None): self.mem_efficient_linear = mem_efficient_linear self.enabled = enabled - self._set_dtype(config, dtype) - assert self.dtype in [torch.half, torch.float], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.float]" + self._set_dtype(ds_config, dtype) + assert self.dtype in [ + torch.half, torch.float], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.float]" def __enter__(self): if not self.enabled: @@ -259,7 +271,7 @@ def _init_subclass(cls, **kwargs): if self.mem_efficient_linear: print_rank_0( "nn.functional.linear has been overridden with a more memory efficient version. This will persist unless manually reset.", - force=True) + force=False) self.linear_bk = torch.nn.functional.linear torch.nn.functional.linear = LinearFunctionForZeroStage3.apply @@ -280,8 +292,8 @@ def _disable_class(cls): torch.Tensor.__new__ = torch.Tensor.__old_new__ torch.empty = _orig_torch_empty - #un doing it here will undo it during training - #if self.mem_efficient_linear: + # un doing it here will undo it during training + # if self.mem_efficient_linear: # torch.nn.functional.linear = self.linear_bk # if self.mem_efficient_linear: # torch.nn.functional.linear = self.linear_bk @@ -296,8 +308,7 @@ def _post_init_method(self, module): def _set_dtype(self, ds_config, dtype): if ds_config is not None and dtype is None: - _ds_config = DeepSpeedConfig(ds_config) - self.dtype = torch.half if _ds_config.fp16_enabled else torch.float + self.dtype = torch.half if ds_config.fp16_enabled else torch.float elif dtype is None: self.dtype = torch.half else: @@ -314,9 +325,11 @@ def __init__(self, mem_efficient_linear=True, remote_device=None, pin_memory=False, + config_dict_or_path=None, config=None, enabled=True, - dtype=None): + dtype=None, + mpu=None): """A context to enable massive model construction for training with ZeRO-3. Models are automatically partitioned (or, sharded) across the system and converted to half precision. @@ -336,12 +349,14 @@ def __init__(self, pin_memory (bool, optional): Potentially increase performance by using pinned memory for model weights. ``remote_device`` must be ``"cpu"``. Defaults to ``False``. - config (``json file`` or dict, optional): If provided, provides configuration + config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration for swapping fp16 params to NVMe. + config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. enabled (bool, optional): If ``False``, this context has no effect. Defaults to ``True``. dtype (``dtype``, optional): Can be used to change the data type of the parameters. Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` + mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,wolrd_size} This context accelerates model initialization and enables models that are too large to allocate in their entirety in CPU memory. It has the @@ -413,9 +428,11 @@ def get_model(): model = deepspeed.zero.Init(module=model) """ + _ds_config = DeepSpeedConfig(config_dict_or_path, + mpu) if config_dict_or_path is not None else None super().__init__(enabled=enabled, mem_efficient_linear=mem_efficient_linear, - config=config, + ds_config=_ds_config, dtype=dtype) if not torch.distributed.is_initialized(): init_distributed() @@ -428,21 +445,20 @@ def get_model(): self.rank = torch.distributed.get_rank(group=self.ds_process_group) self.world_size = torch.distributed.get_world_size(group=self.ds_process_group) - #Local device is the device where the parameters are consumed - #It is the device where parameters are fully instantiated using allgather + # Local device is the device where the parameters are consumed + # It is the device where parameters are fully instantiated using allgather self.local_device = torch.device('cuda:{}'.format(os.environ["LOCAL_RANK"])) - self._validate_remote_device(remote_device, config) + self._validate_remote_device(remote_device, _ds_config) - #Remote device is the device where parameter partiitons are stored - #It can be same as local_device or it could be CPU or NVMe. + # Remote device is the device where parameter partiitons are stored + # It can be same as local_device or it could be CPU or NVMe. self.remote_device = self.local_device if remote_device is None else remote_device self.pin_memory = pin_memory if ( self.remote_device == OFFLOAD_CPU_DEVICE) else False # Enable fp16 param swapping to NVMe if self.remote_device == OFFLOAD_NVME_DEVICE: - _ds_config = DeepSpeedConfig(config) self.param_swapper = AsyncPartitionedParameterSwapper(_ds_config) else: self.param_swapper = None @@ -456,22 +472,21 @@ def get_model(): self._convert_to_deepspeed_param(param) param.partition() - def _validate_remote_device(self, remote_device, ds_config): - if ds_config is not None: - _ds_config = DeepSpeedConfig(ds_config) + def _validate_remote_device(self, remote_device, _ds_config): + if _ds_config is not None: if remote_device in [None, OFFLOAD_CPU_DEVICE]: if _ds_config.zero_config.offload_param is not None: offload_param_device = _ds_config.zero_config.offload_param[ OFFLOAD_PARAM_DEVICE] assert offload_param_device != OFFLOAD_NVME_DEVICE, \ - f"{OFFLOAD_PARAM_DEVICE} in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}." + f"{OFFLOAD_PARAM_DEVICE} in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}." if remote_device == OFFLOAD_NVME_DEVICE: assert _ds_config.zero_config.offload_param is not None, \ - f'{OFFLOAD_PARAM} must be defined in DeepSpeed Config if remote device is {OFFLOAD_NVME_DEVICE}.' + f'{OFFLOAD_PARAM} must be defined in DeepSpeed Config if remote device is {OFFLOAD_NVME_DEVICE}.' assert _ds_config.zero_config.offload_param[OFFLOAD_PARAM_NVME_PATH] is not None, \ - f'{OFFLOAD_PARAM_NVME_PATH} in DeepSpeed Config cannot be None if remote device is {OFFLOAD_NVME_DEVICE}' + f'{OFFLOAD_PARAM_NVME_PATH} in DeepSpeed Config cannot be None if remote device is {OFFLOAD_NVME_DEVICE}' def _post_init_method(self, module): #see_memory_usage(f"Before converting parmas in {module.__class__.__name__}", force=False) @@ -481,12 +496,12 @@ def _post_init_method(self, module): force=False) global param_count - for name, param in module.named_parameters(recurse=False): + for param in module.parameters(recurse=False): param_count += param.numel() if not is_zero_param(param): self._convert_to_deepspeed_param(param) print_rank_0( - f"Partitioning param with ds id {param.ds_id} and shape {param.data.shape}" + f"Partitioning param {debug_param2name_id_shape(param)} module={debug_module2name(module)}" ) param.partition() see_memory_usage( @@ -617,7 +632,7 @@ def _ensure_availability_of_partitioned_params(self, params): def _all_gather(self, param_list, async_op=False, hierarchy=None): - #fetches from nvme if the partition is not available and in nvme + # fetches from nvme if the partition is not available and in nvme self._ensure_availability_of_partitioned_params(param_list) handles = [] @@ -644,10 +659,10 @@ def _all_gather(self, param_list, async_op=False, hierarchy=None): def _partition(self, param_list, force=False, has_been_updated=False): for param in param_list: #print_rank_0(f"Before Partitioning Param {param.ds_id}") - #self._param_status(param) + # self._param_status(param) self._partition_param(param, has_been_updated=has_been_updated) param.ds_status = ZeroParamStatus.NOT_AVAILABLE - #if param.ds_tensor is not None: + # if param.ds_tensor is not None: # assert id(param.data) == id(param.ds_tensor.data), \ # "After the parameters are initially partitioned, make sure we are not recreating the partition." #print_rank_0(f"After Partitioning Param {param.ds_id}") @@ -671,7 +686,7 @@ def _partition_param(self, param, buffer=None, has_been_updated=False): # if numel in empty_buffers: # empty_buffers[numel].append(buffer) - #if torch.distributed.get_rank(): + # if torch.distributed.get_rank(): # print(f"Releasing {param.data.numel()}") if param.ds_tensor is not None and not has_been_updated: @@ -680,7 +695,7 @@ def _partition_param(self, param, buffer=None, has_been_updated=False): see_memory_usage( f'Before partitioning param {param.ds_id} {param.shape}', force=False) - #param.data does not store anything meaningful in partitioned state + # param.data does not store anything meaningful in partitioned state param.data = torch.ones(1, dtype=self.dtype).to(param.device) see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) @@ -758,7 +773,7 @@ def _partition_param(self, param, buffer=None, has_been_updated=False): #param.data = param.ds_tensor.data - #param.data does not store anything meaningful in partitioned state + # param.data does not store anything meaningful in partitioned state see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) @@ -797,23 +812,23 @@ def _allgather_param(self, param, async_op=False, hierarchy=0): assert tensor_size == aligned_param_size, f'param id {param.ds_id} aligned size {aligned_param_size} does not match tensor size {tensor_size}' print_rank_0( - f"{'--'* hierarchy}---- Before allocating Allgather param with id {param.ds_id} and status {param.ds_status} Partition Size {partition_size} and data shape {param.ds_shape}" + f"{'--'* hierarchy}---- Before allocating allgather param {debug_param2name_id_shape_status(param)} partition size={partition_size}" ) see_memory_usage( - f'Before allocate allgather param {param.ds_id} {param.ds_status} {aligned_param_size} {partition_size} {param.ds_shape}', + f'Before allocate allgather param {debug_param2name_id_shape_status(param)} partition_size={partition_size} ', force=False) flat_tensor = torch.zeros(aligned_param_size, dtype=param.dtype, device=param.device).view(-1) see_memory_usage( - f'After allocate allgather param {param.ds_id} {param.ds_status} {aligned_param_size} {partition_size} {param.ds_shape}', + f'After allocate allgather param {debug_param2name_id_shape_status(param)} {aligned_param_size} {partition_size} ', force=False) torch.cuda.synchronize() print_rank_0( - f"{'--'* hierarchy}----Allgather param with id {param.ds_id} and status {param.ds_status} Partition Size {partition_size} and data shape {param.ds_shape}" + f"{'--'* hierarchy}----allgather param with {debug_param2name_id_shape_status(param)} partition size={partition_size}" ) # if not flat_tensor.numel() > 100000: # replicated_tensor = flat_tensor.narrow(0, @@ -995,7 +1010,8 @@ def _partition_gradient(self, param, partition_buffer=None, accumulate=False): dtype=param.dtype, device=param.device) else: - assert partition_buffer.numel() >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}" + assert partition_buffer.numel( + ) >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}" rank = torch.distributed.get_rank(group=self.ds_process_group) start = partition_size * rank diff --git a/deepspeed/runtime/zero/stage1.py b/deepspeed/runtime/zero/stage1.py index 7660c9917b84..20a6c5a21944 100755 --- a/deepspeed/runtime/zero/stage1.py +++ b/deepspeed/runtime/zero/stage1.py @@ -5,7 +5,7 @@ from deepspeed.runtime.zero.utils import _initialize_parameter_parallel_groups from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler -from deepspeed.runtime.utils import get_grad_norm, CheckOverflow +from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION_OPTIMIZER_STATES from deepspeed.utils import logger, log_dist from deepspeed.ops.op_builder import UtilsBuilder @@ -104,6 +104,7 @@ def __init__(self, self.postscale_gradients = postscale_gradients self.gradient_predivide_factor = gradient_predivide_factor self.gradient_average = gradient_average + self._global_grad_norm = 0. # TODO: automatically turn off if #params > some_limit self.all_gather_partitions = all_gather_partitions @@ -683,8 +684,11 @@ def step(self, closure=None): local_sub_partitions_grad_groups.append(local_grad_sub_partitions) + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + #RS: update unscale/clip with sub partitions - self.unscale_and_clip_grads(local_sub_partitions_grad_groups, norm_groups) + self.unscale_and_clip_grads(local_sub_partitions_grad_groups, + self._global_grad_norm) self.optimizer.step() @@ -720,12 +724,7 @@ def step(self, closure=None): return self.overflow - def unscale_and_clip_grads(self, grad_groups_flat, norm_groups): - total_norm = 0.0 - for norm in norm_groups: - total_norm += norm**2.0 - total_norm = math.sqrt(total_norm) - + 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.: diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py index a9216160f2c5..09cdfe80b59b 100755 --- a/deepspeed/runtime/zero/stage2.py +++ b/deepspeed/runtime/zero/stage2.py @@ -13,7 +13,7 @@ import collections from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler -from deepspeed.runtime.utils import see_memory_usage, is_model_parallel_parameter +from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, get_global_norm, see_memory_usage, is_model_parallel_parameter from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION_GRADIENTS from deepspeed.ops.adam import DeepSpeedCPUAdam from deepspeed.ops.op_builder import UtilsBuilder @@ -99,12 +99,14 @@ def __init__(self, gradient_predivide_factor=1.0, gradient_accumulation_steps=1, ignore_unused_parameters=True, - partition_grads=True): + partition_grads=True, + round_robin_gradients=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: {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 @@ -144,12 +146,14 @@ def __init__(self, self.is_gradient_accumulation_boundary = True + self._global_grad_norm = 0. + if mpu is None: self.model_parallel_group = None self.model_parallel_rank = 0 else: self.model_parallel_group = mpu.get_model_parallel_group() - self.model_parallel_rank = mpu.get_model_parallel_rank() + self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu) self.overflow = False self.clip_grad = clip_grad @@ -159,6 +163,7 @@ def __init__(self, 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 @@ -232,10 +237,15 @@ def __init__(self, # 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). - round_robin_tensors, round_robin_indices = self._round_robin_reorder( - self.fp16_groups[i], - dist.get_world_size(group=self.dp_process_group) - ) + if self.round_robin_gradients: + round_robin_tensors, round_robin_indices = self._round_robin_reorder( + self.fp16_groups[i], + dist.get_world_size(group=self.dp_process_group) + ) + else: + round_robin_tensors = self.fp16_groups[i] + round_robin_indices = list(range(len(self.fp16_groups[i]))) + self.round_robin_fp16_groups.append(round_robin_tensors) self.round_robin_fp6_indices.append(round_robin_indices) @@ -462,7 +472,7 @@ def initialize_optimizer_states(self): if not self.cpu_offload: for group in self.single_partition_of_fp32_groups: - group.grad = None + group.grad = None #class init return @@ -486,7 +496,8 @@ def reduce_gradients(self, pipeline_parallel=False): if not self.overlap_comm: for i, group in enumerate(self.fp16_groups): for param in group: - self.reduce_ready_partitions_and_remove_grads(param, i) + if param.grad 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() @@ -964,12 +975,16 @@ def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param): src_tensor = param.grad.view(-1).narrow(0, source_offset, num_elements).float() dest_tensor.copy_(src_tensor, non_blocking=True) - param.grad = None + 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, 'ds_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, @@ -1082,17 +1097,18 @@ def reduce_ipg_grads(self): self.params_already_reduced[param_id] = True - 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: - param.grad = None - elif self.contiguous_gradients: - self.copy_grads_in_partition(param) + 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: + param.grad = None #only if self.partition_gradients + elif self.contiguous_gradients: + self.copy_grads_in_partition(param) self.grads_in_ipg_bucket = [] self.params_in_ipg_bucket = [] @@ -1112,7 +1128,7 @@ def are_all_related_partitions_reduced(params_id): 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 + 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) @@ -1201,7 +1217,7 @@ def allreduce_bucket(self, bucket, allreduce_always_fp32=False, rank=None, log=N def _clear_previous_reduced_grads(self): if self.previous_reduced_grads is not None: for param in self.previous_reduced_grads: - param.grad = None + param.grad = None # overlap enabled self.previous_reduced_grads = None #if rank is specified do a reduction instead of an allreduce @@ -1316,7 +1332,7 @@ def zero_grad(self, set_grads_to_None=True): for group in self.fp16_groups: for p in group: if set_grads_to_None: - p.grad = None + p.grad = None # epilogue and in step else: if p.grad is not None: p.grad.detach_() @@ -1366,6 +1382,9 @@ def get_grad_norm_direct(self, gradients, params, norm_type=2): #if dist.get_rank() == 0: # logger.info(f"Total Norm begining {total_norm}") for g, p in zip(gradients, params): + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, 'ds_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 @@ -1442,7 +1461,7 @@ def get_flat_partition(self, def free_grad_in_param_list(self, param_list): for p in param_list: - p.grad = None + p.grad = None # in step def reset_cpu_buffers(self): self.norm_for_param_grads = {} @@ -1499,11 +1518,6 @@ def step(self, closure=None): see_memory_usage('After overflow after clearing gradients') - logger.info( - "[deepspeed] fp16 dynamic loss scale overflow! Rank {} Skipping step. Attempted loss scale: {}, " - "reducing to {}".format(dist.get_rank(), - prev_scale, - self.loss_scale)) self.start_timers(timer_names) self.stop_timers(timer_names) return @@ -1548,7 +1562,8 @@ def step(self, closure=None): single_partition_grad_groups.append(single_grad_partition) - self.unscale_and_clip_grads(single_partition_grad_groups, norm_groups) + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + self.unscale_and_clip_grads(single_partition_grad_groups, self._global_grad_norm) self.stop_timers([OPTIMIZER_GRADIENTS]) self.start_timers([OPTIMIZER_STEP]) @@ -1570,7 +1585,7 @@ def step(self, closure=None): #get rid of the fp32 gradients. Not needed anymore if not self.cpu_offload: for group in self.single_partition_of_fp32_groups: - group.grad = None + group.grad = None # in step for fp16_partitions, fp32_partition in zip(self.parallel_partitioned_fp16_groups, self.single_partition_of_fp32_groups): fp16_partitions[partition_id].data.copy_(fp32_partition.data) @@ -1624,12 +1639,7 @@ def step(self, closure=None): return - def unscale_and_clip_grads(self, grad_groups_flat, norm_groups): - total_norm = 0.0 - for norm in norm_groups: - total_norm += norm**2.0 - total_norm = math.sqrt(total_norm) - + 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.: @@ -2000,3 +2010,106 @@ def _handle_overflow(cpu_sum, x, i): 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 = [] + enabled.append(f"cpu_offload={1 if cpu_offload else 0}") + 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/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index cac206ed759b..9b06d177fefd 100755 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -16,7 +16,7 @@ from deepspeed.utils.logging import logger from deepspeed.runtime.fp16.loss_scaler import LossScaler, DynamicLossScaler -from deepspeed.runtime.utils import see_memory_usage, is_model_parallel_parameter +from deepspeed.runtime.utils import get_global_norm, see_memory_usage, is_model_parallel_parameter from deepspeed.runtime.zero.partition_parameters import * from deepspeed.runtime.zero.partition_parameters import _init_external_params from deepspeed.runtime.zero.constants import ZERO_OPTIMIZATION_WEIGHTS @@ -32,11 +32,18 @@ pg_correctness_test = False FWD_MODULE_STACK = list() +from deepspeed.utils.debug import debug_module2name_id, debug_param2name_id_numel, debug_param2name_id_shape_device, debug_module2name_class, printflock, log_rank_file def print_rank_0(message, debug=False, force=False): - if torch.distributed.get_rank() == 0 and (debug or force): - logger.info(message) + rank = torch.distributed.get_rank() + if rank == 0 and (debug or force): + print(message) + # other variations + # - print for all ranks w/o interleaving + # printflock(f"[{rank}] {message}") + # - print to log file per rank + # log_rank_file(rank, message) def input(msg): @@ -211,7 +218,7 @@ def get_params_to_prefetch(self, sub_module, numel=2000000): # tracing failed. The sub_module passed at the step_id must match with the sub_module during tracing if sub_module.id != self.sub_module_trace[self.step_id]: print_rank_0( - f"Tracing failed. Prefetching is disabled at sub-module: {sub_module.id}" + f"Tracing failed. Prefetching is disabled at sub-module: {debug_module2name_id(sub_module)}" ) return [] @@ -390,11 +397,13 @@ def reset_step(self): def fetch_sub_module(self, sub_module): partitioned_params = [] params_in_flight = False - #print_rank_0(f"{'--' * self.hierarchy}Fetching params in module {sub_module.__class__.__name__}") + # print_rank_0(f"{'--' * self.hierarchy}Fetching params in module {sub_module.__class__.__name__}") params_to_fetch = [ param for _, param in sub_module.named_parameters(recurse=False) ] + # print([n for n,p in sub_module.named_parameters(recurse=False)]) + if hasattr(sub_module, 'ds_external_parameters'): print_rank_0( f"{'--' * self.hierarchy}--Fetching external parameters {sub_module.ds_external_parameters()}" @@ -407,7 +416,7 @@ def fetch_sub_module(self, sub_module): for param in params_to_fetch: param.ds_active_sub_modules += 1 print_rank_0( - f"{'--' * self.hierarchy}--Fetching parameters {param.ds_id} with active sub modules {param.ds_active_sub_modules}" + f"{'--' * self.hierarchy}--Fetching parameters {param.ds_id} {param.ds_shape} with active sub modules {param.ds_active_sub_modules}" ) if param.ds_status == ZeroParamStatus.AVAILABLE: @@ -441,14 +450,14 @@ def fetch_sub_module(self, sub_module): for _, param in sub_module.named_parameters(recurse=False): param.ds_status = ZeroParamStatus.AVAILABLE print_rank_0( - f"Param id {param.ds_id}, Shape {param.shape}, device {param.device} norm {param.norm()}", + f"Param {debug_param2name_id_shape_device(param)} norm={param.norm()}", force=False) #print_rank_0(f"After fetching (id, shape, device): {[(param.ds_id, param.shape, param.device) for param in sub_module.named_parameters(recurse=False)]}") def release_sub_module(self, sub_module): self.hierarchy -= 1 print_rank_0( - f"{'--' * self.hierarchy}Releasing params in module {sub_module.__class__.__name__}" + f"{'--' * self.hierarchy}Releasing params in module {debug_module2name_class(sub_module)}" ) params_to_release = [ param for _, @@ -468,31 +477,31 @@ def release_sub_module(self, sub_module): if not param.ds_active_sub_modules and not self._keep_for_later( sub_module) and not param.ds_persist: print_rank_0( - f"{'--' * self.hierarchy}--Releasing parameters {param.ds_id} with numel {param.numel()} active sub modules {param.ds_active_sub_modules} and keep for later {self._keep_for_later(sub_module)}", + f"{'--' * self.hierarchy}--Releasing parameter {debug_param2name_id_numel(param)} active sub modules {param.ds_active_sub_modules} and keep for later {self._keep_for_later(sub_module)}", force=False) # Keeping track of number of elements that are consumed by available parameters self._decrement_available_parameter_numel(param.ds_numel) see_memory_usage( - f"Before releasing param {param.ds_id} with numel {param.numel()}", + f"Before releasing param {debug_param2name_id_numel(param)}", force=False) param.partition(hierarchy=self.hierarchy) see_memory_usage( - f"After releasing param {param.ds_id} has numel {param.numel()} ", + f"After releasing param {debug_param2name_id_numel(param)}", force=False) param.ds_status = ZeroParamStatus.NOT_AVAILABLE else: print_rank_0( - f"{'--' * self.hierarchy}--Did not release parameters {param.ds_id} with numel {param.numel()} with active sub modules {param.ds_active_sub_modules}, keep for later {self._keep_for_later(sub_module)} and persistence {param.ds_persist}", + f"{'--' * self.hierarchy}--Did not release param {debug_param2name_id_numel(param)} with active sub modules {param.ds_active_sub_modules}, keep for later={self._keep_for_later(sub_module)} and persistence={param.ds_persist}", force=False) def release_and_reset_parameter(self, param): param.ds_active_sub_modules = 0 if param.ds_status == ZeroParamStatus.AVAILABLE: print_rank_0( - f"Releasing unpartitioned {param.ds_id} active sub-modules {param.ds_active_sub_modules} size {param.ds_numel} and persisitence {param.ds_persist}" + f"Releasing unpartitioned param {debug_param2name_id_numel(param)} active sub-modules {param.ds_active_sub_modules} and persisitence {param.ds_persist}" ) self._decrement_available_parameter_numel(param.ds_numel) param.partition() @@ -618,7 +627,7 @@ def __init__(self, elastic_checkpoint=False, aio_config=None): - see_memory_usage("Stage 3 initialize beginning", force=True) + see_memory_usage("Stage 3 initialize beginning", force=False) if dist.get_rank() == 0: logger.info(f"Reduce bucket size {reduce_bucket_size}") @@ -641,6 +650,7 @@ def __init__(self, self.flatten = util_ops.flatten self.unflatten = util_ops.unflatten self.dtype = self.optimizer.param_groups[0]['params'][0].dtype + self._global_grad_norm = 0. if not all(is_zero_param(p) for p in module.parameters()): group = None @@ -693,7 +703,7 @@ def __init__(self, self.max_params_in_cpu = offload_param_config[OFFLOAD_PARAM_MAX_IN_CPU] print_rank_0( f"FP16 params swapping is {self.params_in_nvme_and_cpu}, Max params in CPU is {self.max_params_in_cpu}", - force=True) + force=False) self.deepspeed_adam_offload = (self.offload_optimizer and type(init_optimizer) == DeepSpeedCPUAdam) @@ -789,7 +799,7 @@ def __init__(self, self.sub_group_size = sub_group_size self.sub_group_to_group_id = {} - see_memory_usage("Before creating fp16 partitions", force=True) + see_memory_usage("Before creating fp16 partitions", force=False) self._create_fp16_partitions_with_defragmentation() num_fp16_subgroups = len(self.fp16_partitioned_groups_flat) see_memory_usage(f"After creating fp16 partitions: {num_fp16_subgroups}", @@ -852,10 +862,13 @@ def __init__(self, count = count + 1 #Largest partitioned param - largest_partitioned_param_numel = max(self.fp16_partitioned_groups_flat_numel) + largest_partitioned_param_numel = max([ + max([tensor.numel() for tensor in fp16_partitioned_group]) + for fp16_partitioned_group in self.fp16_partitioned_groups + ]) print_rank_0( f'Largest partitioned param numel = {largest_partitioned_param_numel}', - force=True) + force=False) see_memory_usage(f"Before Set Grad positions", force=False) @@ -910,7 +923,7 @@ def __init__(self, self.debug_fp16_grads = [{} for _ in self.fp16_groups] if dist.get_rank(group=self.dp_process_group) == 0: - see_memory_usage(f"After initializing ZeRO optimizer", force=True) + see_memory_usage(f"After initializing ZeRO optimizer", force=False) def _configure_tensor_swapping(self, offload_optimizer_config, aio_config): nvme_swap_folder = os.path.join( @@ -1084,7 +1097,7 @@ def _create_fp16_partitions_with_defragmentation(self): for j, param_group in enumerate(self.optimizer.param_groups): sub_groups = self._create_fp16_sub_groups(param_group['params']) - print_rank_0(f'fp16 group {j} has {len(sub_groups)} subgroups', force=True) + print_rank_0(f'fp16 group {j} has {len(sub_groups)} subgroups', force=False) flat_offset = 0 for sub_group in sub_groups: @@ -1319,19 +1332,19 @@ def _create_fp32_partitions(self): nvme_gigabytes = nvme_memory_usage / GIGA_BYTES print_rank_0( f'Swappable FP32 Partitions: count={num_swappable_partitions} size={nvme_gigabytes:5.2f} GB', - force=True) + force=False) if self.params_in_nvme_and_cpu: print_rank_0( f'Swap from NVMe Partitions: count = {num_swap_from_nvme_partitions}, size = {swap_from_nvme_memory_usage/GIGA_BYTES:5.2f}GB', - force=True) + force=False) print_rank_0( f'Swap from CPU Partitions: count = {num_swap_from_cpu_partitions}, size = {swap_from_cpu_memory_usage/GIGA_BYTES:5.2f}GB', - force=True) + force=False) cpu_memory_gigabytes = cpu_memory_usage / GIGA_BYTES print_rank_0( f'In-Memory FP32 Partitions: count={cpu_memory_sub_groups} size={cpu_memory_gigabytes:5.2f} GB', - force=True) + force=False) # Clear for on-the-fly population before the optimizer step for param_group in self.optimizer.param_groups: @@ -2752,6 +2765,7 @@ def step(self, closure=None): return norm_groups = self._get_norm_groups() + self._global_grad_norm = get_global_norm(norm_list=norm_groups) timer_names = set() @@ -2765,7 +2779,7 @@ def step(self, closure=None): self._prepare_sub_group(sub_group_id, timer_names) #scale the fp32 gradients - self.unscale_and_clip_grads(sub_group_id, norm_groups) + self.unscale_and_clip_grads(sub_group_id, self._global_grad_norm) #apply the optimizer step on the sub group and copy fp32 parameters to fp16 self._optimizer_step(sub_group_id) @@ -2814,15 +2828,9 @@ def dump_post_step_gradients(self): norm_list = [param_norm, ds_norm] + unflat_norm print(f'Post-Step Norms {i} {param_id} = {norm_list}') - def unscale_and_clip_grads(self, sub_group_id, norm_groups): - + def unscale_and_clip_grads(self, sub_group_id, total_norm): grad_groups_flat = [self.fp32_partitioned_groups_flat[sub_group_id].grad] - total_norm = 0.0 - for norm in norm_groups: - total_norm += norm**2.0 - total_norm = math.sqrt(total_norm) - # compute combined scale factor for this group combined_scale = self.loss_scale if self.clip_grad > 0.: @@ -3255,3 +3263,155 @@ def _handle_overflow(cpu_sum, x, i): logger.info( f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}" ) + + +def estimate_zero3_model_states_mem_needs(total_params, + largest_layer_params, + num_gpus_per_node=1, + num_nodes=1, + cpu_offload=True, + cpu_offload_params=True, + zero_init=True, + additional_buffer_factor=1.5): + + total_gpus = num_nodes * num_gpus_per_node + gpus_factor = 1 / num_nodes + largest_layer_memory = (4 * largest_layer_params) + + if cpu_offload: + if cpu_offload_params: + gpu_mem = largest_layer_memory + + if zero_init: + cpu_mem = total_params * 18 * gpus_factor * additional_buffer_factor + else: + cpu_mem = total_params * max(4 * num_gpus_per_node, + 18 * gpus_factor) * additional_buffer_factor + + else: + gpu_mem = largest_layer_memory + int(2 * total_params / total_gpus) + + if zero_init: + cpu_mem = total_params * 16 * gpus_factor * additional_buffer_factor + else: + cpu_mem = total_params * max(4 * num_gpus_per_node, + 16 * gpus_factor) * additional_buffer_factor + else: + gpu_mem = largest_layer_memory + int(18 * total_params / total_gpus) + if zero_init: + cpu_mem = largest_layer_params * 4 * num_gpus_per_node * additional_buffer_factor + else: + cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor + + return int(cpu_mem), int(gpu_mem), largest_layer_memory + + +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()) + + largest_layer_params = 0 + for m in model.modules(): + # assuming no shared params within a single layer + layer_params = sum(p.numel() for p in m.parameters(recurse=False)) + largest_layer_params = max(largest_layer_params, layer_params) + + return total_params, largest_layer_params + + +import math + + +def estimate_zero3_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 3 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_zero3_model_states_mem_needs_all_cold`` where you have to pass + the ``total_params`` and ``largest_layer_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, largest_layer_params = model_to_params(model) + + estimate_zero3_model_states_mem_needs_all_cold( + total_params=total_params, + largest_layer_params=largest_layer_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + additional_buffer_factor=additional_buffer_factor) + + +def estimate_zero3_model_states_mem_needs_all_cold(total_params, + largest_layer_params, + num_gpus_per_node=1, + num_nodes=1, + additional_buffer_factor=1.5): + """ + Print out estimates on memory usage requirements for ZeRO 3 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_zero3_model_states_mem_needs_all_live`` and everything + will be derived automatically. + + Args: + - ``total_params``: total model params + - ``largest_layer_params``: largest layer's 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, cpu_offload_params, zero_init): + enabled = [] + enabled.append(f"cpu_offload={1 if cpu_offload else 0}") + enabled.append(f"cpu_offload_params={1 if cpu_offload_params else 0}") + enabled.append(f"zero_init={1 if zero_init else 0}") + 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, {int(largest_layer_params/1e6)}M largest layer params." + ) + print(" per CPU | per GPU | Options") + for cpu_offload in [True, False]: + for cpu_offload_params in [True, False]: + if not cpu_offload and cpu_offload_params: + continue + for zero_init in [True, False]: + cpu_mem, gpu_mem, largest_layer_memory = estimate_zero3_model_states_mem_needs( + total_params=total_params, + largest_layer_params=largest_layer_params, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, + cpu_offload=cpu_offload, + cpu_offload_params=cpu_offload_params, + zero_init=zero_init, + additional_buffer_factor=additional_buffer_factor + ) + + options_str = format_options(cpu_offload=cpu_offload, + cpu_offload_params=cpu_offload_params, + zero_init=zero_init) + print( + f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/deepspeed/runtime/zero/utils.py b/deepspeed/runtime/zero/utils.py index 8873c8db55d5..f90497069901 100755 --- a/deepspeed/runtime/zero/utils.py +++ b/deepspeed/runtime/zero/utils.py @@ -40,7 +40,7 @@ def _initialize_parameter_parallel_groups(parameter_parallel_size=None): def is_zero_supported_optimizer(optimizer): if dist.get_rank() == 0: - print( + logger.info( f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}' ) return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS diff --git a/deepspeed/utils/debug.py b/deepspeed/utils/debug.py new file mode 100644 index 000000000000..bb2b5f7ef75b --- /dev/null +++ b/deepspeed/utils/debug.py @@ -0,0 +1,122 @@ +""" debug utils """ + +import fcntl + +# for debug purposes map module and param objects to their fully qualified names +module_names = {} +param_names = {} + + +def debug_extract_module_and_param_names(model): + # extract the fully qualified names as soon as the model is acquired + global module_names + global param_names + # XXX: can probably make a map of param2module and vice-versa + module_names = {module: name for name, module in model.named_modules()} + param_names = {param: name for name, param in model.named_parameters()} + + +def debug_module2name(module): + if module in module_names: + return module_names[module] + else: + return "unknown" + + +def debug_module2name_id(module): + return f"name={debug_module2name(module)} id={module.id}" + + +def debug_module2name_class(module): + return f"name={debug_module2name(module)} {module.__class__.__name__}" + + +def debug_param2name(param): + if param in param_names: + return param_names[param] + else: + return "unknown" + + +def debug_param2name_id(param): + return f"name={debug_param2name(param)} id={param.ds_id}" + + +def debug_param2name_id_shape(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.data.shape}" + + +def debug_param2name_id_shape_device(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.data.shape} device={param.device}" + + +def debug_param2name_id_numel(param): + return f"name={debug_param2name(param)} id={param.ds_id} numel={param.numel()}" + + +def debug_param2name_id_shape_status(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.data.shape} status={param.ds_status}" + + +def printflock(*msgs): + """ + + For printing messages for all concurrent gpus w/o getting interleaved text. + + This is useful when debugging issues where multi-gpus don't sync. + + 1. Enable the force debug in say partitioning and zero3 files + 2. Override the usual versions with :: + + def print_rank_0(message, debug=False, force=False): + rank = torch.distributed.get_rank() + printflock(f"[{rank}] {message}") + 3. run the program and you get both logs non-interleaved + + But this makes it very difficult to make sense of the output, so the ``log_rank_file`` helper + function might be more useful, as it's easier to send each log stream into a separate file and + then compare those. + + """ + + with open(__file__, "r") as fh: + fcntl.flock(fh, fcntl.LOCK_EX) + try: + print(*msgs) + finally: + fcntl.flock(fh, fcntl.LOCK_UN) + + +fh = None + + +def log_rank_file(rank, *msgs): + """ + Print to a log file of the given rank + + This is useful for debugging hanging in sync processes. Here is a possible workflow: + + 1. Enable the force debug in say partitioning and zero3 files + 2. Override the usual versions of print_rank_0 in those files with :: + + def print_rank_0(message, debug=False, force=False): + rank = torch.distributed.get_rank() + log_rank_file(rank, message) + + 3. run the program + 4. fix up the expected differences, e.g. different cuda numbers :: + + perl -pi -e 's|cuda:1|cuda:0|' log_rank_* + + 5. now diff and see where names and ids diverge - you will find where the gpus don't do the same + work (e.g. when some layers get conditionally skipped on one gpu but not all) + + diff -u log_rank_0.txt log_rank_1.txt | less + + """ + global fh + if fh is None: + fh = open(f"log_rank_{rank}.txt", "w") + for m in msgs: + fh.write(f"{m}\n") + fh.flush() diff --git a/deepspeed/utils/logging.py b/deepspeed/utils/logging.py index d16682f902e8..ca694514678e 100644 --- a/deepspeed/utils/logging.py +++ b/deepspeed/utils/logging.py @@ -3,6 +3,14 @@ import torch.distributed as dist +log_levels = { + "debug": logging.DEBUG, + "info": logging.INFO, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, +} + class LoggerFactory: @staticmethod @@ -58,3 +66,32 @@ def log_dist(message, ranks=None, level=logging.INFO): if should_log: final_message = "[Rank {}] {}".format(my_rank, message) logger.log(level, final_message) + + +def get_current_level(): + """ + Return logger's current log level + """ + return logger.getEffectiveLevel() + + +def should_log_le(max_log_level_str): + """ + Args: + max_log_level_str: maximum log level as a string + + Returns ``True`` if the current log_level is less or equal to the specified log level. Otherwise ``False``. + + Example: + + ``should_log_le("info")`` will return ``True`` if the current log level is either ``logging.INFO`` or ``logging.DEBUG`` + """ + + if not isinstance(max_log_level_str, str): + raise ValueError(f"{max_log_level_str} is not a string") + + max_log_level_str = max_log_level_str.lower() + if max_log_level_str not in log_levels: + raise ValueError(f"{max_log_level_str} is not one of the `logging` levels") + + return get_current_level() <= log_levels[max_log_level_str] diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index 2d98a39e3fc7..1b58c403618c 100644 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -20,13 +20,27 @@ debug = 0 +def get_model_state_file(checkpoint_dir): + + if not os.path.isdir(checkpoint_dir): + raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") + + # there should be only one file + file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") + + if not os.path.exists(file): + raise FileNotFoundError(f"can't find '{file}' in directory '{checkpoint_dir}'") + + return file + + def get_optim_files(checkpoint_dir): if not os.path.isdir(checkpoint_dir): raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") # XXX: need to test that this simple glob rule works for multi-node setup too - optim_files = sorted(glob.glob(f"{checkpoint_dir}/*_optim_states.pt")) + optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt"))) if len(optim_files) == 0: raise FileNotFoundError( @@ -35,13 +49,34 @@ def get_optim_files(checkpoint_dir): return optim_files +def parse_model_state(file): + + # load to cpu + device = torch.device('cpu') + state_dict = torch.load(file, map_location=device) + + if "buffer_names" not in state_dict: + raise ValueError(f"{file} is not a model state checkpoint") + buffer_names = state_dict["buffer_names"] + if debug: + print(buffer_names) + + # recover just the buffers while restoring them to fp32 if they were saved in fp16 + buffers = { + k: v.float() + for k, + v in state_dict["module"].items() if k in buffer_names + } + return buffers + + def parse_optim_states(files): state_dicts = [] for f in files: state_dicts.append(torch.load(f)) if not "zero_stage" in state_dicts[0]['optimizer_state_dict']: - raise ValueError(f"non zero checkpoint") + raise ValueError(f"{files[0]} is not a zero checkpoint") zero_stage = state_dicts[0]['optimizer_state_dict']["zero_stage"] world_size = state_dicts[0]['optimizer_state_dict']["partition_count"] param_shapes = state_dicts[0]["param_shapes"] @@ -86,7 +121,9 @@ def convert_zero_chkpt_to_fp32_consolid_state_dict(checkpoint_dir, output_file): """ print(f"Processing zero checkpoint '{checkpoint_dir}'") + model_file = get_model_state_file(checkpoint_dir) optim_files = get_optim_files(checkpoint_dir) + buffers = parse_model_state(model_file) zero_stage, world_size, param_shapes, fp32_flat_groups = parse_optim_states(optim_files) print( f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}") @@ -108,9 +145,16 @@ def convert_zero_chkpt_to_fp32_consolid_state_dict(checkpoint_dir, output_file): # XXX: memory usage doubles here (zero2) full_single_fp32_vector = torch.cat(fp32_flat_groups, 0) + state_dict = OrderedDict() + + # buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + # params # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support # out-of-core computing solution - state_dict = OrderedDict() offset = 0 total_numel = 0 for name, shape in param_shapes.items(): diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index f49d55a9a2a4..0a847b2fa5e5 100755 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -557,7 +557,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s | Description | Default | | ------------------------------ | ------- | -| Print train loss every N steps | `10` | +| Print progress report every N training steps. The report includes the number of training steps, number of skipped optimizer updates (likely due to overflows in mixed-precision training), current learning rate, and current momentum. | `10` | **wall_clock_breakdown**: [boolean] diff --git a/docs/_tutorials/advanced-install.md b/docs/_tutorials/advanced-install.md index b1c54325c3a7..a1493a9e9e24 100755 --- a/docs/_tutorials/advanced-install.md +++ b/docs/_tutorials/advanced-install.md @@ -11,10 +11,6 @@ just-in-time (JIT) using [torch's JIT C++ extension loader that relies on ninja](https://pytorch.org/docs/stable/cpp_extension.html) to build and dynamically link them at runtime. -**Note:** [PyTorch](https://pytorch.org/) must be installed _before_ installing -DeepSpeed. -{: .notice--info} - ```bash pip install deepspeed ``` @@ -30,6 +26,9 @@ ds_report ## Pre-install DeepSpeed Ops +**Note:** [PyTorch](https://pytorch.org/) must be installed _before_ pre-compiling any DeepSpeed c++/cuda ops. However, this is not required if using the default mode of JIT compilition of ops. +{: .notice--info} + Sometimes we have found it useful to pre-install either some or all DeepSpeed C++/CUDA ops instead of using the JIT compiled path. In order to support pre-installation we introduce build environment flags to turn on/off building diff --git a/docs/_tutorials/inference-tutorial.md b/docs/_tutorials/inference-tutorial.md index ab13b3dfe264..1686e75be9a2 100644 --- a/docs/_tutorials/inference-tutorial.md +++ b/docs/_tutorials/inference-tutorial.md @@ -71,6 +71,7 @@ DeepSpeed inference can be used in conjunction with HuggingFace `pipeline`. Belo ```python # Filename: gpt-neo-2.7b-generation.py +import os import deepspeed import torch import transformers @@ -82,10 +83,10 @@ generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B', device= -deepspeed.init_inference(generator.model, - mp_size=world_size, - dtype=torch.float, - replace_method='auto') +generator.model = deepspeed.init_inference(generator.model, + mp_size=world_size, + dtype=torch.float, + replace_method='auto') string = generator("DeepSpeed is", do_sample=True, min_length=50) if torch.distributed.get_rank() == 0: diff --git a/docs/code-docs/source/index.rst b/docs/code-docs/source/index.rst index 2f39da002d5e..5175209cc1c4 100644 --- a/docs/code-docs/source/index.rst +++ b/docs/code-docs/source/index.rst @@ -79,6 +79,14 @@ Flops Profiler flops-profiler + +Memory Usage +------------------ +.. toctree:: + :maxdepth: 2 + + memory + Indices and tables ------------------ diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst new file mode 100644 index 000000000000..7f810309a80b --- /dev/null +++ b/docs/code-docs/source/memory.rst @@ -0,0 +1,288 @@ +Memory Requirements +----------------------- + + +API To Estimate Memory Usage +============================ + +ZeRO2: + +.. autofunction:: deepspeed.runtime.zero.stage2.estimate_zero2_model_states_mem_needs_all_live + +.. autofunction:: deepspeed.runtime.zero.stage2.estimate_zero2_model_states_mem_needs_all_cold + +Examples: + +Let's try a 3B model with just 1 node with 8 gpus, using live model: + +.. code-block:: bash + + python -c 'from transformers import AutoModel; \ + from deepspeed.runtime.zero.stage2 import estimate_zero2_model_states_mem_needs_all_live; \ + model = AutoModel.from_pretrained("t5-3b"); \ + estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)' + Estimated memory needed for params, optim states and gradients for a: + HW: Setup with 1 node, 8 GPUs per node. + SW: Model with 2851M total params. + per CPU | per GPU | Options + 127.48GB | 5.31GB | cpu_offload=1 + 127.48GB | 15.93GB | cpu_offload=0 + +Now, without the actual model, which requires us to know ``total_params`` and +``largest_layer_params``, but we got those from the run above, so future estimators are now much +faster as we don't need to load the model. + +.. code-block:: bash + + python -c 'from deepspeed.runtime.zero.stage2 import estimate_zero2_model_states_mem_needs_all_cold; \ + estimate_zero2_model_states_mem_needs_all_cold(total_params=2851e6, num_gpus_per_node=8, num_nodes=1)' + Estimated memory needed for params, optim states and gradients for a: + HW: Setup with 1 node, 8 GPUs per node. + SW: Model with 2851M total params. + per CPU | per GPU | Options + 127.45GB | 5.31GB | cpu_offload=1 + 127.45GB | 15.93GB | cpu_offload=0 + +There is a slight difference due to rounding - the actual live model has a few more params + + +ZeRO3: + +.. autofunction:: deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_live + +.. autofunction:: deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_cold + +Examples: + +Let's try a 3B model with just 1 node with 8 gpus, using live model: + +.. code-block:: bash + + python -c 'from transformers import AutoModel; \ + from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \ + model = AutoModel.from_pretrained("t5-3b"); \ + estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)' + + Estimated memory needed for params, optim states and gradients for a: + HW: Setup with 1 node, 8 GPUs per node. + SW: Model with 2851M total params, 32M largest layer params. + per CPU | per GPU | Options + 71.71GB | 0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=1 + 127.48GB | 0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=0 + 63.74GB | 0.79GB | cpu_offload=1, cpu_offload_params=0, zero_init=1 + 127.48GB | 0.79GB | cpu_offload=1, cpu_offload_params=0, zero_init=0 + 1.47GB | 6.10GB | cpu_offload=0, cpu_offload_params=0, zero_init=1 + 127.48GB | 6.10GB | cpu_offload=0, cpu_offload_params=0, zero_init=0 + +Now, without the actual model, which requires us to know ``total_params`` and +``largest_layer_params``, but we got those from the run above, so future estimators are now much +faster as we don't need to load the model. + +.. code-block:: bash + + python -c 'from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_cold; \ + estimate_zero3_model_states_mem_needs_all_cold(total_params=2851e6, largest_layer_params=32e6, num_gpus_per_node=8, num_nodes=1)' + + Estimated memory needed for params, optim states and gradients for a: + HW: Setup with 1 node, 8 GPUs per node. + SW: Model with 2851M total params, 32M largest layer params. + per CPU | per GPU | Options + 71.69GB | 0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=1 + 127.45GB | 0.12GB | cpu_offload=1, cpu_offload_params=1, zero_init=0 + 63.72GB | 0.78GB | cpu_offload=1, cpu_offload_params=0, zero_init=1 + 127.45GB | 0.78GB | cpu_offload=1, cpu_offload_params=0, zero_init=0 + 1.43GB | 6.09GB | cpu_offload=0, cpu_offload_params=0, zero_init=1 + 127.45GB | 6.09GB | cpu_offload=0, cpu_offload_params=0, zero_init=0 + +There is a slight difference due to rounding - the actual live model has a few more params + + + +Discussion +========== + +Let's look in detail how the memory estimator API calculates these numbers and also discuss some additional numbers that aren't covered by the API. + +In the following discussion: + +- ``params`` - total number of model params, which can be calculated as: + +.. code-block:: python + + print(sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())) + +Some models already include the number of params in the model name, e.g. t5-11b (11B params), gpt-neo-1.3B (1.3B params), etc. + +Also if the model weights are stored in ``fp32`` the other quick way to calculate the size of the model is to simply divide the size of the ``state_dict`` file by 4 (fp32 == 4 bytes). For example, you can see that `t5-11b's pytorch_model.bin `__ is 42.1GB in size, so if we divide it by 4, we can immediately tell it's an 11B model. + +The following calculations show how much memory is required by model params, gradients and optimizer states. In addition to those you will need enough memory to fit activation calculations and any temporary memory for intermediate calculations, which for long sequences could be very significant (e.g. could take the same amount of memory as params+grads+optim_states combined). + +The optimizer states assume that ``Adam`` is used, where 4 bytes per parameter are used by momentum and another 4 by variance (8 in total). + +Gradients at ``fp32`` take 4 bytes, and parameters take 2 bytes at ``fp16` and 4 bytes at ``fp32``. + +**GPU RAM** + +The big question is how big of a model you can fit on the hardware you have? Or rather what size of a GPU RAM do you need to fit the desired model. + + +* ZeRO-2: + + - ``"cpu_offload": true``: 2 * params + + Example: a 40GB GPU can fit ~11B param model (regardless of how many GPUs are used). Here the model is loaded in ``fp16`` so just the model weights take about 22GB and the remaining 18GB are used by other components. You can barely fit a very small batch size in this scenario. + + - ``"cpu_offload": false``: 4 params + 16 params/ (total number of gpus) + +* ZeRO-3: + +largest_layer_memory = 4*largest_layer_params - GPU memory needed to gather the largest layer on a single GPU. 2 bytes fp16 params are gathered and 2 bytes fp16 grads are computed (total 4x). The optimizer states and fp32 parameters are updated in partitioned form and copied to fp16 params in partitioned form. This happens during the optimizer step. After that the fp16 params are sufficient. + + - case 1: ``"cpu_offload": false, "cpu_offload_params": false`` - largest_layer_memory + 18 * params / total number of gpus across all nodes + - case 2: ``"cpu_offload": true, "cpu_offload_params": true``- largest_layer_memory. The main limit here is general RAM. + - case 3: ``"cpu_offload": true, "cpu_offload_params": false``- largest_layer_memory + 2 * params / total number of gpus across all nodes + + Example: + +.. code-block:: python + +from transformers import AutoModel +model = AutoModel.from_pretrained("t5-large") + +# shared params calculated only ones +total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) + +largest_layer_params = 0 +for m in model.modules(): + # assuming no shared params within a single layer + layer_params = sum(p.numel() for p in m.parameters(recurse=False)) + largest_layer_params = max(largest_layer_params, layer_params) + +largest_layer_memory = (4*largest_layer_params) + +total_gpus = 4 + +case1 = largest_layer_memory + int(18*total_params/total_gpus) +case2 = largest_layer_memory +case3 = largest_layer_memory + int(2*total_params/total_gpus) + +print(f"total params: {total_params/1e6:6.2f}M") +print(f"largest layer params: {largest_layer_params/1e6:6.2f}M") +print(f"largest layer memory: {largest_layer_memory>>20:6}MB") +print(f"case1 gpu memory: {(case1)>>20:6}MB") +print(f"case2 gpu memory: {(case2)>>20:6}MB") +print(f"case3 gpu memory: {(case3)>>20:6}MB") + +total params: 737.67M +largest layer params: 32.90M +largest layer memory: 125MB +case1 gpu memory: 3291MB +case2 gpu memory: 125MB +case3 gpu memory: 477MB + + +**General RAM**: + +One of the key features of ZeRO is its CPU offload which can dramatically extend the total memory pool accessible to the project by using general RAM. One can easily expand their general RAM by 10x times, at a significantly lower cost than what it'd take to have the same GPU RAM. And often, it's not even possible to buy GPUs with a lot of RAM (112GB GPU anybody?) since they simply don't yet exist. + +In the following calculations we will use: + +- ``additional_buffer_factor=1.5`` as an additional buffer factor to be conservative +- ``n_gpus`` the number of GPUs on a single node (machine) +- ``total_gpus`` the total number of GPUs across all nodes +- ``params`` - total number of model params (see above for how to get this number) + +* ZeRO-2: + + - ``"cpu_offload": false``: + + params * 4 * n_gpus * additional_buffer_factor - this is the memory needed only at the beginning to initialize the model on CPU memory + + - ``"cpu_offload": true``: + + params * max(4 * n_gpus, 16) * additional_buffer_factor + + Example: xxx + +* ZeRO-3: + + gpus_factor = n_gpus / total_gpus + + - case 1: ``"cpu_offload": false``: + + Without ``zero.Init``: + + params * 4 * n_gpus * additional_buffer_factor + + this is the memory needed only at the beginning to initialize the model on CPU memory. Once the model is transferred to GPUs this memory is freed. + + With ``zero.Init``: + + largest_layer_params * 4 * n_gpus * additional_buffer_factor + + assuming Pytorch is deallocating the memory once the tensors are moved to the GPU by ZeRO.Init + + - case 2: ``"cpu_offload": true, cpu_offload_params true``: + + Without ``zero.Init``: + + params * max(4 * n_gpus, 18 * gpus_factor) * additional_buffer_factor + + With ``zero.Init``: + + params * 18 * gpus_factor * additional_buffer_factor + + - case 3: ``"cpu_offload": true, cpu_offload_params false``: + + Without ``zero.Init``: + + params * max(4 * n_gpus, 16 * gpus_factor) * additional_buffer_factor + + With ``zero.Init``: + + params * 16 * gpus_factor * additional_buffer_factor + + +Here is a breakdown for the 16 and 18 multipliers (b = bytes): + +4 (in ``4*n_gpus``): + +- when pytorch creates a model it creates it in fp32 by default (4 bytes) + +16: + +- 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter + +18: + +- 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter +- +2b for fp16 params + + +**Pinned Memory** + +Pinned general RAM is included in normal general RAM allocations (i.e. this is not extra memory allocations but simply shows how much of the general RAM is pinned) + +* ZeRO-2: can't be controlled + +* ZeRO-3 + +To enable add: ``"cpu_offload_use_pin_memory" : true`` + +Now there are 2 sub-cases: + +1. ``"cpu_offload_params": true``: + + - 6 * params (2b for fp16 params + 4b for fp32 gradients) + - if ``gradient_accumulation_steps > 1`` an additional 2b for fp16 gradients are pinned + +2. ``"cpu_offload_params": false``: + + - 4b for fp32 gradients + + +**Activation Memory** + +XXX: For Transformers is probably around (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu + +This needs to be completed. diff --git a/op_builder/builder.py b/op_builder/builder.py index 50dfa4f9c550..70f0fd3d3e55 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -4,7 +4,6 @@ import os import sys import time -import torch import importlib from pathlib import Path import subprocess @@ -17,6 +16,13 @@ DEFAULT_TORCH_EXTENSION_PATH = "/tmp/torch_extensions" DEFAULT_COMPUTE_CAPABILITIES = "6.0;6.1;7.0" +try: + import torch +except ImportError: + print( + f"{WARNING} unable to import torch, please install it if you want to pre-compile any deepspeed ops." + ) + def installed_cuda_version(): import torch.utils.cpp_extension @@ -57,7 +63,8 @@ def get_default_compute_capatabilities(): 11: ["11.0", "11.1", "11.2", - "11.3"], + "11.3", + "11.4"], } @@ -364,6 +371,18 @@ def cxx_args(self): else: return ['-O3', '-std=c++14', '-g', '-Wno-reorder'] + def nvcc_args(self): + args = [ + '-O3', + '--use_fast_math', + '-std=c++17' if sys.platform == "win32" else '-std=c++14', + '-U__CUDA_NO_HALF_OPERATORS__', + '-U__CUDA_NO_HALF_CONVERSIONS__', + '-U__CUDA_NO_HALF2_OPERATORS__' + ] + + return args + self.compute_capability_args() + def libraries_args(self): if sys.platform == "win32": return ['cublas', 'curand'] diff --git a/op_builder/cpu_adam.py b/op_builder/cpu_adam.py index 75fa042613fc..129ddeea3a29 100644 --- a/op_builder/cpu_adam.py +++ b/op_builder/cpu_adam.py @@ -3,7 +3,6 @@ """ import os import sys -import torch import subprocess from .builder import CUDAOpBuilder @@ -26,6 +25,7 @@ def sources(self): return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/custom_cuda_kernel.cu'] def include_paths(self): + import torch CUDA_INCLUDE = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include") return ['csrc/includes', CUDA_INCLUDE] @@ -47,6 +47,7 @@ def simd_width(self): return '-D__SCALAR__' def cxx_args(self): + import torch CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64") SIMD_WIDTH = self.simd_width() @@ -62,15 +63,3 @@ def cxx_args(self): '-fopenmp', SIMD_WIDTH ] - - def nvcc_args(self): - args = [ - '-O3', - '--use_fast_math', - '-std=c++14', - '-U__CUDA_NO_HALF_OPERATORS__', - '-U__CUDA_NO_HALF_CONVERSIONS__', - '-U__CUDA_NO_HALF2_OPERATORS__' - ] - args += self.compute_capability_args() - return args diff --git a/op_builder/fused_adam.py b/op_builder/fused_adam.py index 0340ed02a8fb..c9a0d4436d01 100644 --- a/op_builder/fused_adam.py +++ b/op_builder/fused_adam.py @@ -1,7 +1,6 @@ """ Copyright 2020 The Microsoft DeepSpeed Team """ -import torch from .builder import CUDAOpBuilder diff --git a/op_builder/fused_lamb.py b/op_builder/fused_lamb.py index 4c73ecb404e3..169654809d06 100644 --- a/op_builder/fused_lamb.py +++ b/op_builder/fused_lamb.py @@ -1,7 +1,6 @@ """ Copyright 2020 The Microsoft DeepSpeed Team """ -import torch from .builder import CUDAOpBuilder diff --git a/op_builder/quantizer.py b/op_builder/quantizer.py index f77e43c62408..43bc5778ea20 100644 --- a/op_builder/quantizer.py +++ b/op_builder/quantizer.py @@ -1,4 +1,3 @@ -import torch from .builder import CUDAOpBuilder @@ -21,15 +20,3 @@ def sources(self): def include_paths(self): return ['csrc/includes'] - - def nvcc_args(self): - args = [ - '-O3', - '--use_fast_math', - '-std=c++14', - '-U__CUDA_NO_HALF_OPERATORS__', - '-U__CUDA_NO_HALF_CONVERSIONS__', - '-U__CUDA_NO_HALF2_OPERATORS__' - ] - - return args + self.compute_capability_args() diff --git a/op_builder/sparse_attn.py b/op_builder/sparse_attn.py index 58b8ca4f0e6b..6f30cc84da23 100644 --- a/op_builder/sparse_attn.py +++ b/op_builder/sparse_attn.py @@ -1,7 +1,6 @@ """ Copyright 2020 The Microsoft DeepSpeed Team """ -import torch import warnings from .builder import OpBuilder @@ -28,6 +27,12 @@ def is_compatible(self): #command_status = list(map(self.command_exists, required_commands)) #deps_compatible = all(command_status) + try: + import torch + except ImportError: + self.warning(f"unable to import torch, please install it first") + return False + # torch-cpu will not have a cuda version if torch.version.cuda is None: cuda_compatible = False diff --git a/op_builder/stochastic_transformer.py b/op_builder/stochastic_transformer.py index b7e2f3845117..aa47c13c49e4 100644 --- a/op_builder/stochastic_transformer.py +++ b/op_builder/stochastic_transformer.py @@ -1,7 +1,6 @@ """ Copyright 2020 The Microsoft DeepSpeed Team """ -import torch from .transformer import TransformerBuilder diff --git a/op_builder/transformer.py b/op_builder/transformer.py index a94c6da62e39..2d48e2421b82 100644 --- a/op_builder/transformer.py +++ b/op_builder/transformer.py @@ -1,7 +1,6 @@ """ Copyright 2020 The Microsoft DeepSpeed Team """ -import torch from .builder import CUDAOpBuilder @@ -30,15 +29,3 @@ def sources(self): def include_paths(self): return ['csrc/includes'] - - def nvcc_args(self): - args = [ - '-O3', - '--use_fast_math', - '-std=c++14', - '-U__CUDA_NO_HALF_OPERATORS__', - '-U__CUDA_NO_HALF_CONVERSIONS__', - '-U__CUDA_NO_HALF2_OPERATORS__' - ] - - return args + self.compute_capability_args() diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py index 76ee22e0b2c6..94db63711914 100755 --- a/op_builder/transformer_inference.py +++ b/op_builder/transformer_inference.py @@ -1,4 +1,3 @@ -import torch from .builder import CUDAOpBuilder @@ -24,15 +23,3 @@ def sources(self): def include_paths(self): return ['csrc/transformer/inference/includes'] - - def nvcc_args(self): - args = [ - '-O3', - '--use_fast_math', - '-std=c++14', - '-U__CUDA_NO_HALF_OPERATORS__', - '-U__CUDA_NO_HALF_CONVERSIONS__', - '-U__CUDA_NO_HALF2_OPERATORS__', - ] - - return args + self.compute_capability_args() diff --git a/requirements/requirements.txt b/requirements/requirements.txt index afe7d231d4bf..c0e94ec4dabe 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,5 +1,4 @@ torch>=1.2 -torchvision>=0.4.0 tqdm tensorboardX==1.8 ninja diff --git a/setup.py b/setup.py index 654b983eec81..2424fdc6c46a 100755 --- a/setup.py +++ b/setup.py @@ -22,12 +22,14 @@ from setuptools import setup, find_packages import time +torch_available = True try: import torch from torch.utils.cpp_extension import BuildExtension except ImportError: - raise ImportError('Unable to import torch, please visit https://pytorch.org/ ' - 'to see how to properly install torch on your system.') + torch_available = False + print('[WARNING] Unable to import torch, pre-compiling ops will be disabled. ' \ + 'Please visit https://pytorch.org/ to see how to properly install torch on your system.') from op_builder import ALL_OPS, get_default_compute_capatabilities @@ -45,7 +47,7 @@ def fetch_requirements(path): } # If MPI is available add 1bit-adam requirements -if torch.cuda.is_available(): +if torch_available and torch.cuda.is_available(): if shutil.which('ompi_info') or shutil.which('mpiname'): cupy = f"cupy-cuda{torch.version.cuda.replace('.','')[:3]}" extras_require['1bit_adam'].append(cupy) @@ -60,12 +62,17 @@ def fetch_requirements(path): cmdclass = {} # For any pre-installed ops force disable ninja -cmdclass['build_ext'] = BuildExtension.with_options(use_ninja=False) +if torch_available: + cmdclass['build_ext'] = BuildExtension.with_options(use_ninja=False) -TORCH_MAJOR = torch.__version__.split('.')[0] -TORCH_MINOR = torch.__version__.split('.')[1] +if torch_available: + TORCH_MAJOR = torch.__version__.split('.')[0] + TORCH_MINOR = torch.__version__.split('.')[1] +else: + TORCH_MAJOR = "0" + TORCH_MINOR = "0" -if not torch.cuda.is_available(): +if torch_available and not torch.cuda.is_available(): # Fix to allow docker builds, similar to https://github.com/NVIDIA/apex/issues/486 print( "[WARNING] Torch did not find cuda available, if cross-compiling or running with cpu only " @@ -81,6 +88,9 @@ def fetch_requirements(path): BUILD_OP_DEFAULT = int(os.environ.get('DS_BUILD_OPS', BUILD_OP_PLATFORM)) print(f"DS_BUILD_OPS={BUILD_OP_DEFAULT}") +if BUILD_OP_DEFAULT: + assert torch_available, "Unable to pre-compile ops without torch installed. Please install torch before attempting to pre-compile ops." + def command_exists(cmd): if sys.platform == "win32": @@ -109,6 +119,7 @@ def op_enabled(op_name): # If op install enabled, add builder to extensions if op_enabled(op_name) and op_compatible: + assert torch_available, f"Unable to pre-compile {op_name}, please first install torch" install_ops[op_name] = op_enabled(op_name) ext_modules.append(builder.builder()) @@ -170,7 +181,7 @@ def create_dir_symlink(src, dest): torch_version = ".".join([TORCH_MAJOR, TORCH_MINOR]) # Set cuda_version to 0.0 if cpu-only cuda_version = "0.0" -if torch.version.cuda is not None: +if torch_available and torch.version.cuda is not None: cuda_version = ".".join(torch.version.cuda.split('.')[:2]) torch_info = {"version": torch_version, "cuda_version": cuda_version} diff --git a/tests/unit/simple_model.py b/tests/unit/simple_model.py index 9c6062d79faa..15c40976b6a1 100755 --- a/tests/unit/simple_model.py +++ b/tests/unit/simple_model.py @@ -17,10 +17,7 @@ def __init__(self, hidden_dim, empty_grad=False): def forward(self, x, y): hidden_dim = x - if self.empty_grad and torch.distributed.get_rank() == 0: - hidden_dim = self.linear(hidden_dim) + self.linear2(hidden_dim) - else: - hidden_dim = self.linear(hidden_dim) + hidden_dim = self.linear(hidden_dim) return self.cross_entropy_loss(hidden_dim, y) diff --git a/tests/unit/test_fp16.py b/tests/unit/test_fp16.py index b2e76f0b7b82..0c0ef3edd3a8 100755 --- a/tests/unit/test_fp16.py +++ b/tests/unit/test_fp16.py @@ -856,3 +856,38 @@ def _go(args): model.step() _go(args=args) + + +@pytest.mark.parametrize('stage', [1, 2, 3]) +def test_zero_empty_grad(tmpdir, stage): + config_dict = { + "train_batch_size": 1, + "steps_per_print": 1, + "fp16": { + "enabled": True + }, + "zero_optimization": { + "stage": stage + } + } + args = args_from_dict(tmpdir, config_dict) + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + + @distributed_test(world_size=[1]) + def _go(args, model, hidden_dim): + optimizer = torch.optim.Adam(model.parameters()) + model, _, _, _ = deepspeed.initialize(args=args, + model=model, + optimizer=optimizer) + data_loader = random_dataloader(model=model, + total_samples=50, + hidden_dim=hidden_dim, + device=model.device) + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + _go(args=args, model=model, hidden_dim=hidden_dim) diff --git a/tests/unit/test_zero_context.py b/tests/unit/test_zero_context.py index 8b8d51131503..98ee0c7ad00b 100644 --- a/tests/unit/test_zero_context.py +++ b/tests/unit/test_zero_context.py @@ -281,3 +281,32 @@ def test_stage_3_output_type(output_type): loss = loss['loss'] engine.backward(loss) engine.step() + + +class ConvX(torch.nn.Conv1d): + def __init__(self, *args): + super().__init__(*args) + # This would not be partitioned before bugfix 5ca8167 + self.param_in = torch.nn.Parameter(torch.FloatTensor(5).uniform_()) + + def forward(self, x): + return x + + +class ConvNet(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = ConvX(1, 3, 4) + self.param = torch.nn.Parameter(torch.FloatTensor(5).uniform_()) + + def forward(self, x): + return x + + +def test_subclass_param(): + setup_serial_env() + with deepspeed.zero.Init(config=config): + model = ConvNet() + + assert model.param.ds_status == ZeroParamStatus.NOT_AVAILABLE + assert model.conv1.param_in.ds_status == ZeroParamStatus.NOT_AVAILABLE diff --git a/version.txt b/version.txt index 267577d47e49..2b7c5ae01848 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -0.4.1 +0.4.2