diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index beb36865..f2d4a981 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/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) diff --git a/testing/test_compile_coordinate_descent.py b/testing/test_compile_coordinate_descent.py new file mode 100644 index 00000000..ba27cc15 --- /dev/null +++ b/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 diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index a84449b9..18168c73 100644 --- a/toolkit/config_modules.py +++ b/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