Skip to content

Stabilize DeepCompile ZeRO-3 memory scheduling #8169

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tohtana wants to merge 11 commits into from
Closed

Stabilize DeepCompile ZeRO-3 memory scheduling #8169

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
8f8f693
Scope DeepCompile compiler lifecycle ownership
tohtana Jul 19, 2026
7b536c8
Stabilize DeepCompile ZeRO-3 memory scheduling
tohtana Jul 23, 2026
e0996bf
Merge upstream master into DeepCompile scheduler branch
tohtana Aug 2, 2026
7c9a82b
Start gather recovery on sustained allocator pressure
tohtana Aug 2, 2026
844b8a9
Clean up DeepCompile native state on compile failures
tohtana Aug 2, 2026
5577985
Reduce DeepCompile tests to critical contracts
tohtana Aug 3, 2026
6596188
Reduce DeepCompile product scaffolding
tohtana Aug 3, 2026
6418bec
Merge master and document DeepCompile diagnostics
tohtana Aug 3, 2026
8828019
Merge latest master test contract
tohtana Aug 4, 2026
6308b7a
Remove redundant DeepCompile distributed wrappers
tohtana Aug 5, 2026
34317ef
Document DeepCompile Inductor workaround
tohtana Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions csrc/compile/deepcompile.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
// DeepSpeed Team

#include "deepcompile.h"
#include "z3.h"

#define USE_C10D_NCCL

Expand All @@ -15,7 +16,8 @@ std::shared_ptr<DoubleBufferedReduceBucket> reduce_buckets = nullptr;

c10::intrusive_ptr<c10d::ProcessGroup> process_group = nullptr;
c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem = nullptr;
ncclComm_t nccl_comm;
ncclComm_t nccl_comm = nullptr;
bool nccl_comm_initialized = false;
bool use_symm_mem;
bool profile = false;
bool pre_div_reduce = true;
Expand Down Expand Up @@ -84,10 +86,21 @@ void reset()
void cleanup()
{
reset();
reset_z3_gather_buffer_pool();
if (reduce_buckets) {
reduce_buckets->clear();
reduce_buckets.reset();
}
param_registry.reset();

ncclCommDestroy(nccl_comm);
if (nccl_comm_initialized) {
ncclCommDestroy(nccl_comm);
nccl_comm = nullptr;
nccl_comm_initialized = false;
}
process_group = nullptr;
symm_mem = nullptr;
profile = false;
}

at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id)
Expand Down Expand Up @@ -150,6 +163,7 @@ void init(c10::intrusive_ptr<c10d::ProcessGroup> pg,
// create a new nccl communicator
std::memcpy(&ncclID, tensor.to(torch::Device(torch::kCPU)).data_ptr(), NCCL_UNIQUE_ID_BYTES);
ncclCommInitRank(&nccl_comm, process_group->getSize(), ncclID, process_group->getRank());
nccl_comm_initialized = true;

param_registry = std::make_shared<DSParamRegistry>();
reduce_buckets = std::make_shared<DoubleBufferedReduceBucket>(
Expand Down
15 changes: 15 additions & 0 deletions csrc/compile/init.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,21 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
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",

@pengdurice pengdurice Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wondering if these new items are just test purpose only?

@tohtana tohtana Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@pengdurice pengdurice Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@tohtana tohtana Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@pengdurice pengdurice Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@tohtana tohtana Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

&dc::set_z3_gather_buffer_pool_budget_for_test,
"Override the adaptive ZeRO-3 gather-buffer-pool budget for tests");
m.def("update_z3_gather_buffer_pool_allocator_pressure_for_test",
&dc::update_z3_gather_buffer_pool_allocator_pressure_for_test,
"Simulate allocator pressure for ZeRO-3 gather-buffer-pool tests");
m.def("get_z3_gather_buffer_pool_state_for_test",
&dc::get_z3_gather_buffer_pool_state_for_test,
"Inspect ZeRO-3 gather-buffer-pool accounting for tests");
m.def("set_z3_param_valid_for_test",
&dc::set_z3_param_valid_for_test,
"Set a registered ZeRO-3 parameter validity bit for tests");
m.def("set_z3_prefetch_fail_after_exclusions_for_test",
&dc::set_z3_prefetch_fail_after_exclusions_for_test,
"Inject a prefetch preparation failure after admission exclusions for tests");
m.def("start_forward", &dc::start_forward, "Start forward pass");
m.def("end_forward", &dc::end_forward, "End forward pass");
m.def("start_backward", &dc::start_backward, "Start backward pass");
Expand Down
Loading
Loading