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

@@ -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,29 +2080,18 @@ 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 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
 
                 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)
+                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
@@ -2112,7 +2102,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     )
                     compile_fullgraph = False
 
-                    cache_info = f", cache_size_limit={cache_size_limit}" if cache_size_limit != 8 else ""
+                cache_info = ""
                 # ====================================================
                 # BLOCK COMPILE
                 # ====================================================
@@ -2156,6 +2146,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                             if hasattr(block, '_hf_hook'):
                                 continue
 
+                            compiled_refs.append((block_list, i, block))
                             block_list[i] = torch.compile(
                                 block,
                                 mode=compile_mode,
@@ -2165,6 +2156,17 @@ class BaseSDTrainProcess(BaseTrainProcess):
                             compiled_block_count += 1
 
                     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:
+                                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})."
@@ -2244,6 +2246,18 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     unet_module.to = lambda *args, **kwargs: unet_module
 
             except Exception as e:
+                # 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")
 

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", {})