Stabilize DeepCompile ZeRO-3 memory scheduling

csrc/compile/deepcompile.cpp

@@ -4,6 +4,7 @@
 // DeepSpeed Team
 
 #include "deepcompile.h"
+#include "z3.h"
 
 #define USE_C10D_NCCL
 
@@ -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;
@@ -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)
@@ -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>(

csrc/compile/init.cpp

@@ -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",
+          &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");