Scope DeepCompile compiler lifecycle ownership

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

Author: tohtana

deepspeed/compile/backend.py

@@ -77,6 +77,22 @@ def clear(self):
 fwd_real_inputs = []
 
 
+def cleanup_compiled_backward_state(frame_id=None, owned_frames=None):
+    """Release engine-owned process-global compiled-backward state."""
+    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 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 +113,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 +126,7 @@ def launch_compile_passes(global_steps: int):
         graph_order_with_frame_id.clear()
         profiling_results.clear()
         param_manager.clear()
+        cleanup_compiled_backward_state(owned_frames=owned_frames)
         frames_partitioned.clear()
 
 
@@ -201,6 +219,21 @@ 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 graph-local real inputs from the one-shot queue or persistent storage."""
+    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()
+
+    raise RuntimeError(f"No real inputs available for graph_id {graph_id}. "
+                       f"Local queue size: {len(local_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 +276,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()
+    owner_token = object()
+
     def backend_fn(gm: GraphModule, real_inputs):
         graph_id = id(gm.graph)
 
@@ -263,6 +300,7 @@ def backend_fn(gm: GraphModule, real_inputs):
         # This check cannot be placed here because autograd creates the fw/bw compiler callables before graph
         # partitioning. It is thus postponed to the point where the fw compiler is called.
         frame_id = gm.meta["dynamo_compile_id"].frame_id
+        frame_key = (owner_token, frame_id)
         graph_order_with_frame_id.add_graph(graph_id, frame_id)
 
         z3_partition = any(hasattr(v, "ds_id") for v in real_inputs)
@@ -275,17 +313,15 @@ 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
@@ -304,20 +340,10 @@ def make_fw_graph(gm, sample_inputs):
             if needs_backward:
                 if len(frames_needing_bwd) == 0:
                     patch_compiled_func()
-                frames_needing_bwd.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()}")
+                frames_needing_bwd.add(frame_key)
+                owned_frames.add(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)
 
             param_manager[graph_id] = DSGraphParamManager(gm.graph, real_inputs, param_indices)
@@ -385,9 +411,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_key, owned_frames)
 
             log_rank0(
                 f"Bwd end {graph_index} graph_id={graph_id} alloc_mem={get_accelerator().memory_allocated()} graph={gm.graph}",
@@ -415,10 +439,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. Replace the currently installed
+    # DeepCompile patch before taking ownership for this graph.
+    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)

deepspeed/compile/init_z3.py

@@ -3,6 +3,9 @@
 
 # DeepSpeed Team
 
+from functools import partial
+from threading import Lock
+
 import torch
 
 from deepspeed import comm as dist
@@ -19,6 +22,54 @@
 WARMUP = 5
 
 _MISSING = object()
+_DYNAMO_CONFIG_NAMES = ("force_parameter_static_shapes", "force_nn_module_property_static_shapes")
+_DYNAMO_CONFIG_OWNERS = {}
+_DYNAMO_CONFIG_LOCK = Lock()
+
+
+def _allow_dynamo_dynamic_parameter_shapes_for_z3(compile_kwargs):
+    """Acquire process-wide ZeRO-3 Dynamo config ownership and return its release callback."""
+    dynamo = getattr(torch, "_dynamo", None)
+    if dynamo is None:
+        try:
+            import torch._dynamo as dynamo
+        except ImportError:
+            return None
+
+    dynamo_config = getattr(dynamo, "config", None)
+    if dynamo_config is None:
+        return None
+
+    owner_token = object()
+    config_key = id(dynamo_config)
+    with _DYNAMO_CONFIG_LOCK:
+        state = _DYNAMO_CONFIG_OWNERS.get(config_key)
+        if state is None or state["config"] is not dynamo_config:
+            previous_values = {
+                config_name: getattr(dynamo_config, config_name)
+                for config_name in _DYNAMO_CONFIG_NAMES if hasattr(dynamo_config, config_name)
+            }
+            if not previous_values:
+                return None
+            state = {"config": dynamo_config, "previous_values": previous_values, "owner_tokens": set()}
+            _DYNAMO_CONFIG_OWNERS[config_key] = state
+        state["owner_tokens"].add(owner_token)
+        for config_name in state["previous_values"]:
+            setattr(dynamo_config, config_name, False)
+
+    def restore():
+        with _DYNAMO_CONFIG_LOCK:
+            state = _DYNAMO_CONFIG_OWNERS.get(config_key)
+            if state is None or state["config"] is not dynamo_config or owner_token not in state["owner_tokens"]:
+                return
+            state["owner_tokens"].remove(owner_token)
+            if state["owner_tokens"]:
+                return
+            for config_name, previous_value in state["previous_values"].items():
+                setattr(dynamo_config, config_name, previous_value)
+            del _DYNAMO_CONFIG_OWNERS[config_key]
+
+    return restore
 
 
 def _resolve_expected_grad_dtype(param):
@@ -102,9 +153,21 @@ def set_grad_buffer(_is_gradient_accumulation_boundary):
             if move_opt_states in passes or move_opt_states_sync in passes:
                 init_offload_opt_states(optimizer, dc)
 
-    engine.launch_compile_passes = launch_compile_passes
+    engine._deepcompile_owned_frames = set()
+    engine.launch_compile_passes = partial(launch_compile_passes, owned_frames=engine._deepcompile_owned_frames)
 
     patch_fake_tensor()
     torch._inductor.config.size_asserts = False
 
-    return make_backend(backend, compile_config, compile_kwargs=compile_kwargs)
+    previous_restore = getattr(engine, "_deepcompile_dynamo_config_restore", None)
+    if previous_restore is not None:
+        previous_restore()
+        del engine._deepcompile_dynamo_config_restore
+    restore_dynamo_config = _allow_dynamo_dynamic_parameter_shapes_for_z3(compile_kwargs)
+    if restore_dynamo_config is not None:
+        engine._deepcompile_dynamo_config_restore = restore_dynamo_config
+
+    return make_backend(backend,
+                        compile_config,
+                        compile_kwargs=compile_kwargs,
+                        owned_frames=engine._deepcompile_owned_frames)

deepspeed/compile/patch_compiled_func.py

@@ -82,12 +82,21 @@ class PatchedFunction(torch.autograd.Function, metaclass=FunctionMeta):
 
 
 def unpatch_compiled_func():
+    """Restore torch.autograd.Function and discard inputs captured for this compile cycle."""
     global enabled_patched_func
     enabled_patched_func = False
 
     global original_grad_fn
-    torch.autograd.Function = original_grad_fn
+    if original_grad_fn is not None:
+        torch.autograd.Function = original_grad_fn
+        original_grad_fn = None
+    clear_backward_inputs()
 
 
 def get_backward_inputs():
     return backward_inputs
+
+
+def clear_backward_inputs():
+    """Drop captured real backward inputs before the next graph compilation."""
+    backward_inputs.clear()