diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 9f3554fbb78c..476f71de2120 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -300,9 +300,8 @@ def __init__( monitor_memory=False, ) - if dist.get_rank() == 0: - logger.info( - f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}") + log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", + ranks=[0]) if self.flops_profiler_enabled(): self.flops_profiler = FlopsProfiler(self.module, self) @@ -688,6 +687,9 @@ def zero_prefetch_bucket_size(self): def zero_param_persistence_threshold(self): return self._config.zero_config.param_persistence_threshold + def zero_model_persistence_threshold(self): + return self._config.zero_config.model_persistence_threshold + def zero_gather_16bit_weights_on_model_save(self): return self._config.zero_config.gather_16bit_weights_on_model_save @@ -779,18 +781,17 @@ def _configure_lr_scheduler(self, client_lr_scheduler): # First check for scheduler in json configuration lr_scheduler = self._scheduler_from_config(self.optimizer) if lr_scheduler: - if self.global_rank == 0: - logger.info( - f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}") + log_dist( + f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}", + ranks=[0]) self.lr_scheduler = lr_scheduler else: if isinstance(client_lr_scheduler, Callable): - if self.global_rank == 0: - logger.info('DeepSpeed using client callable to create LR scheduler') + log_dist('DeepSpeed using client callable to create LR scheduler', + ranks=[0]) self.lr_scheduler = client_lr_scheduler(self.basic_optimizer) else: - if self.global_rank == 0: - logger.info('DeepSpeed using client LR scheduler') + log_dist('DeepSpeed using client LR scheduler', ranks=[0]) self.lr_scheduler = client_lr_scheduler log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) @@ -1093,31 +1094,26 @@ def _configure_optimizer(self, client_optimizer, model_parameters): client_optimizer.param_groups[:] = [ pg for pg in client_optimizer.param_groups if len(pg["params"]) != 0 ] - if self.global_rank == 0: - logger.info( - "Removing param_group that has no 'params' in the client Optimizer" - ) + log_dist( + "Removing param_group that has no 'params' in the client Optimizer", + ranks=[0]) basic_optimizer = client_optimizer - if self.global_rank == 0: - logger.info('Using client Optimizer as basic optimizer') + log_dist('Using client Optimizer as basic optimizer', ranks=[0]) else: basic_optimizer = client_optimizer(model_parameters) - if self.global_rank == 0: - logger.info('Using client callable to create basic optimizer') + log_dist('Using client callable to create basic optimizer', ranks=[0]) else: basic_optimizer = self._configure_basic_optimizer(model_parameters) - if self.global_rank == 0: - logger.info( - "Using DeepSpeed Optimizer param name {} as basic optimizer".format( - self.optimizer_name())) + log_dist( + f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", + ranks=[0]) self._check_for_duplicates(basic_optimizer) self.basic_optimizer = basic_optimizer - if self.global_rank == 0: - logger.info("DeepSpeed Basic Optimizer = {}".format( - basic_optimizer.__class__.__name__)) + log_dist("DeepSpeed Basic Optimizer = {basic_optimizer.__class__.__name__}", + ranks=[0]) if self.zero_optimization(): assert ( @@ -1138,8 +1134,7 @@ def _configure_optimizer(self, client_optimizer, model_parameters): elif self.amp_enabled(): assert not (self.fp16_enabled() or self.bfloat16_enabled()), "Cannot enable both amp with (legacy) fp16 or bfloat16 mode" amp_params = self.amp_params() - if self.global_rank == 0: - logger.info(f"Initializing AMP with these params: {amp_params}") + log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) try: logger.info("Initializing Apex amp from: {}".format(amp.__path__)) except NameError: @@ -1340,8 +1335,8 @@ def _configure_bf16_optimizer(self, optimizer): if optimizer is None: optimizer = DummyOptim(list(self.module.parameters())) - if self.global_rank == 0: - logger.info('Creating unfused BF16 optimizer') + log_dist('Creating BF16 optimizer', ranks=[0]) + timers = self.timers if self.wall_clock_breakdown() else None optimizer = BF16_Optimizer( optimizer, @@ -1356,7 +1351,6 @@ def _configure_bf16_optimizer(self, optimizer): def _configure_zero_optimizer(self, optimizer): zero_stage = self.zero_optimization_stage() - log_dist('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage), ranks=[0]) assert self.communication_data_type in (torch.float16, torch.bfloat16), "ZeRO supports only 'communication_data_type': ['fp16', 'bfp16']" timers = self.timers if self.wall_clock_breakdown() else None @@ -1374,6 +1368,8 @@ def _configure_zero_optimizer(self, optimizer): round_robin_gradients = self.zero_round_robin_gradients() assert not isinstance(optimizer, DummyOptim), "zero stage 2 requires an optimizer" + log_dist('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage), + ranks=[0]) # Overlap and contiguous grads are meaningless in stage 1 and are ignored if zero_stage == ZeroStageEnum.optimizer_states: overlap_comm = False @@ -1419,10 +1415,8 @@ def _configure_zero_optimizer(self, optimizer): elif zero_stage == ZeroStageEnum.weights: assert not self.has_moe_layers, "MoE not supported with Stage 3" - logger.info("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None - from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 - if isinstance(optimizer, DummyOptim): + log_dist("Creating ZeRO Offload", ranks=[0]) optimizer = DeepSpeedZeRoOffload( self.module, timers=timers, @@ -1432,10 +1426,13 @@ def _configure_zero_optimizer(self, optimizer): max_reuse_distance=self.zero_max_reuse_distance(), max_live_parameters=self.zero_max_live_parameters(), param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), offload_param_config=self.zero_offload_param(), mpu=self.mpu) else: - + log_dist('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage), + ranks=[0]) + from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 optimizer = DeepSpeedZeroOptimizer_Stage3( self.module, optimizer, @@ -1451,6 +1448,7 @@ def _configure_zero_optimizer(self, optimizer): max_reuse_distance=self.zero_max_reuse_distance(), max_live_parameters=self.zero_max_live_parameters(), param_persistence_threshold=self.zero_param_persistence_threshold(), + model_persistence_threshold=self.zero_model_persistence_threshold(), dp_process_group=self.data_parallel_group, reduce_scatter=self.zero_reduce_scatter(), overlap_comm=self.zero_overlap_comm(), diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py index aaca3cbda887..7da8824ed3a6 100644 --- a/deepspeed/runtime/zero/config.py +++ b/deepspeed/runtime/zero/config.py @@ -4,6 +4,7 @@ """ from pydantic import Field, validator +import sys from typing import Optional from enum import Enum from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigModel @@ -114,6 +115,9 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel): param_persistence_threshold: int = Field(1e5, ge=0, alias="stage3_param_persistence_threshold") + model_persistence_threshold: int = Field(sys.maxsize, + ge=0, + alias="stage3_model_persistence_threshold") max_live_parameters: int = Field(1e9, ge=0, alias="stage3_max_live_parameters") max_reuse_distance: int = Field(1e9, ge=0, alias="stage3_max_reuse_distance") gather_16bit_weights_on_model_save: bool = Field( diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index e56043b01167..cd62015cf5a6 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/deepspeed/runtime/zero/parameter_offload.py @@ -3,6 +3,7 @@ Licensed under the MIT license. """ +import sys import torch from torch.cuda import Stream from collections import OrderedDict @@ -173,10 +174,11 @@ def __init__(self, max_reuse_distance=1000000000, max_live_parameters=1000000000, param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, offload_param_config=None, mpu=None): - see_memory_usage("TensorOffload initialize beginning", force=True) + see_memory_usage("DeepSpeedZeRoOffload initialize [begin]", force=True) print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) @@ -196,8 +198,11 @@ def __init__(self, _inject_parameters(module, ZeROOrderedDict) - self.persistence_threshold = int(param_persistence_threshold) - self.persistent_parameters = self.mark_persistent_parameters() + self.param_numel_persistence_threshold = int(param_persistence_threshold) + self.model_persistence_threshold = int(model_persistence_threshold) + self.persistent_parameters = self.mark_persistent_parameters( + self.param_numel_persistence_threshold, + self.model_persistence_threshold) self.param_coordinators = {} self._prefetch_bucket_sz = int(prefetch_bucket_size) @@ -213,6 +218,8 @@ def __init__(self, f'Created module hooks: forward = {len(self.forward_hooks)}, backward = {len(self.backward_hooks)}', force=False) + see_memory_usage("DeepSpeedZeRoOffload initialize [end]", force=True) + @instrument_w_nvtx def partition_all_parameters(self): """Partitioning Parameters that were not partitioned usually if parameters @@ -291,12 +298,15 @@ def _end_of_forward_hook(module, *args): global FWD_MODULE_STACK FWD_MODULE_STACK.append(self.module) - def mark_persistent_parameters(self): + def mark_persistent_parameters(self, param_threshold, model_threshold): persistent_params = [] total_persistent_parameters = 0 params_count = 0 for _, param in self.module.named_parameters(recurse=True): - if param.ds_numel < self.persistence_threshold: + if param.ds_numel + total_persistent_parameters > model_threshold: + continue + + if param.ds_numel < param_threshold: params_count += 1 param.ds_persist = True persistent_params.append(param) @@ -304,7 +314,7 @@ def mark_persistent_parameters(self): print_rank_0( f"Parameter Offload: Total persistent parameters: {total_persistent_parameters} in {params_count} params", - force=False) + force=True) return persistent_params diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index db135e0cf182..b6bd5ed645f9 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -683,9 +683,13 @@ def get_model(): # Remote device is the device where parameter partitions 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 - == OffloadDeviceEnum.cpu) else False + self.remote_device = self.local_device if remote_device in [ + None, + OffloadDeviceEnum.none + ] else remote_device + self.pin_memory = pin_memory if ( + self.remote_device in [OffloadDeviceEnum.cpu, + OffloadDeviceEnum.nvme]) else False # Enable fp16 param swapping to NVMe if self.remote_device == OffloadDeviceEnum.nvme: diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py index 6a923805eb90..1dcff3f1c12f 100644 --- a/deepspeed/runtime/zero/partitioned_param_coordinator.py +++ b/deepspeed/runtime/zero/partitioned_param_coordinator.py @@ -400,6 +400,16 @@ def __all_gather_params(self, params: Set[Parameter]) -> None: assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary() self.__inflight_param_registry[param] = handle + # Release swap buffers for persisted params on nvme since they will never be partitioned or evicted from GPU + swap_persisted_params = [ + p for p in partitioned_params + if p.ds_persist and p.ds_tensor.final_location == OffloadDeviceEnum.nvme + ] + if swap_persisted_params: + swap_persisted_params[ + 0].nvme_swapper.remove_partition_and_release_buffers( + swap_persisted_params) + @instrument_w_nvtx def __release_param(self, param: Parameter) -> None: if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules: diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index fd1322ad2f77..9c603be0d2d2 100755 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -3,6 +3,7 @@ Licensed under the MIT license. """ +import sys import gc import collections from typing import Deque, Dict, Tuple @@ -88,6 +89,7 @@ def __init__(self, max_reuse_distance=1000000000, max_live_parameters=1000000000, param_persistence_threshold=100000, + model_persistence_threshold=sys.maxsize, dp_process_group=None, reduce_scatter=True, overlap_comm=False, @@ -146,15 +148,18 @@ def __init__(self, self.params_in_nvme_and_cpu = False self.max_params_in_cpu = 0 - self.parameter_offload = DeepSpeedZeRoOffload(module, - timers, - ds_config, - overlap_comm, - prefetch_bucket_size, - max_reuse_distance, - max_live_parameters, - param_persistence_threshold, - offload_param_config) + self.parameter_offload = DeepSpeedZeRoOffload( + module=module, + timers=timers, + ds_config=ds_config, + overlap_comm=overlap_comm, + prefetch_bucket_size=prefetch_bucket_size, + max_reuse_distance=max_reuse_distance, + max_live_parameters=max_live_parameters, + param_persistence_threshold=param_persistence_threshold, + model_persistence_threshold=model_persistence_threshold, + offload_param_config=offload_optimizer_config) + self.persistent_parameters = self.parameter_offload.persistent_parameters self._configure_offloading(offload_optimizer_config, offload_param_config)