Harden DeepCompile compiler lifecycle cleanup

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Author: tohtana

deepspeed/compile/backend.py

@@ -58,6 +58,12 @@ def set_needs_backward(self, frame_id: int, needs_backward: bool):
         if frame_id in self.frames:
             self.frames[frame_id] = (self.frames[frame_id][0], needs_backward)
 
+    def remove_graph(self, graph_id: int, frame_id: int) -> bool:
+        if frame_id in self.frames and self.frames[frame_id][0] == graph_id:
+            del self.frames[frame_id]
+            return True
+        return False
+
     def get_graph_order(self) -> List[Tuple[int, bool]]:
         assert all(isinstance(needs_backward, bool) for _, needs_backward in self.frames.values())
         return list(self.frames.values())
@@ -77,6 +83,66 @@ def clear(self):
 fwd_real_inputs = []
 
 
+def cleanup_compiled_backward_state(frame_id=None, owned_frames=None):
+    """Release process-global compiled-backward state after completion or failure."""
+    if frame_id is None:
+        if owned_frames is None:
+            frames_needing_bwd.clear()
+        else:
+            frames_needing_bwd.difference_update(owned_frames)
+            owned_frames.clear()
+    else:
+        frames_needing_bwd.discard(frame_id)
+        if owned_frames is not None:
+            owned_frames.discard(frame_id)
+    if len(frames_needing_bwd) == 0:
+        unpatch_compiled_func()
+
+
+def cleanup_backend_state(frame_id, graph_id, owned_frames=None):
+    """Release all process-global state owned by a failed graph attempt."""
+    if graph_order_with_frame_id.remove_graph(graph_id, frame_id):
+        frames_partitioned.discard(frame_id)
+    profiling_results.pop(graph_id, None)
+    param_manager.pop(graph_id, None)
+    opt_pass_times[:] = [timing for timing in opt_pass_times if timing[2] != graph_id]
+    cleanup_compiled_backward_state(frame_id, owned_frames)
+
+
+def _cleanup_compiled_backward_state_on_error(frame_id, graph_id, owned_frames=None):
+
+    def decorator(fn):
+
+        def wrapped(*args, **kwargs):
+            try:
+                return fn(*args, **kwargs)
+            except Exception:
+                cleanup_backend_state(frame_id, graph_id, owned_frames)
+                raise
+
+        return wrapped
+
+    return decorator
+
+
+def _cleanup_compiled_backward_backend_state_on_error(owned_frames=None):
+
+    def decorator(fn):
+
+        def wrapped(gm, *args, **kwargs):
+            frame_id = gm.meta["dynamo_compile_id"].frame_id
+            graph_id = id(gm.graph)
+            try:
+                return fn(gm, *args, **kwargs)
+            except Exception:
+                cleanup_backend_state(frame_id, graph_id, owned_frames)
+                raise
+
+        return wrapped
+
+    return decorator
+
+
 def register_compile_pass(name: str, opt_pass_fn, contract=None):
     from .passes.contract import register_pass_contract
     opt_passes[name] = opt_pass_fn
@@ -97,7 +163,8 @@ def init_schedule(schedule):
     remaining_schedule = deque(schedule)
 
 
-def launch_compile_passes(global_steps: int):
+def launch_compile_passes(global_steps: int, owned_frames=None):
+    """Advance the pass schedule and discard state owned by the previous compile cycle."""
     global next_pass_step, next_passes
 
     if len(remaining_schedule) > 0 and global_steps == remaining_schedule[0][0]:
@@ -109,6 +176,8 @@ def launch_compile_passes(global_steps: int):
         graph_order_with_frame_id.clear()
         profiling_results.clear()
         param_manager.clear()
+        fwd_real_inputs.clear()
+        cleanup_compiled_backward_state(owned_frames=owned_frames)
         frames_partitioned.clear()
 
 
@@ -201,6 +270,25 @@ def set_example_values_to_symints(real_inputs, param_indices=None):
     return tuple(real_inputs_ret)
 
 
