Merge pull request #3155 from bghira/fix/reduce-overhead-checkpoint-compat

Handle reduce-overhead with activation checkpointing

Author: bghira

simpletuner/helpers/training/dynamo.py

@@ -11,6 +11,7 @@
 
 _PEFT_LORA_CUDAGRAPH_PATCHED = False
 _PEFT_TE_LORA_CUDAGRAPH_PATCHED = False
+_CHECKPOINT_CUDAGRAPH_ISSUE = "https://github.com/pytorch/pytorch/issues/154306"
 
 
 @torch.compiler.disable
@@ -52,19 +53,43 @@ def run_with_dynamo_config(config: Any, fn: Callable[..., Any], *args: Any, **kw
 def _normalise_text(value: Any) -> str:
     if value is None:
         return ""
+    enum_value = getattr(value, "value", None)
+    if isinstance(enum_value, str):
+        value = enum_value
     return str(value).strip().lower().replace("_", "-")
 
 
+def apply_checkpointing_cudagraph_compatibility(config: Any) -> bool:
+    backend = _normalise_text(getattr(config, "dynamo_backend", None))
+    mode = _normalise_text(getattr(config, "dynamo_mode", None))
+    checkpointing = _coerce_flag(getattr(config, "gradient_checkpointing", False))
+    if backend != "inductor" or mode not in {"cudagraphs", "reduce-overhead"} or not checkpointing:
+        return False
+
+    config.dynamo_mode = "default"
+    os.environ["TRAINING_DYNAMO_MODE"] = "default"
+    os.environ["ACCELERATE_DYNAMO_MODE"] = "default"
+    logger.warning(
+        "Activation checkpointing is incompatible with Inductor CUDA Graph mode '%s' due to PyTorch issue %s; "
+        "using dynamo_mode='default' while preserving regional compilation.",
+        mode,
+        _CHECKPOINT_CUDAGRAPH_ISSUE,
+    )
+    return True
+
+
 def _inductor_cudagraphs_enabled(config: Any = None) -> bool:
     config_backend = _normalise_text(getattr(config, "dynamo_backend", None))
+    training_backend = _normalise_text(os.environ.get("TRAINING_DYNAMO_BACKEND"))
     accelerate_backend = _normalise_text(os.environ.get("ACCELERATE_DYNAMO_BACKEND"))
-    backend = accelerate_backend or config_backend
+    backend = config_backend or training_backend or accelerate_backend
     if backend != "inductor":
         return False
 
     config_mode = _normalise_text(getattr(config, "dynamo_mode", None))
+    training_mode = _normalise_text(os.environ.get("TRAINING_DYNAMO_MODE"))
     accelerate_mode = _normalise_text(os.environ.get("ACCELERATE_DYNAMO_MODE"))
-    mode = accelerate_mode or config_mode
+    mode = config_mode or training_mode or accelerate_mode
 
     try:
         import torch._inductor.config as inductor_config

simpletuner/helpers/training/trainer.py

@@ -69,7 +69,11 @@
 from simpletuner.helpers.training.deepspeed_optimizers import DEFAULT_OPTIMIZER as DS_DEFAULT_OPTIMIZER
 from simpletuner.helpers.training.deepspeed_optimizers import sanitize_optimizer_block
 from simpletuner.helpers.training.default_settings.safety_check import safety_check
-from simpletuner.helpers.training.dynamo import install_cudagraph_workarounds, mark_cudagraph_step_begin
+from simpletuner.helpers.training.dynamo import (
+    apply_checkpointing_cudagraph_compatibility,
+    install_cudagraph_workarounds,
+    mark_cudagraph_step_begin,
+)
 from simpletuner.helpers.training.evaluation import ModelEvaluator
 from simpletuner.helpers.training.exceptions import GPUHealthError
 from simpletuner.helpers.training.gpu_circuit_breaker import get_current_gpu_index, get_gpu_circuit_breaker, is_cuda_error
@@ -893,12 +897,16 @@ def _coerce_flag(value: object) -> bool:
         except ValueError:
             raise
 
+        apply_checkpointing_cudagraph_compatibility(self.config)
+
         dynamo_backend_env = "no"
         if resolved_dynamo_backend and resolved_dynamo_backend != DynamoBackend.NO:
             dynamo_backend_env = resolved_dynamo_backend.value.lower()
         elif isinstance(dynamo_backend_value, str) and dynamo_backend_value.strip():
             dynamo_backend_env = dynamo_backend_value.strip().lower()
         os.environ["TRAINING_DYNAMO_BACKEND"] = dynamo_backend_env
+        dynamo_mode_env = str(getattr(self.config, "dynamo_mode", "") or "").strip().lower()
+        os.environ["TRAINING_DYNAMO_MODE"] = dynamo_mode_env
         self._configure_inductor_dynamic_training_passes(dynamo_backend_env)
         install_cudagraph_workarounds(self.config)
 

tests/test_dynamo.py

@@ -20,6 +20,42 @@ def _autograd_graph_contains(fn, needle: str) -> bool:
 
 
 class DynamoCudagraphWorkaroundTests(unittest.TestCase):
+    def test_activation_checkpointing_downgrades_reduce_overhead_to_default(self):
+        from simpletuner.helpers.training import dynamo
+
+        config = SimpleNamespace(
+            dynamo_backend="inductor",
+            dynamo_mode="reduce-overhead",
+            dynamo_use_regional_compilation=True,
+            gradient_checkpointing=True,
+        )
+        with (
+            unittest.mock.patch.object(dynamo.logger, "warning") as warning,
+            unittest.mock.patch.dict(os.environ, {}, clear=True),
+        ):
+            self.assertTrue(dynamo.apply_checkpointing_cudagraph_compatibility(config))
+            self.assertEqual(os.environ["TRAINING_DYNAMO_MODE"], "default")
+            self.assertEqual(os.environ["ACCELERATE_DYNAMO_MODE"], "default")
+
+        self.assertEqual(config.dynamo_mode, "default")
+        self.assertTrue(config.dynamo_use_regional_compilation)
+        warning.assert_called_once()
+        self.assertEqual(warning.call_args.args[2], "https://github.com/pytorch/pytorch/issues/154306")
+
+    def test_reduce_overhead_is_preserved_without_activation_checkpointing(self):
+        from simpletuner.helpers.training import dynamo
+
+        config = SimpleNamespace(
+            dynamo_backend="inductor",
+            dynamo_mode="reduce-overhead",
+            gradient_checkpointing=False,
+        )
+        with unittest.mock.patch.dict(os.environ, {}, clear=True):
+            self.assertFalse(dynamo.apply_checkpointing_cudagraph_compatibility(config))
+            self.assertNotIn("TRAINING_DYNAMO_MODE", os.environ)
+
+        self.assertEqual(config.dynamo_mode, "reduce-overhead")
+
     def test_peft_lora_cudagraph_patch_clones_base_result(self):
         from peft import LoraConfig
         from peft.tuners.lora.layer import Linear
@@ -123,6 +159,25 @@ def test_inductor_cudagraph_tree_mode_enabled_from_accelerate_env(self):
         ):
             self.assertTrue(dynamo._inductor_cudagraphs_enabled(SimpleNamespace(dynamo_backend=None)))
 
+    def test_configured_inductor_overrides_outer_accelerate_no_backend(self):
+        import torch._inductor.config as inductor_config
+
+        from simpletuner.helpers.training import dynamo
+
+        config = SimpleNamespace(dynamo_backend="inductor", dynamo_mode="reduce-overhead")
+        with (
+            unittest.mock.patch.dict(
+                os.environ,
+                {
+                    "ACCELERATE_DYNAMO_BACKEND": "NO",
+                    "ACCELERATE_DYNAMO_MODE": "default",
+                },
+            ),
+            unittest.mock.patch.object(inductor_config.triton, "cudagraphs", False),
+            unittest.mock.patch.object(inductor_config.triton, "cudagraph_trees", True),
+        ):
+            self.assertTrue(dynamo._inductor_cudagraphs_enabled(config))
+
 
 if __name__ == "__main__":
     unittest.main()