RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
This commit is contained in:
PlagueKind 2026-06-24 15:59:41 +02:00 committed by GitHub
parent d586125b40
commit 089e41dd1c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 143 additions and 132 deletions
Show Stats 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
# ============================================================
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")

5
toolkit/config_modules.py
View File

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