ZeRO-3: adaptive prefetch bucket size - #8431
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c18a8e20bd
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| 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 |
There was a problem hiding this comment.
When adaptive prefetching is enabled in multi-rank training, each rank derives new_sz from its own unsynchronized wall-clock measurements and assigns it locally. Timing noise can therefore give ranks different bucket sizes; the subsequent prefetch loop selects a different parameter set and invokes all_gather_coalesced with different tensor sizes or collective sequences, which can hang or fail the job. Derive a common decision through the data-parallel group or broadcast the selected size before applying it.
Useful? React with 👍 / 👎.
| 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 |
There was a problem hiding this comment.
On the normal CUDA/NCCL asynchronous path, Work.wait() and wait_stream() establish stream dependencies without waiting for the GPU operation to finish on the host, so this perf_counter() interval primarily measures Python bookkeeping and enqueue latency rather than GPU time stalled on the all-gather. The inter-call interval similarly omits asynchronous device compute, making the ratio unrelated to the quantities the controller claims to compare and typically causing it to shrink the bucket despite GPU stalls. Use device events or another measurement that observes completion.
Useful? React with 👍 / 👎.
| if self.__adaptive_prefetch: | ||
| self.__update_adaptive_prefetch(_wait_t0) |
There was a problem hiding this comment.
For the initial invalid/recording passes, this call still updates the controller even though the immediately following is_complete_trace() guard prevents all prefetching. Models with at least ten module fetches therefore classify unavoidable demand-fetch waits as evidence that the prefetch bucket is too small and can drive it toward the maximum before prefetching has ever run. Only collect samples once the trace is complete, and reset the timing state when entering that mode.
Useful? React with 👍 / 👎.
| else: | ||
| return | ||
|
|
||
| new_sz = max(self.__adaptive_prefetch_min_sz, min(self.__adaptive_prefetch_max_sz, new_sz)) |
There was a problem hiding this comment.
When stage3_prefetch_bucket_size starts outside the configured adaptive range, it is clamped only after a grow or shrink branch. If the measured ratio remains between 0.05 and 0.15, the early return leaves the bucket outside its bounds indefinitely—for example, a 1-billion-element starting bucket with a 500-million maximum can retain the excessive allocation. Clamp the starting value during initialization or before the threshold return, and validate that the minimum does not exceed the maximum.
Useful? React with 👍 / 👎.
| ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference. | ||
| """ | ||
|
|
||
| adaptive_prefetch_bucket_size: bool = Field(False, alias="stage3_adaptive_prefetch_bucket_size") |
There was a problem hiding this comment.
This introduces three public configuration keys and a distributed runtime controller, but the commit changes no tests or user-facing documentation, and a repository-wide search finds no other references to the new keys. Add coverage for adaptation, trace lifecycle, bounds, and multi-rank agreement, plus document the options in the configuration reference, as required for new features.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
| ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference. | ||
| """ | ||
|
|
||
| adaptive_prefetch_bucket_size: bool = Field(False, alias="stage3_adaptive_prefetch_bucket_size") |
There was a problem hiding this comment.
This is a non-merge commit, but git show -s --format=%B 2c89e1ef519612d5d043a0b221a2059804907fe7 contains no Signed-off-by trailer, so it violates the repository's commit/DCO requirement. Recreate the commit metadata with --signoff before merging.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
c18a8e2 to
8115ae1
Compare
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 deepspeedai#8193 Signed-off-by: promptsmith1990 <transformertitan@gmail.com> Signed-off-by: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com>
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 <transformertitan@gmail.com> Signed-off-by: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com>
b965120 to
481c21a
Compare
|
Hi @promptsmith1990, I built a follow-up on this PR's current head ( The DCO-signed-off follow-up commit is 6e70233. It fixes raw engine-config plumbing (including inference), makes the resize decision consistent across data-parallel ranks, samples CUDA compute-stream events without synchronizing each fetch, resets stale trace samples, and enforces the initial/runtime bounds. The feature remains opt-in. It can be cherry-picked onto this branch; I have kept this as a follow-up to your PR. Validation on that exact source: 43 CPU/Gloo tests passed, both single-GPU overlap configurations passed, and 5 two-GPU NCCL tests passed. The native training cases compare 22 steps and final weights against an unpartitioned SGD reference. Added boundary regressions reproduce 3 failures on the original head. I also completed a BERT-base/SQuAD 1.1 fixed-subset comparison on 2 × RTX 4000 Ada (20 GiB each, PCIe, no NVLink), Python 3.12.14 / PyTorch 2.13.0+cu130 / Transformers 4.57.6. Each policy has three paired trials (seeds 1337/1338/1339), with identical initial weights, data and RNG; policy order rotates across seeds. BF16, ZeRO-3, overlap enabled, AdamW, global batch 4, sequence length 256, 20 warmup + 80 measured training steps. Training/validation subsets contain 512/64 QA examples.
All three policies have identical initial/final validation loss, identical global-mean loss at every training step, and bitwise-identical final BF16 parameters within each seed. Final validation losses are 3.011719 / 2.895264 / 3.077637 for seeds 1337/1338/1339, respectively. Both ranks use the same window throughout each run; the adaptive window grows to about 9.31M by the final training step. Relative to static 1M, the paired adaptive throughput differences are +1.48%, +4.95%, +2.25% (median +2.25%). Relative to static 50M, they are -0.88%, +0.93%, -1.67% (median -0.88%). There is no allocated-memory reduction versus static 1M, and reserved memory is 160 MiB higher. These measurements support correctness on this workload, not a general speedup claim. Measurement limits: this is a custom fixed-subset harness, not an execution of the full DeepSpeedExamples training script. Inputs are GPU-resident before timing; throughput includes indexing, training and controller work, excludes preprocessing/evaluation, and uses the slower rank. Memory is the measured PyTorch allocator peak, not every driver/NCCL allocation. Controller overhead is not separately isolated. Initial validation advances the controller cadence before warmup. This is a short-run comparison, not a full convergence result. Would you be open to incorporating this follow-up and collaborating on the remaining design/validation? If you prefer to squash it into your work, please retain my contribution as |
2 participants
Problem
The ZeRO-3 prefetch bucket (
stage3_prefetch_bucket_size) is set once at startupand never changes. Getting it right requires trial and error: too small and the GPU
sits idle waiting for allgathers to finish; too large and it wastes memory and can
actually hurt throughput by crowding out live parameters.
Issue #8193 asked for a way to adapt the bucket automatically.
Solution
This PR adds an optional runtime adaptation loop inside
PartitionedParameterCoordinatorthat adjusts__prefetch_bucket_szbased onobserved timing:
allgathers). High values mean we didn't prefetch enough.
_fetch_sub_module_implcalls.Used as a proxy for per-module compute time.
Every 10 steps the ratio
wait_ema / step_emais checked:> 0.15→ grow the bucket by 25% (GPU is stalling, need more prefetch)< 0.05→ shrink by 10% (barely any wait, we're over-fetching)The bucket is clamped to
[min_size, max_size]to stay within sensible bounds.Both EMAs use α=0.1, so the controller reacts to trends rather than noise.
New config keys (all opt-in, backward-compatible defaults)
stage3_prefetch_bucket_sizecontinues to work as before and is used as thestarting value when adaptation is enabled.
Files changed
deepspeed/runtime/zero/config.py— three newDeepSpeedZeroConfigfieldsdeepspeed/runtime/zero/parameter_offload.py— reads config, passes to coordinatordeepspeed/runtime/zero/partitioned_param_coordinator.py— EMA tracking and__update_adaptive_prefetchmethod;import timeadded; timing point addedaround the fetch-wait loop
Notes
time.perf_counter()is CPU-only and adds no CUDA synchronisationResolves #8193