ZeRO-3: adaptive prefetch bucket size

deepspeed/runtime/zero/config.py

@@ -218,6 +218,25 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
     ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference.
     """
 
+    adaptive_prefetch_bucket_size: bool = Field(False, alias="stage3_adaptive_prefetch_bucket_size")
+    """
+    Dynamically adjust the prefetch bucket size at runtime based on observed fetch-wait and compute
+    times. When enabled, ``prefetch_bucket_size`` is used as the starting point and is clamped
+    between ``adaptive_prefetch_min_size`` and ``adaptive_prefetch_max_size``. Disabled by default.
+    """
+
+    adaptive_prefetch_min_size: int = Field(pp_int(1e7), ge=0, alias="stage3_adaptive_prefetch_min_size")
+    """
+    Lower bound for the adaptive prefetch bucket size (in parameter elements). Only used when
+    ``adaptive_prefetch_bucket_size`` is enabled.
+    """
+
+    adaptive_prefetch_max_size: int = Field(pp_int(5e8), ge=0, alias="stage3_adaptive_prefetch_max_size")
+    """
+    Upper bound for the adaptive prefetch bucket size (in parameter elements). Only used when
+    ``adaptive_prefetch_bucket_size`` is enabled.
+    """
+
     param_persistence_threshold: int = Field(pp_int(1e5), ge=0, alias="stage3_param_persistence_threshold")
     """
     Do not partition parameters smaller than this threshold. Smaller values use

deepspeed/runtime/zero/parameter_offload.py

