Stabilize DeepCompile ZeRO-3 memory scheduling by tohtana · Pull Request #8169 · deepspeedai/DeepSpeed

tohtana · 2026-07-23T20:53:11Z

Depends on #8159.

Overview

DeepCompile's ZeRO-3 scheduler does not consistently account for memory pressure across gather/release ordering, prefetch, selective gathering, and native gather-buffer reuse. This can retain excess gathered storage or trigger a full-parameter gather while Dynamo evaluates guards.

This PR adds rank-consistent scheduler budgeting and diagnostics, coordinates prefetch and selective gathering with graph profiling, reuses native gather storage within a bounded pressure-aware lifecycle, and avoids guard-time full-parameter gathers while the DeepCompile eager fallback is active.

Results

The end-to-end measurements were collected.

Model: Qwen3-14B
Hardware: one node with 8 H100 GPUs
Activation checkpointing: enabled
Micro-batch size: 4
Sequence length: 4096
Measurement: 10 warmup iterations followed by 20 measured iterations

Mode	Allocator retries per rank	Average step time	Throughput
ZeRO-3-only pass	18–19 → 0	8.4% faster	9.2% higher
Full DeepCompile pass	20 → 2–4	9.2% faster	10.1% higher

In a longer observation window covering steps 31–300, no allocator retries occurred.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 55779855e3

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

pengdurice · 2026-08-04T05:06:47Z

    m.def("register_graph_z3",
          &dc::register_graph_z3,
          "Register graph with a list of ds parameter ids");
