From 55fa522ef9085b83ac9a7792c29f1696d547e0ea Mon Sep 17 00:00:00 2001 From: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:18:28 -0700 Subject: [PATCH 1/2] ZeRO-3: adaptive prefetch bucket size The prefetch bucket size is currently a static value set once at startup. If it's too small, the GPU stalls waiting for parameters; if it's too large, it wastes memory and may hurt throughput. This adds a runtime adaptation loop that tracks two exponential moving averages per training step: - time spent blocked waiting for in-flight allgathers to finish - total inter-step wall time (proxy for compute time) When the wait/step ratio exceeds 15% the bucket grows by 25%; when it drops below 5% the bucket shrinks by 10%. The bucket is clamped to [stage3_adaptive_prefetch_min_size, stage3_adaptive_prefetch_max_size] and is only adjusted every 10 steps to avoid thrashing. The feature is opt-in: set stage3_adaptive_prefetch_bucket_size=true in the ZeRO config. The existing stage3_prefetch_bucket_size is still used as the starting value and remains the only knob needed when disabled. Closes #8193 Signed-off-by: promptsmith1990 Signed-off-by: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com> --- deepspeed/runtime/zero/config.py | 19 +++++ deepspeed/runtime/zero/parameter_offload.py | 8 ++ .../zero/partitioned_param_coordinator.py | 73 +++++++++++++++++++ 3 files changed, 100 insertions(+) diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py index f696152a3843..11482dec089e 100644 --- a/deepspeed/runtime/zero/config.py +++ b/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 diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index f9fc7c2c0e64..f414e5b061a0 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/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 = [] diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py index 68dd3e654d72..29dcab81fdf3 100644 --- a/deepspeed/runtime/zero/partitioned_param_coordinator.py +++ b/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(): From 481c21a7a72fcdf6752d6c50bbf9c1d63cf9d53d Mon Sep 17 00:00:00 2001 From: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:17:35 -0700 Subject: [PATCH 2/2] tests: adaptive prefetch bucket size for ZeRO-3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two test additions: 1. test_zero_config.py — two new cases checking that the three new adaptive prefetch config fields parse correctly from both their canonical names and their stage3_* JSON aliases. 2. test_zero3_adaptive_prefetch.py — unit tests for the EMA adaptation logic in PartitionedParameterCoordinator. All tests run on CPU only (no GPU required). They drive __update_adaptive_prefetch directly with mocked perf_counter values to verify: - high wait ratio grows the bucket - low wait ratio shrinks the bucket - mid-range ratio leaves bucket unchanged - bucket is clamped to [min, max] bounds - no resize before the 10-step update interval Signed-off-by: promptsmith1990 Signed-off-by: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com> --- .../zero/test_zero3_adaptive_prefetch.py | 132 ++++++++++++++++++ tests/unit/runtime/zero/test_zero_config.py | 18 +++ 2 files changed, 150 insertions(+) create mode 100644 tests/unit/runtime/zero/test_zero3_adaptive_prefetch.py diff --git a/tests/unit/runtime/zero/test_zero3_adaptive_prefetch.py b/tests/unit/runtime/zero/test_zero3_adaptive_prefetch.py new file mode 100644 index 000000000000..d253535b7cb6 --- /dev/null +++ b/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}" diff --git a/tests/unit/runtime/zero/test_zero_config.py b/tests/unit/runtime/zero/test_zero_config.py index 5ced5f69170d..74ea885cf55a 100644 --- a/tests/unit/runtime/zero/test_zero_config.py +++ b/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