Fix DeepCompile profiling memory cleanup (deepspeedai#8106)

tohtana · web-flow · commit 131779cf84ad · 2026-07-01T20:05:23.000Z
This PR reduces DeepCompile ZeRO-3 memory pressure in profiling and
compiled backward by fixing three related issues:

- Let the partitioner keep using its min-cut rematerialization policy
for non-parameter activations, while still forcing ZeRO-3 parameter
aliases/casts to be recomputed.
- Compute selective-gather persistence capacity from profiled transient
headroom instead of current allocator availability.
- Avoid retaining temporary profiling outputs across warmup/measured
profiling iterations, and clear gathered ZeRO-3 parameters when
profiling exits through an exception.
- Mark incomplete memory profiles explicitly, synchronize the completion
status across ranks, and skip profile-dependent
prefetch/selective-gather decisions when profiling data is incomplete.

Signed-off-by: Masahiro Tanaka &lt;mtanaka@anyscale.com&gt;
diff --git a/deepspeed/compile/backend.py b/deepspeed/compile/backend.py
@@ -21,6 +21,7 @@
 except ImportError:
     pass
 
+import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
 
 from .fx import add_free_activations
@@ -107,7 +108,7 @@ def launch_compile_passes(global_steps: int):
         frames_partitioned.clear()
 
 
-def set_time_and_tensor_size(graph_id, graph: Graph, mem, bwd, profiling_results):
+def set_time_and_tensor_size(graph_id, graph: Graph, mem, bwd, profiling_results, mem_complete=True):
     node_time = []
     tensor_sizes = []
 
@@ -121,11 +122,22 @@ def set_time_and_tensor_size(graph_id, graph: Graph, mem, bwd, profiling_results
         profiling_results[graph_id].bwd_time = node_time
         profiling_results[graph_id].bwd_tensor_sizes = tensor_sizes
         profiling_results[graph_id].bwd_mem = mem
+        profiling_results[graph_id].bwd_mem_complete = mem_complete
     else:
         profiling_results[graph_id].fwd_graph = graph
         profiling_results[graph_id].fwd_time = node_time
         profiling_results[graph_id].fwd_tensor_sizes = tensor_sizes
         profiling_results[graph_id].fwd_mem = mem
+        profiling_results[graph_id].fwd_mem_complete = mem_complete
+
+
+def _sync_memory_profile_complete(profile_complete: bool) -> 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)
+    return bool(complete.item())
 
 
 def evaluate_symint_from_shape_env(sym_int_v):
@@ -213,9 +225,13 @@ def run_opt_passes(opt_passes: List[Callable],
 
             mem_prof = MemoryProfilingInterpreter(gm, debug_log=debug_log)
             mem_prof.run(*create_inputs_fn())
-            mem = [(name, current_alloc, delta, peak) for name, current_alloc, delta, peak in mem_prof.mem_record]
+            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:
+                mem = []
 
-            set_time_and_tensor_size(graph_id, gm.graph, mem, bwd, profiling_results)
+            set_time_and_tensor_size(graph_id, gm.graph, mem, bwd, profiling_results, profile_complete)
 
         with unset_fake_temporarily():
             get_accelerator().synchronize()
diff --git a/deepspeed/compile/partitioner.py b/deepspeed/compile/partitioner.py
@@ -75,10 +75,9 @@ def need_recompute(n: Node) -> bool:
             any([(isinstance(a, Node) and (a in ds_param_inputs or a in recomputed_nodes)) for a in node.args]):
             node.meta["recompute"] = CheckpointPolicy.MUST_RECOMPUTE
             recomputed_nodes.add(node)
-        else:
-            # If checkpointing is not enabled for this graph, assume all
-            # activations required by the backward pass should be saved.
-            node.meta.setdefault("recompute", CheckpointPolicy.MUST_SAVE)
+        # Leave non-parameter activations to the default min-cut policy. Forcing
+        # every other node to MUST_SAVE prevents safe activation rematerialization
+        # and can make long-sequence compiled backward graphs OOM.
 
 
 def get_wrapped_partitioner(
diff --git a/deepspeed/compile/passes/prefetch.py b/deepspeed/compile/passes/prefetch.py
@@ -39,8 +39,10 @@ def schedule_prefetch(gm: GraphModule, graph_id: int, graph_order: List[Tuple[in
                       create_inputs_fn, mem_budget: float, param_manager: DSGraphParamManager,
                       bwd: bool) -> GraphModule:
 
-    profile_graph = profiling_results[graph_id].bwd_graph if bwd else profiling_results[graph_id].fwd_graph
-    if is_profile_incomplete(profile_graph):
+    profile = profiling_results[graph_id]
+    profile_graph = profile.bwd_graph if bwd else profile.fwd_graph
+    mem_complete = profile.bwd_mem_complete if bwd else profile.fwd_mem_complete
+    if is_profile_incomplete(profile_graph) or not mem_complete:
         print_rank_0(f"schedule_prefetch graph_id={graph_id} incomplete profiling data; skipping prefetch")
         return gm
 
diff --git a/deepspeed/compile/passes/selective_gather.py b/deepspeed/compile/passes/selective_gather.py
@@ -54,21 +54,22 @@ def _compute_persistence_budget(all_graph_mem_records: List[List[Tuple[str, int,
             "profiled_list_count": 0,
         }
 
-    # Persistent parameters add to live allocations that remain resident past an op boundary.
+    # Persistent parameters stay live during transient allocations inside an op.
     peak_resident_alloc = max(record[1] for mem_records in non_empty_records for record in mem_records)
     transient_peak = max(record[3] for mem_records in non_empty_records for record in mem_records)
 
     return {
         "usable_mem": usable_mem,
         "peak_resident_alloc": peak_resident_alloc,
         "transient_peak": transient_peak,
-        "available_mem": max(0, usable_mem - peak_resident_alloc),
+        "available_mem": max(0, usable_mem - transient_peak),
         "profiled_list_count": len(non_empty_records),
     }
 
 
 def _profile_result_incomplete(prof) -> bool:
-    return is_profile_incomplete(prof.fwd_graph) or is_profile_incomplete(prof.bwd_graph)
+    return (is_profile_incomplete(prof.fwd_graph) or is_profile_incomplete(prof.bwd_graph) or not prof.fwd_mem_complete
+            or (prof.needs_backward and not prof.bwd_mem_complete))
 
 
 def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int, bool]], profiling_results,
@@ -171,7 +172,8 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int
     current_available_mem = vals_to_bcast[1].item()
 
     budget = _compute_persistence_budget(all_graph_mem_records, total_mem, MEM_MARGIN)
-    available_mem = int(current_available_mem * (1 - MEM_MARGIN))
+    profiled_available_mem = budget["available_mem"]
+    available_mem = profiled_available_mem
 
     ds_id_to_param = {}
     for g_id, g_pm in param_manager.items():
@@ -185,7 +187,7 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int
         f"selective_gather target_graph_id={target_graph_id} profiled_mem_lists={budget['profiled_list_count']} "
         f"total_mem={total_mem} usable_mem={budget['usable_mem']} peak_resident_alloc={budget['peak_resident_alloc']} "
         f"transient_peak={budget['transient_peak']} current_available_mem={current_available_mem} "
-        f"usable_available_mem={available_mem} "
+        f"profiled_transient_available_mem={profiled_available_mem} "
         f"persistent_count={len(persistent_ds_ids)} persistent_bytes={persistent_bytes} "
         f"candidate_count={len(ds_ids)} candidate_bytes={candidate_bytes}")
 
@@ -198,7 +200,7 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int
         return gm
 
     if available_mem == 0:
-        print_rank_0("selective_gather no currently available memory for new persistent params")
+        print_rank_0("selective_gather no profiled headroom for new persistent params")
         return gm
 
     persistent_mem = 0
diff --git a/deepspeed/compile/profilers/__init__.py b/deepspeed/compile/profilers/__init__.py
@@ -16,6 +16,8 @@ class ProfilingResult:
     needs_backward: bool = False
     fwd_mem: List[Tuple[str, int, int, int]] = field(default_factory=list)  # name, current_alloc, delta, peak
     bwd_mem: List[Tuple[str, int, int, int]] = field(default_factory=list)
+    fwd_mem_complete: bool = True
+    bwd_mem_complete: bool = True
     fwd_time: List[Tuple[str, int, int]] = field(default_factory=list)  # name, device_time, wall_time
     bwd_time: List[Tuple[str, int, int]] = field(default_factory=list)
     fwd_tensor_sizes: List[Tuple[str, int]] = field(default_factory=list)  # name, size
diff --git a/deepspeed/compile/profilers/graph_profile.py b/deepspeed/compile/profilers/graph_profile.py
@@ -91,6 +91,25 @@ def _backfill_missing_profile_metadata(graph: Graph, profile_complete: bool = Tr
             node.meta.setdefault(key, default)
 
 
+def _run_warmup_for_profile(call_fn, warmup):
+    for _ in range(warmup):
+        warmup_out = call_fn()
+        del warmup_out
+
+
+def _run_repeatedly_for_profile(call_fn, iteration, start_events, end_events):
+    out = None
+    for i in range(iteration):
+        start_events[i].record()
+        out = call_fn()
+        end_events[i].record()
+        if i + 1 < iteration:
+            del out
+            out = None
+
+    return out
+
+
 def _get_mem_usage_out_of_torch():
 
     adjust = 0
@@ -218,19 +237,18 @@ def rebuild_param_if_necessary(v):
         alloc_mem_start = get_accelerator().memory_allocated()
         max_mem_start = get_accelerator().max_memory_allocated()
 
-        if not run_only_once:
-            for i in range(self.warmup):
-                out = getattr(self, n.op)(n.target, args, kwargs)
+        def run_target():
+            return getattr(self, n.op)(n.target, args, kwargs)
+
+        warmup = 0 if run_only_once else self.warmup
+        _run_warmup_for_profile(run_target, warmup)
 
         if is_comm_op(n):
             assert self.distributed, f"Distributed environment is not initialized but comm operator {n.name} {n.target} is used."
             dist.barrier()
 
         start = time.time()
-        for i in range(iteration):
-            start_events[i].record()
-            out = getattr(self, n.op)(n.target, args, kwargs)
-            end_events[i].record()
+        out = _run_repeatedly_for_profile(run_target, iteration, start_events, end_events)
         accelerator.synchronize()
         walltime_sum = time.time() - start
 
@@ -295,13 +313,15 @@ def __init__(self, gm: GraphModule, debug_log=False):
         self.device = torch.device(get_accelerator().current_device())
         self.mem_record = []
         self.last_alloc = get_accelerator().memory_allocated()
+        self.profile_complete = True
 
         self.node_counter = 0
         self.node_num = len(gm.graph.nodes)
         self.debug_log = debug_log
 
     def run(self, *args) -> Any:
         return_val = None
+        self.profile_complete = True
         try:
             assert _all_real_if_tensor(args), "Inputs must be real tensors"
             self.nz3.enable_profiling(True)
@@ -311,9 +331,14 @@ def run(self, *args) -> Any:
                 with get_accelerator().random().fork_rng(devices=[self.device]):
                     return_val = super().run(*args)
         except Exception as e:
+            self.profile_complete = False
+            self.mem_record.clear()
             print(f"MemoryProfiling error {e}")
         finally:
-            self.nz3.enable_profiling(False)
+            try:
+                self.nz3.clear_all_gathered_params()
+            finally:
+                self.nz3.enable_profiling(False)
 
         return return_val
 
diff --git a/tests/unit/compile/test_list_schedule.py b/tests/unit/compile/test_list_schedule.py
@@ -11,6 +11,7 @@
 from torch.fx import Graph, GraphModule
 
 import deepspeed.compile.util as compile_util
+from deepspeed.compile import backend as backend_mod
 from deepspeed.compile import inductor as inductor_mod
 from deepspeed.compile import list_schedule as schedule_mod
 from deepspeed.compile.passes import prefetch as prefetch_mod
@@ -66,6 +67,31 @@ def _placeholder(graph, name):
     return _with_meta(graph.placeholder(name))
 
 
+def test_sync_memory_profile_complete_noops_without_distributed(monkeypatch):
+    monkeypatch.setattr(backend_mod.dist, "is_initialized", lambda: False)
+
+    def fail_all_reduce(*args, **kwargs):
+        raise AssertionError("all_reduce should not run without distributed init")
+
+    monkeypatch.setattr(backend_mod.dist, "all_reduce", fail_all_reduce)
+
+    assert backend_mod._sync_memory_profile_complete(True)
+    assert not backend_mod._sync_memory_profile_complete(False)
+
+
+def test_sync_memory_profile_complete_reduces_asymmetric_failure(monkeypatch):
+    monkeypatch.setattr(backend_mod.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(backend_mod, "get_accelerator", lambda: SimpleNamespace(current_device=lambda: "cpu"))
+
+    def mark_any_rank_failed(tensor, op):
+        assert op == backend_mod.dist.ReduceOp.MIN
+        tensor[0] = 0
+
+    monkeypatch.setattr(backend_mod.dist, "all_reduce", mark_any_rank_failed)
+
+    assert not backend_mod._sync_memory_profile_complete(True)
+
+
 def _allgather(graph, arg, ds_id, name, tensor_size=1, device_time=1):
     return _with_meta(
         graph.call_function(torch.ops.dc.allgather_param.default, (arg, 0, ds_id), {"dtype": torch.float16},
@@ -301,6 +327,42 @@ def available_memory(self):
     assert any("incomplete profiling data" in message for message in logs)
 
 
+def test_schedule_prefetch_skips_when_memory_profile_incomplete(monkeypatch):
+    graph = Graph()
+
+    param = _placeholder(graph, "mem_incomplete_param")
+    ag = _allgather(graph, param, 91, "mem_incomplete")
+    wait = _wait(graph, ag, 91, "mem_incomplete")
+    use = _neg(graph, wait, "mem_incomplete_use")
+    release = _release(graph, use, 91, "mem_incomplete")
+
+    graph.output((release, ))
+    graph.lint()
+
+    profiling_results = {
+        0:
+        ProfilingResult(fwd_graph=graph,
+                        bwd_graph=None,
+                        fwd_mem=[("profiled_before_abort", 0, 0, 0)],
+                        fwd_mem_complete=False)
+    }
+    gm = GraphModule(torch.nn.Module(), graph)
+    logs = []
+
+    monkeypatch.setattr(prefetch_mod, "print_rank_0", lambda message: logs.append(message))
+
+    assert prefetch_mod.schedule_prefetch(gm,
+                                          graph_id=0,
+                                          graph_order=[(0, False)],
+                                          profiling_results=profiling_results,
+                                          create_inputs_fn=lambda: (),
+                                          mem_budget=0,
+                                          param_manager={},
+                                          bwd=False) is gm
+    assert gm.graph is graph
+    assert any("incomplete profiling data" in message for message in logs)
+
+
 def test_graphsafe_rng_state_outputs_are_registered_no_reuse():
     graphsafe_run_with_rng_state = inductor_mod._get_graphsafe_run_with_rng_state()
     if graphsafe_run_with_rng_state is None:
diff --git a/tests/unit/v1/compile/test_graph_profile.py b/tests/unit/v1/compile/test_graph_profile.py
diff --git a/tests/unit/v1/compile/test_partitioner.py b/tests/unit/v1/compile/test_partitioner.py
diff --git a/tests/unit/v1/compile/test_selective_gather.py b/tests/unit/v1/compile/test_selective_gather.py