+    m.def("set_z3_gather_buffer_pool_budget_for_test",


wondering if these new items are just test purpose only?

Yes, these three new APIs are for testing purpose only. We need them to set or inspect the internal state of the memory pool. It doesn't look great to have these as public APIs, but I didn't have an idea about other approaches.

just curious if it is possible to do integration test rather than unit test on this level? For example, instead of writing the scaffold to test GatherBufferPool, can do a test of functions that call the natural creation of this class, e.g. calling from Z3CustomOpExecutor?

Yes, that is actually how the current tests are structured. They do not construct or call GatherBufferPool directly. register_graph_z3() creates Z3CustomOpExecutor with the shared pool, and the CUDA test runs the real allgather_param, wait_allgather, and release_param ops.
The _for_test APIs are limited to making the budget and allocator-pressure inputs deterministic and exposing accounting of internal values.

OK, your last sentence explains my confusion. OK, maybe we can call it deterministic or pre determined size etc?

Yes. I think “fixed budget” is more precise because the test fixes the pool’s byte budget rather than individual buffer sizes. If we keep for_test, how about these

set_z3_gather_buffer_pool_budget_for_test -> set_z3_gather_buffer_pool_fixed_budget_for_test

setBudgetForTest -> setFixedBudgetForTest

update_z3_gather_buffer_pool_allocator_pressure_for_test -> update_z3_gather_buffer_pool_allocator_pressure_for_test

pengdurice · 2026-08-04T17:36:35Z

+        if (!enabled_ || pressure_recovery_in_progress_) { return at::Tensor(); }
+
+        Entry* best = nullptr;
+        for (auto& entry : entries_) {


is entries_ ordered? if not, is the full loop expensive vs an max heap like solution?

entries_ is currently an unsorted std::vector as we don't simply choose the maximum one (best fit to capacity, LRU for eviction, search by storage identity, etc.)
However, I don't think the scan is significant overhead. I expect the pool to contain tens of entries rather than thousands because it is byte-bounded and the dominant gather buffers are relatively large. Given an 80GB GPU, the hard cap is 80 GiB / 32 = 2.5 GiB. If we the average buffer size is 64MB, we will have 80 entries. This should be negligible in the C++ code.

Small correction to my previous example: a 2.5 GiB cap with an average 64 MiB buffer gives about 40 entries, not 80.

pengdurice · 2026-08-04T20:05:41Z

general question, do we have experiments results showcasing the changes' benefits?

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

pengdurice · 2026-08-05T21:21:27Z

+        pool = std::make_shared<GatherBufferPool>();
+        weak_gather_buffer_pool = pool;
+        if (gather_buffer_pool_test_budget) {
+            pool->setBudgetForTest(gather_buffer_pool_test_budget.value());


setBudgetForTest is only called when it is test purpose?

Yes. gather_buffer_pool_test_budget is unset by default and is populated only by set_z3_gather_buffer_pool_budget_for_test().

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

tohtana · 2026-08-05T22:27:04Z

general question, do we have experiments results showcasing the changes' benefits?

Yes. We have some measurements although I have not rerun the exact current head yet.

The experiment used

Qwen3-14B
One node with 8 H100s
Activation checkpointing enabled
Micro-batch 4
Sequence length 4096
10 warmup and 20 measured iterations

The results are:

With the ZeRO-3-only pass, the measured PyTorch caching-allocator retry count fell from 18–19 retries per rank to zero on every rank. Average step time improved by 8.4%, and throughput improved by 9.2%.
With the full DeepCompile pass, the retry count fell from 20 per rank to 2–4 per rank. Average step time improved by 9.2%, and throughput improved by 10.1%.
For a longer run, from 31-300 steps, no allocator retry happened.

pengdurice · 2026-08-06T00:10:57Z

general question, do we have experiments results showcasing the changes' benefits?

Yes. We have some measurements although I have not rerun the exact current head yet.

The experiment used

Qwen3-14B

One node with 8 H100s

Activation checkpointing enabled

Micro-batch 4

Sequence length 4096

10 warmup and 20 measured iterations

The results are:

With the ZeRO-3-only pass, the measured PyTorch caching-allocator retry count fell from 18–19 retries per rank to zero on every rank. Average step time improved by 8.4%, and throughput improved by 9.2%.

With the full DeepCompile pass, the retry count fell from 20 per rank to 2–4 per rank. Average step time improved by 9.2%, and throughput improved by 10.1%.
For a longer run, from 31-300 steps, no allocator retry happened.

That's great, maybe add that to the PR's description?

pengdurice · 2026-08-06T00:17:21Z

the gather buffer pool and the memory budget seem to be two separate optimization, worth considering splitting to to PRs?

tohtana · 2026-08-07T23:01:11Z

the gather buffer pool and the memory budget seem to be two separate optimization, worth considering splitting to to PRs?

I split this PR and opened #8233 and #8234. Can you review these? @pengdurice
I'm closing this now.

tohtana added 2 commits July 22, 2026 17:50

Scope DeepCompile compiler lifecycle ownership

8f8f693

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Stabilize DeepCompile ZeRO-3 memory scheduling

7b536c8

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

tohtana changed the title ~~Tohtana/deepcompile zero3 memory stability scheduler standalone~~ Stabilize DeepCompile ZeRO-3 memory scheduling Jul 23, 2026

tohtana mentioned this pull request Jul 28, 2026

Make DeepCompile ZeRO-3 scheduling memory-aware #8129

Closed

tohtana added 4 commits August 2, 2026 10:51

Merge upstream master into DeepCompile scheduler branch

e0996bf

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Start gather recovery on sustained allocator pressure

7c9a82b

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Clean up DeepCompile native state on compile failures

844b8a9

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Reduce DeepCompile tests to critical contracts

5577985

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

tohtana marked this pull request as ready for review August 3, 2026 17:56

tohtana requested review from loadams and tjruwase as code owners August 3, 2026 17:56

tohtana requested review from eternalNight and pengdurice August 3, 2026 17:57

chatgpt-codex-connector Bot reviewed Aug 3, 2026

View reviewed changes

Comment thread deepspeed/utils/allocator_telemetry.py Outdated

Comment thread deepspeed/utils/allocator_telemetry.py Outdated

tohtana added 3 commits August 3, 2026 11:23

Reduce DeepCompile product scaffolding

6596188

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Merge master and document DeepCompile diagnostics

6418bec

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Merge latest master test contract

8828019

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

pengdurice reviewed Aug 4, 2026

View reviewed changes

eternalNight reviewed Aug 4, 2026

View reviewed changes

Comment thread deepspeed/compile/inductor.py

eternalNight reviewed Aug 4, 2026

View reviewed changes

Comment thread deepspeed/compile/util.py Outdated

pengdurice reviewed Aug 4, 2026

View reviewed changes

Remove redundant DeepCompile distributed wrappers

6308b7a

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

pengdurice reviewed Aug 5, 2026

View reviewed changes

Document DeepCompile Inductor workaround

34317ef

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

This was referenced Aug 7, 2026

Split DeepCompile ZeRO-3 memory scheduler #8233

Open

Add adaptive gather-buffer pooling for DeepCompile ZeRO-3 #8234

Open

tohtana closed this Aug 7, 2026

Uh oh!

Conversation

tohtana commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Results

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

pengdurice Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pengdurice commented Aug 4, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tohtana commented Aug 5, 2026

Uh oh!

pengdurice commented Aug 6, 2026

Uh oh!

pengdurice commented Aug 6, 2026

Uh oh!

tohtana commented Aug 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tohtana commented Jul 23, 2026 •

edited

Loading

pengdurice Aug 4, 2026 •

edited

Loading