From 089e41dd1c45ae66cd1ba548bc365ffd44a93a62 Mon Sep 17 00:00:00 2001 From: PlagueKind <284434812+PlagueKind@users.noreply.github.com> Date: Wed, 24 Jun 2026 15:59:41 +0200 Subject: [PATCH] Compile improvements - auto cache size, fix fullgraph setting, fix triton detection (#899) * Compile improvements - auto cache size, fix fullgraph setting, fix triton detection * remove forced torchao no longer needed --- jobs/process/BaseSDTrainProcess.py | 270 +++++++++++++++-------------- toolkit/config_modules.py | 5 +- 2 files changed, 143 insertions(+), 132 deletions(-) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index bfd35365..3096c752 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -2068,6 +2068,7 @@ class BaseSDTrainProcess(BaseTrainProcess): # -> block-level compilation # ============================================================ if self.model_config.compile: + compiled_refs = [] # (block_list, index, original_block) for rollback on failure try: inner_unet_check = unwrap_model(self.sd.unet) is_unet_offloaded = hasattr(inner_unet_check, '_memory_manager') @@ -2079,145 +2080,111 @@ 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) - try: - from torch.utils._triton import has_triton - triton_available = has_triton() - except Exception: - triton_available = False + if not is_unet_offloaded: + self.sd.unet.to(self.device_torch) - if not triton_available: - print_acc("WARNING: compile is disabled.") - print_acc("Triton is not available or not working on this system.") - print_acc("Install a working 'triton' package to use compile.") - print_acc("Continuing without compilation.") - else: - - if not is_unet_offloaded: - self.sd.unet.to(self.device_torch) - - cache_size_limit = getattr(self.model_config, 'cache_size_limit', 8) + cache_size_limit = getattr(self.model_config, 'cache_size_limit', None) + user_set_cache_limit = cache_size_limit is not None + if user_set_cache_limit: torch._dynamo.config.cache_size_limit = cache_size_limit - torch._dynamo.config.suppress_errors = False + torch._dynamo.config.suppress_errors = False - compile_mode = getattr(self.model_config, 'compile_mode', 'default') - compile_dynamic = getattr(self.model_config, 'compile_dynamic', True) - compile_fullgraph = getattr(self.model_config, 'compile_fullgraph', True) - block_compile = getattr(self.model_config, 'block_compile', False) + compile_mode = getattr(self.model_config, 'compile_mode', 'default') + compile_dynamic = getattr(self.model_config, 'compile_dynamic', True) + compile_fullgraph = getattr(self.model_config, 'compile_fullgraph', False) + block_compile = getattr(self.model_config, 'block_compile', False) - # quantized + offloaded unet is incompatible with fullgraph; force it off - if is_unet_quantized and is_unet_offloaded and compile_fullgraph: - print_acc( - "Quantized offloaded Transformer detected: fullgraph=True is incompatible, " - "switching to fullgraph=False." - ) - compile_fullgraph = False + # quantized + offloaded unet is incompatible with fullgraph; force it off + if is_unet_quantized and is_unet_offloaded and compile_fullgraph: + print_acc( + "Quantized offloaded Transformer detected: fullgraph=True is incompatible, " + "switching to fullgraph=False." + ) + compile_fullgraph = False - cache_info = f", cache_size_limit={cache_size_limit}" if cache_size_limit != 8 else "" - # ==================================================== - # BLOCK COMPILE - # ==================================================== - if block_compile: - BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names() + cache_info = "" + # ==================================================== + # BLOCK COMPILE + # ==================================================== + if block_compile: + BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names() - if BLOCK_LIST_ATTRS is None or len(BLOCK_LIST_ATTRS) == 0: - BLOCK_LIST_ATTRS = [ - 'layers', - 'transformer_blocks', - 'single_transformer_blocks', - 'double_stream_blocks', - 'single_stream_blocks', - 'double_blocks', - 'single_blocks', - 'blocks', - ] - inner_unet = unwrap_model(self.sd.unet) + if BLOCK_LIST_ATTRS is None or len(BLOCK_LIST_ATTRS) == 0: + BLOCK_LIST_ATTRS = [ + 'layers', + 'transformer_blocks', + 'single_transformer_blocks', + 'double_stream_blocks', + 'single_stream_blocks', + 'double_blocks', + 'single_blocks', + 'blocks', + ] + inner_unet = unwrap_model(self.sd.unet) - compiled_block_count = 0 - - for attr_name in BLOCK_LIST_ATTRS: - # attr_name may be a dotted path for models that nest their - # blocks (e.g. hidream_o1's "model.language_model.layers"). - block_list = inner_unet - for part in attr_name.split('.'): - block_list = getattr(block_list, part, None) - if block_list is None: - break + compiled_block_count = 0 + for attr_name in BLOCK_LIST_ATTRS: + # attr_name may be a dotted path for models that nest their + # blocks (e.g. hidream_o1's "model.language_model.layers"). + block_list = inner_unet + for part in attr_name.split('.'): + block_list = getattr(block_list, part, None) if block_list is None: + break + + if block_list is None: + continue + + if not hasattr(block_list, '__len__'): + continue + + for i, block in enumerate(block_list): + if not isinstance(block, torch.nn.Module): continue - if not hasattr(block_list, '__len__'): + if hasattr(block, '_hf_hook'): continue - for i, block in enumerate(block_list): - if not isinstance(block, torch.nn.Module): - continue - - if hasattr(block, '_hf_hook'): - continue - - block_list[i] = torch.compile( - block, - mode=compile_mode, - dynamic=compile_dynamic, - fullgraph=compile_fullgraph, - ) - compiled_block_count += 1 - - if compiled_block_count > 0: - print_acc( - f"Compiled {compiled_block_count} transformer block(s) " - f"with torch.compile (mode='{compile_mode}', fullgraph={compile_fullgraph}, dynamic={compile_dynamic}{cache_info})." + compiled_refs.append((block_list, i, block)) + block_list[i] = torch.compile( + block, + mode=compile_mode, + dynamic=compile_dynamic, + fullgraph=compile_fullgraph, ) - print_acc("The first forward pass will be slow during compile. This is normal.") - print_acc("If you are experiencing issues, disable block_compile.") - else: - print_acc( - f"No individual transformer blocks found; " - f"falling back to whole-model torch.compile " - f"(mode='{compile_mode}', fullgraph={compile_fullgraph}, dynamic={compile_dynamic}{cache_info})." - ) - print_acc("The first forward pass will hang for a while. This is normal.") + compiled_block_count += 1 - if is_unet_quantized and not is_unet_offloaded and compile_fullgraph: - print_acc( - "Quantized model detected: fullgraph=True is incompatible " - "for whole-model compile, switching to fullgraph=False." - ) - compile_fullgraph = False - - if compile_mode == 'default': - self.sd.unet = torch.compile( - self.sd.unet, - dynamic=compile_dynamic, - fullgraph=compile_fullgraph, - ) + if compiled_block_count > 0: + if user_set_cache_limit: + auto_cache_limit = max(cache_size_limit, compiled_block_count * 2) + if auto_cache_limit != cache_size_limit: + torch._dynamo.config.cache_size_limit = auto_cache_limit + cache_info = f", cache_size_limit={auto_cache_limit} (auto)" else: - self.sd.unet = torch.compile( - self.sd.unet, - mode=compile_mode, - dynamic=compile_dynamic, - fullgraph=compile_fullgraph, - ) - - # ==================================================== - # WHOLE MODEL COMPILE - # ==================================================== + cache_info = f", cache_size_limit={cache_size_limit}" + else: + auto_cache_limit = compiled_block_count * 2 + torch._dynamo.config.cache_size_limit = auto_cache_limit + cache_info = f", cache_size_limit={auto_cache_limit} (auto)" + print_acc( + f"Compiled {compiled_block_count} transformer block(s) " + f"with torch.compile (mode='{compile_mode}', fullgraph={compile_fullgraph}, dynamic={compile_dynamic}{cache_info})." + ) + print_acc("The first forward pass will be slow during compile. This is normal.") + print_acc("If you are experiencing issues, disable block_compile.") else: - print_acc("Compiling model with torch.compile (whole-model compile).") + print_acc( + f"No individual transformer blocks found; " + f"falling back to whole-model torch.compile " + f"(mode='{compile_mode}', fullgraph={compile_fullgraph}, dynamic={compile_dynamic}{cache_info})." + ) print_acc("The first forward pass will hang for a while. This is normal.") - print_acc( - f"Using torch.compile settings: " - f"mode={compile_mode}, " - f"dynamic={compile_dynamic}, " - f"fullgraph={compile_fullgraph}{cache_info}" - ) - - if compile_fullgraph: + if is_unet_quantized and not is_unet_offloaded and compile_fullgraph: print_acc( - "fullgraph=True is incompatible with whole-model compile, " - "switching to fullgraph=False." + "Quantized model detected: fullgraph=True is incompatible " + "for whole-model compile, switching to fullgraph=False." ) compile_fullgraph = False @@ -2235,17 +2202,64 @@ class BaseSDTrainProcess(BaseTrainProcess): fullgraph=compile_fullgraph, ) - if not is_unet_offloaded: - # once compiled, dynamo guards hold weakrefs to the params; - # .to() on quantized params requires swap_tensors, which fails - # on tensors with weakrefs. The model stays on device anyway, - # so make .to() a no-op. - unet_module = self.sd.unet - unet_module.to = lambda *args, **kwargs: unet_module + # ==================================================== + # WHOLE MODEL COMPILE + # ==================================================== + else: + print_acc("Compiling model with torch.compile (whole-model compile).") + print_acc("The first forward pass will hang for a while. This is normal.") + + print_acc( + f"Using torch.compile settings: " + f"mode={compile_mode}, " + f"dynamic={compile_dynamic}, " + f"fullgraph={compile_fullgraph}{cache_info}" + ) + + if compile_fullgraph: + print_acc( + "fullgraph=True is incompatible with whole-model compile, " + "switching to fullgraph=False." + ) + compile_fullgraph = False + + if compile_mode == 'default': + self.sd.unet = torch.compile( + self.sd.unet, + dynamic=compile_dynamic, + fullgraph=compile_fullgraph, + ) + else: + self.sd.unet = torch.compile( + self.sd.unet, + mode=compile_mode, + dynamic=compile_dynamic, + fullgraph=compile_fullgraph, + ) + + if not is_unet_offloaded: + # once compiled, dynamo guards hold weakrefs to the params; + # .to() on quantized params requires swap_tensors, which fails + # on tensors with weakrefs. The model stays on device anyway, + # so make .to() a no-op. + unet_module = self.sd.unet + unet_module.to = lambda *args, **kwargs: unet_module except Exception as e: - print_acc(f"Failed to compile model: {e}") - print_acc("Continuing without compilation") + # undo any block-level compiles that happened before the failure, + # so "continuing without compilation" is actually true + if len(compiled_refs) > 0: + for block_list, i, original_block in compiled_refs: + block_list[i] = original_block + + if 'triton' in str(e).lower(): + print_acc("WARNING: compile is disabled.") + print_acc("Triton is not available or not working on this system.") + print_acc("Install a working 'triton' package to use compile.") + print_acc("Continuing without compilation.") + else: + print_acc(f"Failed to compile model: {e}") + print_acc("Continuing without compilation") if self.has_first_sample_requested and self.step_num <= 1 and not self.train_config.disable_sampling: print_acc("Generating first sample from first sample config") diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index c7eabebf..3b912cab 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -712,14 +712,11 @@ class ModelConfig: if self.compile and self.quantize: print("Quantized model detected - allowing torch.compile (experimental)") - # make it torchao instead of quantio for compatibility with torch compile - if self.qtype == "qfloat8": - self.qtype = "float8" self.block_compile = kwargs.get("block_compile", False) 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.cache_size_limit = kwargs.get("cache_size_limit", 8) + self.cache_size_limit = kwargs.get("cache_size_limit", None) # kwargs to pass to the model self.model_kwargs = kwargs.get("model_kwargs", {})