Split DeepCompile ZeRO-3 memory scheduler

blogs/deepcompile/README.md

@@ -153,6 +153,21 @@ This project is the result of a close collaboration between Microsoft and the Un
 
 # Appendix
 
+## Diagnostics
+
+DeepCompile's low-level scheduler diagnostics are controlled by an opt-in environment variable rather than a
+`compile` configuration field. Prefix the normal launch command with the variable, for example:
+
+```bash
+DEEPSPEED_COMPILE_SCHEDULER_BUDGET_DEBUG=1 deepspeed <training-script>
+```
+
+The flag is disabled when unset or set, case-insensitively, to an empty value, `0`, `false`, or `no`; any other
+value enables the diagnostic. It 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`.
+This debugging stream is 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

@@ -15,7 +15,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 +85,20 @@ void reset()
 void cleanup()
 {
     reset();
+    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 +161,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/z3.cpp

@@ -191,6 +191,7 @@ class Z3CustomOpExecutor : public CustomOpExecutor {
                 }
             }
             auto target_dtype = dtype ? dtype.value() : ds_tensor.scalar_type();
+            at::cuda::CUDAStreamGuard guard(ag_stream_);
             output_bufs[ds_id] =
                 torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype));
         }

deepspeed/compile/backend.py

@@ -27,12 +27,12 @@
 from .fx import add_free_activations
 from .graph_param import DSGraphParamManager
 from .profilers import ProfilingResult
-from .profilers.graph_profile import MemoryProfilingInterpreter
+from .profilers.graph_profile import MemoryProfilingInterpreter, is_profile_incomplete
 from .patch_compiled_func import (clear_backward_inputs, patch_compiled_func, pop_backward_input,
                                   register_backward_frame, unpatch_compiled_func)
 from .util import get_input_nodes, get_activation_node_names, get_index_by_graph_id, get_deepcompile_handle, log_rank0, is_backend_inductor
 from .partitioner import get_wrapped_partitioner
-from .inductor import register_custom_ops, patch_create_aot_dispatcher_function
+from .inductor import register_custom_ops, patch_create_aot_dispatcher_function, deepcompile_z3_inductor_config_patch
 from .input_storage import InputStorage
 
 remaining_schedule = None
@@ -158,12 +158,15 @@ def set_time_and_tensor_size(graph_id, graph: Graph, mem, bwd, profiling_results
         profiling_results[graph_id].fwd_mem_complete = mem_complete
 
 
-def _sync_memory_profile_complete(profile_complete: bool) -> bool:
+def _sync_memory_profile_complete(profile_complete: bool, process_group=None) -> bool:
     if not dist.is_initialized():
         return profile_complete
 
     complete = torch.tensor([1 if profile_complete else 0], device=torch.device(get_accelerator().current_device()))
-    dist.all_reduce(complete, dist.ReduceOp.MIN)
+    if process_group is None:
+        dist.all_reduce(complete, dist.ReduceOp.MIN)
+    else:
+        dist.all_reduce(complete, dist.ReduceOp.MIN, group=process_group)
     return bool(complete.item())
 
 
@@ -174,16 +177,28 @@ def evaluate_symint_from_shape_env(sym_int_v):
     return sym_int_v.node.hint
 
 
-def set_example_values_to_symints(real_inputs, param_indices=None):
+_ZERO_PARAMETER_COMPILE_METADATA = ("ds_id", "ds_shape", "ds_persist", "ds_status", "ds_target_dtype")
+
+
+def set_example_values_to_symints(real_inputs, param_indices=None, real_zero_params=None):
     real_inputs_ret = []
 
     # Create a set of parameter indices for quick lookup
     param_idx_set = set()
     if param_indices is not None:
         param_idx_set = {i for i, _, _ in param_indices}
+    param_ds_ids = {i: ds_id for i, ds_id, _ in (param_indices or [])}
+    real_zero_params = real_zero_params or {}
 
     for i, v in enumerate(real_inputs):
         if isinstance(v, torch.Tensor):
+            real_zero_param = real_zero_params.get(param_ds_ids.get(i))
+            if i in param_idx_set and real_zero_param is not None:
+                # Stored and fake profiling inputs both discard instance-bound
+                # ZeRO methods, so recover the original before materialization.
+                real_inputs_ret.append(real_zero_param)
+                continue
+
             if is_fake(v):
                 shape = []
                 for fs in v.shape:
@@ -207,10 +222,12 @@ def set_example_values_to_symints(real_inputs, param_indices=None):
 
                     # Create Parameter if this input index corresponds to a parameter
                     if i in param_idx_set:
-                        dummy_v = torch.nn.Parameter(dummy_v)
-                        # Copy any additional attributes from the original if they exist
-                        if hasattr(v, 'ds_id'):
-                            dummy_v.ds_id = v.ds_id
+                        dummy_v = torch.nn.Parameter(dummy_v, requires_grad=v.requires_grad)
+                        # Profiling and graph-parameter consumers use these ZeRO
+                        # attributes after symbolic fake inputs are materialized.
+                        for attr in _ZERO_PARAMETER_COMPILE_METADATA:
+                            if hasattr(v, attr):
+                                setattr(dummy_v, attr, getattr(v, attr))
 
                     real_inputs_ret.append(dummy_v)
             else:
@@ -248,7 +265,9 @@ def run_opt_passes(opt_passes: List[Callable],
                    mem_budget: float,
                    param_manager,
                    bwd: bool,
-                   debug_log=False) -> None:
+                   debug_log=False,
+                   process_group=None) -> None:
+    """Apply scheduled graph passes and retain only complete post-pass memory profiles."""
 
     with unset_fake_temporarily():
         get_accelerator().synchronize()
@@ -265,13 +284,23 @@ def run_opt_passes(opt_passes: List[Callable],
             gm.graph.lint()
             gm.recompile()
 
-            mem_prof = MemoryProfilingInterpreter(gm, debug_log=debug_log)
-            mem_prof.run(*create_inputs_fn())
-            profile_complete = _sync_memory_profile_complete(mem_prof.profile_complete)
-            if profile_complete:
-                mem = [(name, current_alloc, delta, peak) for name, current_alloc, delta, peak in mem_prof.mem_record]
-            else:
+            # Re-profiling an already incomplete graph would turn synthetic
+            # backfilled metadata into a seemingly valid memory profile.
+            operator_profile_complete = _sync_memory_profile_complete(not is_profile_incomplete(gm.graph),
+                                                                      process_group)
+            if not operator_profile_complete:
+                profile_complete = False
                 mem = []
+            else:
+                mem_prof = MemoryProfilingInterpreter(gm, debug_log=debug_log, process_group=process_group)
+                mem_prof.run(*create_inputs_fn())
+                profile_complete = _sync_memory_profile_complete(mem_prof.profile_complete, process_group)
+                if profile_complete:
+                    mem = [(name, current_alloc, delta, peak)
+                           for name, current_alloc, delta, peak in mem_prof.mem_record]
+                else:
+                    mem = []
+                del mem_prof
 
             set_time_and_tensor_size(graph_id, gm.graph, mem, bwd, profiling_results, profile_complete)
 
@@ -281,7 +310,7 @@ def run_opt_passes(opt_passes: List[Callable],
             get_accelerator().empty_cache()
 
 
-def make_backend(backend, compile_config, compile_kwargs={}, owned_frames=None):
+def make_backend(backend, compile_config, compile_kwargs={}, process_group=None, owned_frames=None):
 
     register_custom_ops()
 
@@ -312,11 +341,17 @@ def backend_fn(gm: GraphModule, real_inputs):
         if z3_partition:
             param_indices = [(i, input_val.ds_id, input_val.ds_shape) for i, input_val in enumerate(real_inputs)
                              if isinstance(input_val, torch.nn.Parameter)]
+            real_zero_params = {
+                input_val.ds_id: input_val
+                for input_val in real_inputs if isinstance(input_val, torch.nn.Parameter)
+                and hasattr(input_val, "all_gather") and hasattr(input_val, "partition")
+            }
         else:
             assert all(hasattr(v, "param_id") for v in real_inputs
                        if isinstance(v, torch.nn.Parameter)), "All param inputs should have param_id"
             param_indices = [(i, input_val.param_id, input_val.shape) for i, input_val in enumerate(real_inputs)
                              if isinstance(input_val, torch.nn.Parameter)]
+            real_zero_params = {}
 
         # Create an InputStorage instance for this specific graph
         # It will be captured by the make_fw_graph closure, eliminating the need for graph ID management
@@ -331,7 +366,7 @@ def backend_fn(gm: GraphModule, real_inputs):
 
         global profiling_results
         if graph_id not in profiling_results:
-            profiling_results[graph_id] = ProfilingResult()
+            profiling_results[graph_id] = ProfilingResult(process_group=process_group)
             profiling_results[graph_id].param_indices = param_indices
 
         def make_fw_graph(gm, sample_inputs):
@@ -350,7 +385,7 @@ def make_fw_graph(gm, sample_inputs):
                 register_backward_frame(frame_key)
 
             real_inputs = _get_fw_real_inputs(local_fwd_real_inputs, input_storage, graph_id, debug_log=debug_log)
-            real_inputs = set_example_values_to_symints(real_inputs)
+            real_inputs = set_example_values_to_symints(real_inputs, param_indices, real_zero_params=real_zero_params)
 
             param_manager[graph_id] = DSGraphParamManager(gm.graph, real_inputs, param_indices)
 
@@ -365,7 +400,8 @@ def make_fw_graph(gm, sample_inputs):
                 mem_budget=.0,  # unused
                 param_manager=param_manager,
                 bwd=False,
-                debug_log=debug_log)
+                debug_log=debug_log,
+                process_group=process_group)
 
             opt_pass_times.append(("fwd", graph_index, graph_id, time.time() - time_start))
 
@@ -404,7 +440,8 @@ def make_bw_graph(gm, sample_inputs):
                 mem_budget=.0,  # unused
                 param_manager=param_manager,
                 bwd=True,
-                debug_log=debug_log)
+                debug_log=debug_log,
+                process_group=process_group)
 
             # assert graph_id in param_manager, f"Graph {graph_id} not found in param_manager"
 
@@ -446,7 +483,8 @@ def compiler_fn(gm, sample_inputs):
                                                                        make_bw_graph, real_inputs, param_indices,
                                                                        param_manager, frame_id, frames_partitioned)
             try:
-                return torch._inductor.compile(gm, real_inputs)
+                with deepcompile_z3_inductor_config_patch(z3_partition):
+                    return torch._inductor.compile(gm, real_inputs)
             finally:
                 # AotAutograd.__init__ is process-global; never leak this
                 # graph-specific compiler wiring into a later compilation.

deepspeed/compile/inductor.py

@@ -3,6 +3,7 @@
 
 # DeepSpeed Team
 
+from contextlib import nullcontext
 from typing import Set
 
 import torch
@@ -21,6 +22,43 @@
 from .graph_param import DSGraphParamManager
 from .partitioner import get_wrapped_partitioner
 
+# With PyTorch 2.10, Inductor was observed to generate a mix-order persistent
+# reduction for a DeepCompile ZeRO-3 backward graph, and Triton rejected the
+# kernel for exceeding the hardware's per-kernel resource limit. Setting
+# persistent_reductions=False alone is insufficient because mix-order codegen
+# passes override_persistent_reduction=True. Revisit this ZeRO-3 compile
+# workaround when PyTorch's Inductor reduction heuristics change.
+_DEEP_COMPILE_Z3_INDUCTOR_REDUCTION_CONFIG = {
+    "triton.mix_order_reduction": False,
+    "triton.persistent_reductions": False,
+}
+
+
+def deepcompile_z3_inductor_config_patch(enabled: bool):
+    """Disable reduction heuristics that create oversized kernels for DeepCompile ZeRO-3 graphs."""
+    if not enabled:
+        return nullcontext()
+
+    inductor = getattr(torch, "_inductor", None)
+    config = getattr(inductor, "config", None)
+    if config is None or not hasattr(config, "patch"):
+        return nullcontext()
+
+    triton_config = getattr(config, "triton", None)
+    if triton_config is None:
+        return nullcontext()
+
+    overrides = {
+        config_name: value
+        for config_name, value in _DEEP_COMPILE_Z3_INDUCTOR_REDUCTION_CONFIG.items()
+        if hasattr(triton_config,
+                   config_name.split(".", 1)[1])
+    }
+    if not overrides:
+        return nullcontext()
+
+    return config.patch(overrides)
+
 
 def _get_graphsafe_run_with_rng_state():
     try:
@@ -46,6 +84,7 @@ def _mark_output_never_reuse(out, *, enabled):
 
 
 def patch_compiler(original_compiler, dc_compiler, z3_partition: bool, graph_id, graph_param_manager, bwd: bool):
+    """Wrap an AOT compiler with DeepCompile rewrites and ZeRO-3 fake-shape repair."""
 
     def wrapped_compiler(gm, fake_inputs):
         mod_graph = dc_compiler(gm, fake_inputs)
@@ -80,7 +119,8 @@ def wrapped_compiler(gm, fake_inputs):
         else:
             patched_inputs = fake_inputs
 
-        return original_compiler(gm, patched_inputs)
+        with deepcompile_z3_inductor_config_patch(z3_partition):
+            return original_compiler(gm, patched_inputs)
 
     return wrapped_compiler
 

deepspeed/compile/init_z1.py

@@ -11,7 +11,7 @@
 from deepspeed.accelerator import get_accelerator
 from .passes import zero1_compile, zero3_compile
 from .backend import make_backend, launch_compile_passes, init_schedule
-from .util import get_deepcompile_handle, add_pre_backward_hook
+from .util import add_pre_backward_hook
 
 WARMUP = 5
 
@@ -96,8 +96,7 @@ def init_z1(engine, backend, compile_config, compile_kwargs, schedule=None, use_
         hook.remove()
     optimizer._grad_acc_hooks.clear()
 
-    dc = get_deepcompile_handle()
-    dc.init(engine.data_parallel_group, compile_config, engine.zero_reduce_bucket_size())
+    dc = engine._initialize_deepcompile_native(compile_config)
 
     if use_z2:
         grad_buffer = {}

deepspeed/compile/init_z3.py

@@ -16,7 +16,7 @@
 from .passes import zero3_compile, prefetch, selective_gather, offload_parameters
 from .backend import make_backend, launch_compile_passes, init_schedule
 from .patch_fake_tensor import patch_fake_tensor
-from .util import get_deepcompile_handle, add_pre_backward_hook, add_post_backward_hook
+from .util import add_pre_backward_hook, add_post_backward_hook
 from .z3_eager_fallback import DeepCompileZ3EagerFallback
 
 WARMUP = 5
@@ -104,8 +104,7 @@ def init_z3(engine, backend, compile_config, compile_kwargs, schedule=None):
         optimizer.ipg_buckets.clear()
         get_accelerator().empty_cache()
 
-    dc = get_deepcompile_handle()
-    dc.init(engine.data_parallel_group, compile_config, engine.zero_reduce_bucket_size())
+    dc = engine._initialize_deepcompile_native(compile_config)
 
     engine._deepcompile_z3_eager_fallback = DeepCompileZ3EagerFallback(engine)
     add_post_backward_hook(engine._deepcompile_z3_eager_fallback.complete_backward)
@@ -194,4 +193,5 @@ def set_grad_buffer(_is_gradient_accumulation_boundary):
     return make_backend(backend,
                         compile_config,
                         compile_kwargs=compile_kwargs,
+                        process_group=engine.data_parallel_group,
                         owned_frames=engine._deepcompile_owned_frames)