+def _get_fw_real_inputs(local_real_inputs, input_storage: InputStorage, graph_id: int, debug_log: bool = False):
+    """Resolve real inputs from the one-shot queue, persistent storage, then legacy state."""
+    if local_real_inputs:
+        return local_real_inputs.popleft()
+
+    if input_storage.has_data():
+        if debug_log:
+            log_rank0(f"Retrieving real inputs from storage for graph_id={graph_id}", enable=True)
+        return input_storage.get()
+
+    if fwd_real_inputs:
+        return fwd_real_inputs.pop(0)
+
+    raise RuntimeError(f"No real inputs available for graph_id {graph_id}. "
+                       f"Local queue size: {len(local_real_inputs)}, "
+                       f"global queue size: {len(fwd_real_inputs)}, "
+                       f"storage has data: {input_storage.has_data()}")
+
+
 def run_opt_passes(opt_passes: List[Callable],
                    gm: GraphModule,
                    graph_id: int,
@@ -243,14 +331,18 @@ def run_opt_passes(opt_passes: List[Callable],
             get_accelerator().empty_cache()
 
 
-def make_backend(backend, compile_config, compile_kwargs={}):
+def make_backend(backend, compile_config, compile_kwargs={}, owned_frames=None):
 
     register_custom_ops()
 
     # Extract values from compile_config
     debug_log = compile_config.debug_log
     free_activation = compile_config.free_activation and not is_backend_inductor(backend)
 
+    if owned_frames is None:
+        owned_frames = set()
+
+    @_cleanup_compiled_backward_backend_state_on_error(owned_frames)
     def backend_fn(gm: GraphModule, real_inputs):
         graph_id = id(gm.graph)
 
@@ -275,24 +367,23 @@ def backend_fn(gm: GraphModule, real_inputs):
             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)]
 
-        global fwd_real_inputs
-
         # 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
         input_storage = InputStorage(keep_int_input_tensors=compile_config.keep_int_input_tensors,
                                      keep_all_input_tensors=compile_config.keep_all_input_tensors)
 
-        # Store in both list (for backward compatibility) and storage (for persistence)
+        # Store in a closure-local queue and storage (for persistence).
         # The input_storage keeps tensor metadata to handle cases where
         # backend_fn is called once but make_fw_graph is called multiple times
-        fwd_real_inputs.append(real_inputs)
+        local_fwd_real_inputs = deque([real_inputs])
         input_storage.put(real_inputs)
 
         global profiling_results
         if graph_id not in profiling_results:
             profiling_results[graph_id] = ProfilingResult()
             profiling_results[graph_id].param_indices = param_indices
 
+        @_cleanup_compiled_backward_state_on_error(frame_id, graph_id, owned_frames)
         def make_fw_graph(gm, sample_inputs):
             time_start = time.time()
             graph_index = len(graph_order_with_frame_id) - 1
@@ -305,19 +396,9 @@ def make_fw_graph(gm, sample_inputs):
                 if len(frames_needing_bwd) == 0:
                     patch_compiled_func()
                 frames_needing_bwd.add(frame_id)
+                owned_frames.add(frame_id)
 
-            # Try to get real_inputs from the list first, then from storage
-            if fwd_real_inputs:
-                real_inputs = fwd_real_inputs.pop(0)
-            elif input_storage.has_data():
-                # Note: input_storage is captured from the enclosing backend_fn scope
-                # Materialize tensors from storage when list is empty
-                log_rank0(f"Retrieving real inputs from storage for graph_id={graph_id}", enable=debug_log)
-                real_inputs = input_storage.get()
-            else:
-                raise RuntimeError(f"No real inputs available for graph_id {graph_id}. "
-                                   f"List size: {len(fwd_real_inputs)}, Storage has data: {input_storage.has_data()}")
-
+            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)
 
             param_manager[graph_id] = DSGraphParamManager(gm.graph, real_inputs, param_indices)
@@ -342,6 +423,7 @@ def make_fw_graph(gm, sample_inputs):
 
             return gm.graph
 
+        @_cleanup_compiled_backward_state_on_error(frame_id, graph_id, owned_frames)
         def make_bw_graph(gm, sample_inputs):
             time_start = time.time()
 
@@ -385,9 +467,7 @@ def make_bw_graph(gm, sample_inputs):
                 add_free_activations(graph_id, gm.graph,
                                      get_activation_node_names(gm.graph, param_nodes_bw, non_param_input_names))
 
-            frames_needing_bwd.remove(frame_id)
-            if len(frames_needing_bwd) == 0:
-                unpatch_compiled_func()
+            cleanup_compiled_backward_state(frame_id, owned_frames)
 
             log_rank0(
                 f"Bwd end {graph_index} graph_id={graph_id} alloc_mem={get_accelerator().memory_allocated()} graph={gm.graph}",
@@ -415,10 +495,15 @@ def compiler_fn(gm, sample_inputs):
                                             partition_fn=partition_fn)
             return torch._dynamo.optimize(**compile_kwargs)(aot_mod)
         elif backend == "inductor":
