Release swap buffers for persisted params

deepspeed/runtime/engine.py

@@ -678,6 +678,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
 
@@ -1343,7 +1346,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
 
@@ -1361,6 +1363,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
@@ -1406,10 +1410,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):
+                logger.info("Creating ZeRO Offload") if dist.get_rank() == 0 else None
                 optimizer = DeepSpeedZeRoOffload(
                     self.module,
                     timers=timers,
@@ -1419,10 +1421,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,
@@ -1438,6 +1443,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(),

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(

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,20 +298,23 @@ 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)
                 total_persistent_parameters += param.ds_numel
 
         print_rank_0(
             f"Parameter Offload: Total persistent parameters: {total_persistent_parameters} in {params_count} params",
-            force=False)
+            force=True)
 
         return persistent_params
 

deepspeed/runtime/zero/partition_parameters.py

@@ -669,9 +669,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:

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:

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)