github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Compile improvements - auto cache size, fix fullgraph setting, fix triton detection (#899)

Browse Source 
* Compile improvements - auto cache size, fix fullgraph setting, fix triton detection

* remove forced torchao no longer needed
This commit is contained in:
PlagueKind
2026-06-24 15:59:41 +02:00
committed by GitHub
parent d586125b40
commit 089e41dd1c
2 changed files with 143 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files

270
jobs/process/BaseSDTrainProcess.py
View File

@@ -2068,6 +2068,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
# -> block-level compilation # -> block-level compilation
# ============================================================ # ============================================================
if self.model_config.compile: if self.model_config.compile:
compiled_refs = [] # (block_list, index, original_block) for rollback on failure
try: try:
inner_unet_check = unwrap_model(self.sd.unet) inner_unet_check = unwrap_model(self.sd.unet)
is_unet_offloaded = hasattr(inner_unet_check, '_memory_manager') 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_unet_quantized = getattr(self.model_config, 'quantize', False)
is_quantized = is_unet_quantized or getattr(self.model_config, 'quantize_te', False) is_quantized = is_unet_quantized or getattr(self.model_config, 'quantize_te', False)
try: if not is_unet_offloaded:
from torch.utils._triton import has_triton self.sd.unet.to(self.device_torch)
triton_available = has_triton()
except Exception:
triton_available = False
if not triton_available: cache_size_limit = getattr(self.model_config, 'cache_size_limit', None)
print_acc("WARNING: compile is disabled.") user_set_cache_limit = cache_size_limit is not None
print_acc("Triton is not available or not working on this system.") if user_set_cache_limit:
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)
torch._dynamo.config.cache_size_limit = cache_size_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_mode = getattr(self.model_config, 'compile_mode', 'default')
compile_dynamic = getattr(self.model_config, 'compile_dynamic', True) 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) block_compile = getattr(self.model_config, 'block_compile', False)
# quantized + offloaded unet is incompatible with fullgraph; force it off # quantized + offloaded unet is incompatible with fullgraph; force it off
if is_unet_quantized and is_unet_offloaded and compile_fullgraph: if is_unet_quantized and is_unet_offloaded and compile_fullgraph:
print_acc( print_acc(
"Quantized offloaded Transformer detected: fullgraph=True is incompatible, " "Quantized offloaded Transformer detected: fullgraph=True is incompatible, "
"switching to fullgraph=False." "switching to fullgraph=False."
) )
compile_fullgraph = False compile_fullgraph = False
cache_info = f", cache_size_limit={cache_size_limit}" if cache_size_limit != 8 else "" cache_info = ""
# ==================================================== # ====================================================
# BLOCK COMPILE # BLOCK COMPILE
# ==================================================== # ====================================================
if block_compile: if block_compile:
BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names() BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names()
if BLOCK_LIST_ATTRS is None or len(BLOCK_LIST_ATTRS) == 0: if BLOCK_LIST_ATTRS is None or len(BLOCK_LIST_ATTRS) == 0:
BLOCK_LIST_ATTRS = [ BLOCK_LIST_ATTRS = [
'layers', 'layers',
'transformer_blocks', 'transformer_blocks',
'single_transformer_blocks', 'single_transformer_blocks',
'double_stream_blocks', 'double_stream_blocks',
'single_stream_blocks', 'single_stream_blocks',
'double_blocks', 'double_blocks',
'single_blocks', 'single_blocks',
'blocks', 'blocks',
] ]
inner_unet = unwrap_model(self.sd.unet) inner_unet = unwrap_model(self.sd.unet)
compiled_block_count = 0 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
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: 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 continue
if not hasattr(block_list, '__len__'): if hasattr(block, '_hf_hook'):
continue continue
for i, block in enumerate(block_list): compiled_refs.append((block_list, i, block))
if not isinstance(block, torch.nn.Module): block_list[i] = torch.compile(
continue block,
mode=compile_mode,
if hasattr(block, '_hf_hook'): dynamic=compile_dynamic,
continue fullgraph=compile_fullgraph,
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})."
) )
print_acc("The first forward pass will be slow during compile. This is normal.") compiled_block_count += 1
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.")
if is_unet_quantized and not is_unet_offloaded and compile_fullgraph: if compiled_block_count > 0:
print_acc( if user_set_cache_limit:
"Quantized model detected: fullgraph=True is incompatible " auto_cache_limit = max(cache_size_limit, compiled_block_count * 2)
"for whole-model compile, switching to fullgraph=False." if auto_cache_limit != cache_size_limit:
) torch._dynamo.config.cache_size_limit = auto_cache_limit
compile_fullgraph = False cache_info = f", cache_size_limit={auto_cache_limit} (auto)"
if compile_mode == 'default':
self.sd.unet = torch.compile(
self.sd.unet,
dynamic=compile_dynamic,
fullgraph=compile_fullgraph,
)
else: else:
self.sd.unet = torch.compile( cache_info = f", cache_size_limit={cache_size_limit}"
self.sd.unet, else:
mode=compile_mode, auto_cache_limit = compiled_block_count * 2
dynamic=compile_dynamic, torch._dynamo.config.cache_size_limit = auto_cache_limit
fullgraph=compile_fullgraph, 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})."
# WHOLE MODEL COMPILE )
# ==================================================== 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: 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("The first forward pass will hang for a while. This is normal.")
print_acc( if is_unet_quantized and not is_unet_offloaded and compile_fullgraph:
f"Using torch.compile settings: "
f"mode={compile_mode}, "
f"dynamic={compile_dynamic}, "
f"fullgraph={compile_fullgraph}{cache_info}"
)
if compile_fullgraph:
print_acc( print_acc(
"fullgraph=True is incompatible with whole-model compile, " "Quantized model detected: fullgraph=True is incompatible "
"switching to fullgraph=False." "for whole-model compile, switching to fullgraph=False."
) )
compile_fullgraph = False compile_fullgraph = False
@@ -2235,17 +2202,64 @@ class BaseSDTrainProcess(BaseTrainProcess):
fullgraph=compile_fullgraph, fullgraph=compile_fullgraph,
) )
if not is_unet_offloaded: # ====================================================
# once compiled, dynamo guards hold weakrefs to the params; # WHOLE MODEL COMPILE
# .to() on quantized params requires swap_tensors, which fails # ====================================================
# on tensors with weakrefs. The model stays on device anyway, else:
# so make .to() a no-op. print_acc("Compiling model with torch.compile (whole-model compile).")
unet_module = self.sd.unet print_acc("The first forward pass will hang for a while. This is normal.")
unet_module.to = lambda *args, **kwargs: unet_module
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: except Exception as e:
print_acc(f"Failed to compile model: {e}") # undo any block-level compiles that happened before the failure,
print_acc("Continuing without compilation") # 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: 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") print_acc("Generating first sample from first sample config")

5
toolkit/config_modules.py
View File

@@ -712,14 +712,11 @@ class ModelConfig:
if self.compile and self.quantize: if self.compile and self.quantize:
print("Quantized model detected - allowing torch.compile (experimental)") 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.block_compile = kwargs.get("block_compile", False)
self.compile_mode = kwargs.get("compile_mode", "default") self.compile_mode = kwargs.get("compile_mode", "default")
self.compile_fullgraph = kwargs.get("compile_fullgraph", False) self.compile_fullgraph = kwargs.get("compile_fullgraph", False)
self.compile_dynamic = kwargs.get("compile_dynamic", True) 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 # kwargs to pass to the model
self.model_kwargs = kwargs.get("model_kwargs", {}) self.model_kwargs = kwargs.get("model_kwargs", {})