-            patch_create_aot_dispatcher_function(graph_id, z3_partition, make_fw_graph, make_bw_graph, real_inputs,
-                                                 param_indices, param_manager, frame_id, frames_partitioned)
-
-            return torch._inductor.compile(gm, real_inputs)
+            restore_aotautograd = patch_create_aot_dispatcher_function(graph_id, z3_partition, make_fw_graph,
+                                                                       make_bw_graph, real_inputs, param_indices,
+                                                                       param_manager, frame_id, frames_partitioned)
+            try:
+                return torch._inductor.compile(gm, real_inputs)
+            finally:
+                # AotAutograd.__init__ is process-global; never leak this
+                # graph-specific compiler wiring into a later compilation.
+                restore_aotautograd()
 
         raise ValueError(f"Unsupported backend {backend}")
 

deepspeed/compile/inductor.py

@@ -144,36 +144,45 @@ def _patch_deepcompile_aot_kwargs(kwargs: dict, *, graph_id: int, z3_partition:
 
 def patch_create_aot_dispatcher_function(graph_id: int, z3_partition: bool, make_fw_graph, make_bw_graph, real_inputs,
                                          param_indices, param_manager, frame_id: int, frames_partitioned: Set[int]):
+    """Temporarily install graph-specific AOT compilers and return an idempotent restore callback."""
 
     from torch._dynamo.backends.common import AotAutograd
     import functools
 
-    def patch_aotautograd():
-        # Unpatch if it was already patched
-        if hasattr(AotAutograd, "__original_init"):
-            AotAutograd.__init__ = AotAutograd.__original_init
-
-        original_init = AotAutograd.__init__
-
-        @functools.wraps(original_init)
-        def patched_init(self, **kwargs):
-            _patch_deepcompile_aot_kwargs(kwargs,
-                                          graph_id=graph_id,
-                                          z3_partition=z3_partition,
-                                          make_fw_graph=make_fw_graph,
-                                          make_bw_graph=make_bw_graph,
-                                          real_inputs=real_inputs,
-                                          param_indices=param_indices,
-                                          param_manager=param_manager,
-                                          frame_id=frame_id,
-                                          frames_partitioned=frames_partitioned)
-
-            original_init(self, **kwargs)
-
-        AotAutograd.__original_init = original_init
-        AotAutograd.__init__ = patched_init
-
-    patch_aotautograd()
+    # The constructor patch is process-global. Recover first if a previous
+    # compile failed before reaching its restoration callback.
+    if hasattr(AotAutograd, "__original_init"):
+        AotAutograd.__init__ = AotAutograd.__original_init
+        delattr(AotAutograd, "__original_init")
+
+    original_init = AotAutograd.__init__
+
+    @functools.wraps(original_init)
+    def patched_init(self, **kwargs):
+        _patch_deepcompile_aot_kwargs(kwargs,
+                                      graph_id=graph_id,
+                                      z3_partition=z3_partition,
+                                      make_fw_graph=make_fw_graph,
+                                      make_bw_graph=make_bw_graph,
+                                      real_inputs=real_inputs,
+                                      param_indices=param_indices,
+                                      param_manager=param_manager,
+                                      frame_id=frame_id,
+                                      frames_partitioned=frames_partitioned)
+
+        original_init(self, **kwargs)
+
+    AotAutograd.__original_init = original_init
+    AotAutograd.__init__ = patched_init
+
+    def restore_aotautograd():
+        """Restore only this invocation's patch without clobbering a newer owner."""
+        if AotAutograd.__init__ is patched_init:
+            AotAutograd.__init__ = original_init
+            if getattr(AotAutograd, "__original_init", None) is original_init:
+                delattr(AotAutograd, "__original_init")
+
+    return restore_aotautograd
 
 
 def register_custom_ops():

deepspeed/compile/init_z1.py

@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import copy
+from functools import partial
 
 import torch
 
@@ -185,5 +186,9 @@ def release_grad_buffer(group_idx=None):
 
     init_schedule(schedule)
 
-    engine.launch_compile_passes = launch_compile_passes
-    return make_backend(backend, compile_config, compile_kwargs=compile_kwargs)
+    engine._deepcompile_owned_frames = set()
+    engine.launch_compile_passes = partial(launch_compile_passes, owned_frames=engine._deepcompile_owned_frames)
+    return make_backend(backend,
+                        compile_config,
+                        compile_kwargs=compile_kwargs,
+                        owned_frames=engine._deepcompile_owned_frames)