Stabilize DeepCompile ZeRO-3 memory scheduling

blogs/deepcompile/README.md

@@ -153,6 +153,24 @@ This project is the result of a close collaboration between Microsoft and the Un
 
 # Appendix
 
+## Diagnostics
+
+DeepCompile's low-level scheduler and gather-buffer diagnostics are opt-in environment variables rather than
+`compile` configuration fields. Prefix the normal launch command with either variable, for example:
+
+```bash
+DEEPSPEED_COMPILE_SCHEDULER_BUDGET_DEBUG=1 deepspeed <training-script>
+DEEPSPEED_ALLOCATOR_TELEMETRY=1 deepspeed <training-script>
+```
+
+Both flags are disabled when unset or set, case-insensitively, to an empty value, `0`, `false`, or `no`; any other
+value enables the diagnostic. `DEEPSPEED_COMPILE_SCHEDULER_BUDGET_DEBUG` prints rank-zero scheduler budget and
+cross-rank schedule-fingerprint lines beginning with `DeepCompile ZeRO-3 scheduler`, `DeepCompile ZeRO-3
+collective_schedule_projection`, or `DeepCompile ZeRO-3 final_schedule_fingerprint`.
+`DEEPSPEED_ALLOCATOR_TELEMETRY` prints native, line-oriented gather-buffer pool events from each process beginning
+with `DEEPSPEED_Z3_GATHER_BUFFER_POOL`. These debugging streams are not JSON or a stable machine-readable API and
+may add synchronization or logging overhead.
+
 ## Examples and Benchmarks
 
 Our DeepSpeedExamples repository provides [example code](https://github.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks/deepcompile) to enable DeepCompile.

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,15 @@ 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("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");