@@ -186,6 +186,11 @@ def __init__(
         self._prefetch_bucket_sz = int(prefetch_bucket_size)
         self._max_reuse_distance_in_numel = int(max_reuse_distance)
         self._max_available_parameters_in_numel = int(max_live_parameters)
+
+        zero_cfg = getattr(ds_config, 'zero_config', None)
+        self._adaptive_prefetch = getattr(zero_cfg, 'adaptive_prefetch_bucket_size', False)
+        self._adaptive_prefetch_min_sz = int(getattr(zero_cfg, 'adaptive_prefetch_min_size', 10_000_000))
+        self._adaptive_prefetch_max_sz = int(getattr(zero_cfg, 'adaptive_prefetch_max_size', 500_000_000))
         self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream(
         ) if overlap_comm else get_accelerator().default_stream()
 
@@ -215,6 +220,9 @@ def __init__(
             zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights,
             fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module,
             log_trace_cache_warnings=self.log_trace_cache_warnings,
+            adaptive_prefetch=self._adaptive_prefetch,
+            adaptive_prefetch_min_sz=self._adaptive_prefetch_min_sz,
+            adaptive_prefetch_max_sz=self._adaptive_prefetch_max_sz,
         )
 
         self.forward_hooks = []

deepspeed/runtime/zero/partitioned_param_coordinator.py

@@ -7,6 +7,7 @@
 import collections
 from collections import UserDict
 import threading
+import time
 from typing import Deque, Dict, Set
 
 from deepspeed import comm as dist
@@ -99,6 +100,9 @@ def __init__(
         zero_quantized_nontrainable_weights=False,
         fast_sharding_for_leaf_module=False,
         log_trace_cache_warnings=False,
+        adaptive_prefetch: bool = False,
+        adaptive_prefetch_min_sz: int = 10_000_000,
+        adaptive_prefetch_max_sz: int = 500_000_000,
     ) -> None:
         # mapping of param -> handle for each param that is currently in flight
         self.__inflight_param_registry = inflight_param_registry
@@ -122,6 +126,21 @@ def __init__(
         self.__prefetch_bucket_sz: int = prefetch_bucket_sz
         self.__prefetch_nvme: bool = prefetch_nvme
         self.hierarchy: int = 0
+
+        # Adaptive prefetch: adjusts __prefetch_bucket_sz based on measured
+        # fetch-wait time vs compute time using exponential moving averages.
+        self.__adaptive_prefetch: bool = adaptive_prefetch
+        self.__adaptive_prefetch_min_sz: int = adaptive_prefetch_min_sz
+        self.__adaptive_prefetch_max_sz: int = adaptive_prefetch_max_sz
+        # EMA of wall-clock time spent waiting for in-flight fetches to complete
+        self.__fetch_wait_ema: float = 0.0
+        # EMA of wall-clock time between consecutive fetch_sub_module calls
+        # (approximates compute time per module)
+        self.__inter_step_ema: float = 0.0
+        self.__last_step_time: float = 0.0
+        self.__adaptive_step_count: int = 0
+        # Update the bucket size every this many steps to avoid thrashing
+        self.__adaptive_update_interval: int = 10
         self.zero_quantized_weights = zero_quantized_weights
         self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights
 
@@ -192,6 +211,56 @@ def _clean_inflight_param_registry(self) -> None:
             self.__release_param(param)
         self.__inflight_param_registry.clear()
 
+    def __update_adaptive_prefetch(self, wait_t0: float) -> None:
+        """Update EMAs and periodically resize the prefetch bucket.
+
+        wait_t0 is the perf_counter() timestamp taken just before entering the
+        per-param wait loop.  The elapsed time since then is the fetch-wait
+        cost for this step.  The elapsed time since the *previous* call is a
+        proxy for the module's compute time (it includes both compute and any
+        wait, so it slightly over-estimates pure compute; the ratio still
+        captures the trend).
+        """
+        now = time.perf_counter()
+        wait_time = now - wait_t0
+
+        # Compute inter-step interval only after the first call.
+        if self.__last_step_time > 0.0:
+            inter_step = now - self.__last_step_time
+            alpha = 0.1  # EMA smoothing factor
+            self.__fetch_wait_ema = alpha * wait_time + (1 - alpha) * self.__fetch_wait_ema
+            self.__inter_step_ema = alpha * inter_step + (1 - alpha) * self.__inter_step_ema
+        self.__last_step_time = now
+
+        self.__adaptive_step_count += 1
+        if self.__adaptive_step_count % self.__adaptive_update_interval != 0:
+            return
+
+        if self.__inter_step_ema <= 0.0:
+            return
+
+        # Ratio of wait time to total step time.  A high ratio means the GPU is
+        # blocking waiting for prefetched data — increase the bucket.  A low
+        # ratio means we're fetching more than needed — shrink to save memory.
+        wait_ratio = self.__fetch_wait_ema / self.__inter_step_ema
+        old_sz = self.__prefetch_bucket_sz
+        if wait_ratio > 0.15:
+            # Fetch wait is a meaningful fraction of the step — prefetch more.
+            new_sz = int(old_sz * 1.25)
+        elif wait_ratio < 0.05:
+            # Almost no wait time — we're over-fetching; back off a little.
+            new_sz = int(old_sz * 0.90)
+        else:
+            return
+
+        new_sz = max(self.__adaptive_prefetch_min_sz, min(self.__adaptive_prefetch_max_sz, new_sz))
+        if new_sz != old_sz:
+            logger.info(
+                f"[ZeRO-3 adaptive prefetch] bucket size {old_sz} -> {new_sz} "
+                f"(wait_ratio={wait_ratio:.3f}, wait_ema={self.__fetch_wait_ema*1e3:.2f}ms, "
+                f"step_ema={self.__inter_step_ema*1e3:.2f}ms)")
+            self.__prefetch_bucket_sz = new_sz
+
     def _invalidate_trace(self) -> None:
         if self.is_invalid_trace():
             raise RuntimeError("attempted to invalidate already invalid trace")
@@ -387,6 +456,7 @@ def _fetch_sub_module_impl(self, current_submodule: Module, forward: bool, is_le
         fast_fetch = self.fast_sharding_for_leaf_module and is_leaf
         # wait for parameters in the immediately needed submodule to become available
         in_checkpoint_recompute = forward and torch._C._current_graph_task_id() != -1
+        _wait_t0 = time.perf_counter() if self.__adaptive_prefetch else 0.0
         for param in params_to_fetch:
             param.ds_active_sub_modules.add(current_submodule.ds_id)
             # Only frozen params need recompute attribution; trainable ones release via their backward hook.
@@ -417,6 +487,9 @@ def _fetch_sub_module_impl(self, current_submodule: Module, forward: bool, is_le
             AllGatherCoalescedHandle.free_buffer()
         self.__profiler.stop_event(wait_event_name, wait_numel)
 
+        if self.__adaptive_prefetch:
+            self.__update_adaptive_prefetch(_wait_t0)
+
         # kick off parameter prefetches for upcoming modules
         # don't prefetch if we dont have a completed model trace
         if self.is_complete_trace():

tests/unit/runtime/zero/test_zero3_adaptive_prefetch.py

@@ -0,0 +1,132 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+"""Unit tests for ZeRO-3 adaptive prefetch bucket size.
+
+The adaptation logic lives entirely in PartitionedParameterCoordinator and
+uses only CPU wall-clock time, so these tests run without a GPU.
+"""
+
+import time
+import pytest
+from unittest.mock import patch
+
+from deepspeed.runtime.zero.partitioned_param_coordinator import (
+    InflightParamRegistry,
+    PartitionedParameterCoordinator,
+)
+
+
+def _make_coordinator(prefetch_bucket_sz=50_000_000,
+                      adaptive_prefetch=True,
+                      adaptive_prefetch_min_sz=10_000_000,
+                      adaptive_prefetch_max_sz=500_000_000):
+    """Instantiate a coordinator with no GPU resources for unit testing."""
+    return PartitionedParameterCoordinator(
+        prefetch_bucket_sz=prefetch_bucket_sz,
+        max_reuse_distance_in_numel=int(1e9),
+        max_available_parameters_in_numel=int(1e9),
+        allgather_stream=None,
+        inflight_param_registry=InflightParamRegistry(),
+        prefetch_nvme=False,
+        timers=None,
+        adaptive_prefetch=adaptive_prefetch,
+        adaptive_prefetch_min_sz=adaptive_prefetch_min_sz,
+        adaptive_prefetch_max_sz=adaptive_prefetch_max_sz,
+    )
+
+
+def _drive_adaptation(coordinator, wait_ratio, n_steps=10):
+    """Simulate n_steps calls to __update_adaptive_prefetch with a fixed wait/step ratio.
+
+    wait_ratio is the fraction of each step that looks like fetch-wait time.
+    """
+    step_duration = 0.01  # 10 ms per step
+    wait_duration = step_duration * wait_ratio
+
+    update_fn = coordinator._PartitionedParameterCoordinator__update_adaptive_prefetch
+
+    t = time.perf_counter()
+    for _ in range(n_steps):
+        # Simulate: wait started `wait_duration` seconds before "now"
+        wait_t0 = t - wait_duration
+        with patch("deepspeed.runtime.zero.partitioned_param_coordinator.time") as mock_time:
+            mock_time.perf_counter.return_value = t
+            update_fn(wait_t0)
+        t += step_duration
+
+
+class TestAdaptivePrefetchConfig:
+
+    def test_disabled_by_default(self):
+        coord = _make_coordinator(adaptive_prefetch=False)
+        assert coord._PartitionedParameterCoordinator__adaptive_prefetch == False
+
+    def test_enabled(self):
+        coord = _make_coordinator(adaptive_prefetch=True)
+        assert coord._PartitionedParameterCoordinator__adaptive_prefetch == True
+
+    def test_bounds_stored(self):
+        coord = _make_coordinator(adaptive_prefetch_min_sz=5_000_000,
+                                  adaptive_prefetch_max_sz=200_000_000)
+        assert coord._PartitionedParameterCoordinator__adaptive_prefetch_min_sz == 5_000_000
+        assert coord._PartitionedParameterCoordinator__adaptive_prefetch_max_sz == 200_000_000
+
+
+class TestAdaptivePrefetchLogic:
+
+    def test_high_wait_ratio_grows_bucket(self):
+        """When wait time is >15% of step time, bucket should grow."""
+        initial = 50_000_000
+        coord = _make_coordinator(prefetch_bucket_sz=initial)
+        # Drive 20 steps at 30% wait ratio so EMAs converge enough to trigger
+        _drive_adaptation(coord, wait_ratio=0.30, n_steps=20)
+        sz = coord._PartitionedParameterCoordinator__prefetch_bucket_sz
+        assert sz > initial, f"expected bucket to grow from {initial}, got {sz}"
+
+    def test_low_wait_ratio_shrinks_bucket(self):
+        """When wait time is <5% of step time, bucket should shrink."""
+        initial = 50_000_000
+        coord = _make_coordinator(prefetch_bucket_sz=initial)
+        _drive_adaptation(coord, wait_ratio=0.01, n_steps=20)
+        sz = coord._PartitionedParameterCoordinator__prefetch_bucket_sz
+        assert sz < initial, f"expected bucket to shrink from {initial}, got {sz}"
+
+    def test_mid_ratio_leaves_bucket_unchanged(self):
+        """Wait ratio in [5%, 15%] should not trigger a resize."""
+        initial = 50_000_000
+        coord = _make_coordinator(prefetch_bucket_sz=initial)
+        _drive_adaptation(coord, wait_ratio=0.09, n_steps=20)
+        sz = coord._PartitionedParameterCoordinator__prefetch_bucket_sz
+        assert sz == initial, f"expected bucket unchanged at {initial}, got {sz}"
+
+    def test_bucket_clamped_to_min(self):
+        """Bucket should never drop below adaptive_prefetch_min_sz."""
+        min_sz = 40_000_000
+        coord = _make_coordinator(prefetch_bucket_sz=min_sz + 1_000_000,
+                                  adaptive_prefetch_min_sz=min_sz)
+        # Very low wait ratio drives the bucket down
+        _drive_adaptation(coord, wait_ratio=0.001, n_steps=100)
+        sz = coord._PartitionedParameterCoordinator__prefetch_bucket_sz
+        assert sz >= min_sz, f"bucket {sz} dropped below min {min_sz}"
+
+    def test_bucket_clamped_to_max(self):
+        """Bucket should never exceed adaptive_prefetch_max_sz."""
+        max_sz = 60_000_000
+        coord = _make_coordinator(prefetch_bucket_sz=max_sz - 1_000_000,
+                                  adaptive_prefetch_max_sz=max_sz)
+        # Very high wait ratio drives the bucket up
+        _drive_adaptation(coord, wait_ratio=0.99, n_steps=100)
+        sz = coord._PartitionedParameterCoordinator__prefetch_bucket_sz
+        assert sz <= max_sz, f"bucket {sz} exceeded max {max_sz}"
+
+    def test_no_update_before_interval(self):
+        """Bucket should not change before the update interval (10 steps) is reached."""
+        initial = 50_000_000
+        coord = _make_coordinator(prefetch_bucket_sz=initial)
+        # Only 5 steps — not enough to trigger a resize even at high wait ratio
+        _drive_adaptation(coord, wait_ratio=0.99, n_steps=5)
+        sz = coord._PartitionedParameterCoordinator__prefetch_bucket_sz
+        assert sz == initial, f"bucket changed before update interval, got {sz}"

tests/unit/runtime/zero/test_zero_config.py

@@ -69,6 +69,24 @@ def test_zero_config_offload_configs():
     assert isinstance(config.offload_optimizer, DeepSpeedZeroOffloadOptimizerConfig)
 
 
+def test_zero_config_adaptive_prefetch_defaults():
+    config = DeepSpeedZeroConfig()
+    assert config.adaptive_prefetch_bucket_size == False
+    assert config.adaptive_prefetch_min_size == int(1e7)
+    assert config.adaptive_prefetch_max_size == int(5e8)
+
+
+def test_zero_config_adaptive_prefetch_aliases():
+    config = DeepSpeedZeroConfig(**{"stage3_adaptive_prefetch_bucket_size": True})
+    assert config.adaptive_prefetch_bucket_size == True
+
+    config = DeepSpeedZeroConfig(**{"stage3_adaptive_prefetch_min_size": 5_000_000})
+    assert config.adaptive_prefetch_min_size == 5_000_000
+
+    config = DeepSpeedZeroConfig(**{"stage3_adaptive_prefetch_max_size": 1_000_000_000})
+    assert config.adaptive_prefetch_max_size == 1_000_000_000
+
+
 def test_zero_offload_optimizer_config_pipeline():
     config = DeepSpeedZeroOffloadOptimizerConfig()
     assert config.pipeline == False