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

Add a triton check on compile

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-12 12:26:16 -06:00
parent 10e50d5797
commit 43b22b91ee
1 changed files with 12 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
jobs/process/BaseSDTrainProcess.py
View File

@@ -2079,7 +2079,18 @@ class BaseSDTrainProcess(BaseTrainProcess):
is_quantized = getattr(self.model_config, 'quantize', False) or \
getattr(self.model_config, 'quantize_te', False)
if is_quantized and is_unet_offloaded:
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.")
elif is_quantized and is_unet_offloaded:
print_acc("WARNING: compile is disabled.")
print_acc("Quantized models with Transformer offloading are incompatible.")
print_acc("Disable Transformer offload to use compile.")