Fix DeepCompile for PyTorch 2.8/2.9 compatibility (#7755)

PyTorch 2.8 added a new static_lifetime_input_indices parameter to the
partition function. This breaks DeepCompile with ZeRO stage 3.
This PR updates `partition_recompute_ds_params` to accept `**kwargs` and
forward them to the underlying partition function, maintaining backward
compatibility with PyTorch 2.6/2.7.

DeepCompile works with PyTorch v2.9 when using ZeRO Stage 1 or 2.
However, ZeRO Stage 3 is not currently supported on PyTorch v2.9 (it
still works on PyTorch <= v2.8). DeepCompile tests are skipped when
PyTorch version is v2.9 and ZeRO stage is 3.

---------

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

Author: tohtana

deepspeed/compile/inductor.py

@@ -7,13 +7,10 @@
 
 try:
     import torch.utils._pytree as pytree
-    from torch._functorch.aot_autograd import create_aot_dispatcher_function
     from torch._inductor.lowering import register_lowering, fallbacks, add_needs_realized_inputs
     from torch._inductor.ir import TensorBox, FallbackKernel, Layout, IRNode
     from torch._inductor.virtualized import V
     from torch._inductor.scheduler import Scheduler
-
-    original_create_aot_dispatcher_function = create_aot_dispatcher_function
 except ImportError:
     pass
 

deepspeed/compile/partitioner.py

@@ -87,10 +87,10 @@ def get_wrapped_partitioner(
     partition_fn,
 ):
 
-    def partition_recompute_ds_params(joint_module: GraphModule, _joint_inputs, *,
-                                      num_fwd_outputs) -> Tuple[GraphModule, GraphModule]:
+    def partition_recompute_ds_params(joint_module: GraphModule, _joint_inputs, *, num_fwd_outputs,
+                                      **kwargs) -> Tuple[GraphModule, GraphModule]:
         if z3_partition:
             _recompute_param_aliases(joint_module.graph, param_indices)
-        return partition_fn(joint_module, _joint_inputs, num_fwd_outputs=num_fwd_outputs)
+        return partition_fn(joint_module, _joint_inputs, num_fwd_outputs=num_fwd_outputs, **kwargs)
 
     return partition_recompute_ds_params

deepspeed/compile/util.py

@@ -25,7 +25,7 @@
 
 
 def is_deepcompile_supported() -> bool:
-    return required_torch_version(min_version=2.6, max_version=2.7) and get_accelerator().device_name() == "cuda"
+    return required_torch_version(min_version=2.6, max_version=2.9) and get_accelerator().device_name() == "cuda"
 
 
 dc_handle = None

deepspeed/runtime/engine.py

@@ -77,6 +77,7 @@
 from deepspeed.runtime import lr_schedules
 from deepspeed.utils import groups
 from deepspeed.utils import logger, log_dist, log_dist_once, instrument_w_nvtx
+from deepspeed.utils.torch import required_torch_version
 from deepspeed.utils.z3_leaf_module import apply_zero_leaf_module_config
 from deepspeed.utils.timer import NoopTimer, ThroughputTimer, SynchronizedWallClockTimer, \
     FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, \
@@ -4296,6 +4297,10 @@ def passes_name_to_fn(passes):
             elif self.zero_optimization_stage() == ZeroStageEnum.gradients:
                 backend = init_z1(self, backend, compile_config, compile_kwargs, schedule, use_z2=True)
             elif self.zero_optimization_stage() == ZeroStageEnum.weights:
+                if required_torch_version(min_version=2.9):
+                    raise RuntimeError(
+                        "DeepCompile with ZeRO stage 3 is not currently supported on PyTorch >= 2.9. "
+                        "Please use ZeRO stage 1 or 2 with DeepCompile, or disable DeepCompile for ZeRO stage 3.")
                 backend = init_z3(self, backend, compile_config, compile_kwargs, schedule)
 
         # Hook state must align with whether DeepCompile is active.

tests/unit/v1/compile/test_compile_zero.py

@@ -82,6 +82,8 @@ class TestDeepCompile(DistributedTest):
     def test(self, zero_stage, dtype, deepcompile):
         if not required_torch_version(min_version=2.6):
             pytest.skip("DeepCompile requires PyTorch >= v2.6")
+        if zero_stage == 3 and required_torch_version(min_version=2.9):
+            pytest.skip("DeepCompile with ZeRO stage 3 is not currently supported on PyTorch >= 2.9")
 
         if dtype == torch.bfloat16:
             skip_on_arch(min_arch=8)
@@ -123,6 +125,8 @@ def test_padded_shard_handling(self, zero_stage, dtype):
         """Test that parameters with padding (uneven division) work correctly with DeepCompile"""
         if not required_torch_version(min_version=2.6):
             pytest.skip("DeepCompile requires PyTorch >= v2.6")
+        if required_torch_version(min_version=2.9):
+            pytest.skip("DeepCompile with ZeRO stage 3 is not supported on PyTorch >= 2.9")
 
         if get_accelerator().device_name() == "cpu":
             pytest.skip("CPU does not support this test yet")
@@ -156,6 +160,8 @@ def test_free_activation_mode(self, zero_stage, dtype):
         """Test that eagerly free activations work correctly and the threshold is configurable"""
         if not required_torch_version(min_version=2.6):
             pytest.skip("DeepCompile requires PyTorch >= v2.6")
+        if zero_stage == 3 and required_torch_version(min_version=2.9):
+            pytest.skip("DeepCompile with ZeRO stage 3 is not supported on PyTorch >= 2.9")
 
         if get_accelerator().device_name() == "cpu":
             pytest.skip("CPU does not support this test yet")
@@ -187,6 +193,8 @@ def test_fusing_allgather_and_autocast(self, zero_stage, dtype):
         """Test that allgather and autocast can be correctly fused with DeepCompile"""
         if not required_torch_version(min_version=2.6):
             pytest.skip("DeepCompile requires PyTorch >= v2.6")
+        if required_torch_version(min_version=2.9):
+            pytest.skip("DeepCompile with ZeRO stage 3 is not supported on PyTorch >= 2.9")
 
         if get_accelerator().device_name() == "cpu":
             pytest.skip("CPU does not support this test yet")