Disable exhaustive tuning for quantized compile

jobs/process/BaseSDTrainProcess.py

@@ -2101,6 +2101,24 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 is_unet_quantized = getattr(self.model_config, 'quantize', False)
                 is_quantized = is_unet_quantized or getattr(self.model_config, 'quantize_te', False)
 
+                if is_quantized:
+                    # TorchAO quantization enables exhaustive coordinate-descent
+                    # tuning globally. That is unexpectedly expensive with the
+                    # default compile mode, so make it an explicit opt-in.
+                    coordinate_descent = bool(
+                        getattr(
+                            self.model_config,
+                            'compile_coordinate_descent',
+                            False,
+                        )
+                    )
+                    torch._inductor.config.coordinate_descent_tuning = (
+                        coordinate_descent
+                    )
+                    torch._inductor.config.coordinate_descent_check_all_directions = (
+                        coordinate_descent
+                    )
+
                 if not is_unet_offloaded:
                     self.sd.unet.to(self.device_torch)
 

testing/test_compile_coordinate_descent.py

@@ -0,0 +1,14 @@
+from toolkit.config_modules import ModelConfig
+
+
+def test_coordinate_descent_defaults_off():
+    config = ModelConfig(name_or_path="test")
+    assert config.compile_coordinate_descent is False
+
+
+def test_coordinate_descent_can_be_enabled_explicitly():
+    config = ModelConfig(
+        name_or_path="test",
+        compile_coordinate_descent=True,
+    )
+    assert config.compile_coordinate_descent is True

toolkit/config_modules.py

@@ -721,6 +721,9 @@ class ModelConfig:
         self.compile_mode = kwargs.get("compile_mode", "default")
         self.compile_fullgraph = kwargs.get("compile_fullgraph", False)
         self.compile_dynamic = kwargs.get("compile_dynamic", True)
+        self.compile_coordinate_descent = kwargs.get(
+            "compile_coordinate_descent", False
+        )
         self.cache_size_limit = kwargs.get("cache_size_limit", None)
         
         # kwargs to pass to the model