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

Allow nested transformer block names for quantization, lora targeting, quantizing

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-13 14:33:16 -06:00
parent 17c9279828
commit c8cd78b1a4
4 changed files with 20 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/lora_special.py
View File

@@ -370,7 +370,10 @@ class LoRASpecialNetwork(ToolkitNetworkMixin, LoRANetwork):
transformer_block_names = base_model.get_transformer_block_names()
if transformer_block_names is not None:
if not any([name in lora_name for name in transformer_block_names]):
# match against clean_name (dotted) so block names can be
# dotted paths (e.g. "model.language_model.layers"); lora_name
# has dots replaced with "$$"/"_" and wouldn't match.
if not any([block_name in clean_name for block_name in transformer_block_names]):
skip = True
else:
if self.is_pixart:

8
toolkit/util/quantize.py
View File

@@ -297,7 +297,13 @@ def quantize_model(
all_blocks: List[torch.nn.Module] = []
transformer_block_names = base_model.get_transformer_block_names()
for name in transformer_block_names:
block_list = getattr(model_to_quantize, name, None)
# name may be a dotted path for models that nest their blocks
# (e.g. hidream_o1's "model.language_model.layers").
block_list = model_to_quantize
for part in name.split('.'):
block_list = getattr(block_list, part, None)
if block_list is None:
break
if block_list is not None:
all_blocks += list(block_list)
base_model.print_and_status_update(