feat(offload): match balanced offload lists by component name

Entries in the always and never lists are matched against the pipeline
component name (text_encoder, vae) as well as the model class name, so
one entry covers every architecture instead of needing a new class name
per model. Class entries keep working unchanged.

Never is still tested first, so a class entry there exempts a single
model from a component entry in the always list.
This commit is contained in:
CalamitousFelicitousness
2026-08-07 21:25:36 +01:00
parent 251945939b
commit 248c5c6b7b
2 changed files with 13 additions and 6 deletions
+11 -4
View File
@@ -224,6 +224,14 @@ class OffloadHook(accelerate.hooks.ModelHook):
def model_size(self):
return sum(self.offload_map.values())
def matches(self, module, names: list, module_name: str | None = None) -> bool:
"""Match against an always/never list by class name or by pipeline component name.
Component entries such as `text_encoder` cover every architecture without listing each encoder class."""
if module.__class__.__name__ in names:
return True
module_name = module_name or getattr(module, 'module_name', None)
return module_name is not None and module_name in names
def init_hook(self, module):
return module
@@ -249,8 +257,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
for pipe in get_pipe_variants():
for module_name in get_module_names(pipe):
module_instance = getattr(pipe, module_name, None)
module_cls = module_instance.__class__.__name__
if (module_instance is not None) and (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(getattr(module_instance, "device", devices.cpu), devices.cpu)):
if (module_instance is not None) and (_id != id(module_instance)) and (not self.matches(module_instance, self.offload_never, module_name)) and (not devices.same_device(getattr(module_instance, "device", devices.cpu), devices.cpu)):
apply_balanced_offload_to_module(module_instance, op='pre')
self.last_cls = module.__class__.__name__
process_timer.add('offload', time.time() - t0)
@@ -440,9 +447,9 @@ def move_module_to_cpu(module, op='unk', force:bool=False):
op = f'{op}:force'
module = do_move(module)
used_gpu -= module_size
elif module_cls in offload_hook_instance.offload_never:
elif offload_hook_instance.matches(module, offload_hook_instance.offload_never, module_name):
op = f'{op}:never'
elif module_cls in offload_hook_instance.offload_always:
elif offload_hook_instance.matches(module, offload_hook_instance.offload_always, module_name):
op = f'{op}:always'
module = do_move(module)
used_gpu -= module_size
+2 -2
View File
@@ -943,8 +943,8 @@
{"id":"","label":"Model load model direct to GPU","localized":"","hint":"","ui":"settings_sd"},
{"id":"","label":"Model offload mode","localized":"","hint":"Controls how model components move between VRAM and system RAM to fit larger models on less VRAM.<br>- <b>none</b>: keeps everything on the GPU; fastest, but only works if the whole model fits in VRAM<br>- <b>balanced</b>: the recommended default; offloads only when VRAM use crosses a threshold, so it suits almost any GPU (tuned by the watermarks below)<br>- <b>group</b>: offloads groups of layers via diffusers group offloading; an alternative middle ground with optional stream prefetch<br>- <b>model</b>: offloads whole components such as the VAE or text encoder when idle; a more compatible fallback when balanced or group are unsupported, with smaller savings<br>- <b>sequential</b>: offloads layer by layer; the most memory saving but slowest, for when even balanced runs out of memory<br><br>Command-line flags override the automatic choice:<br>- <code>--lowvram</code>: forces <b>sequential</b><br>- <code>--medvram</code>: forces <b>balanced</b> with low watermark <b>0</b><br><br>With no flag, <b>balanced</b> is the automatic default on any GPU, with watermarks set by GPU memory (low / high):<br>- 12 GB or less: <b>0</b> / <b>0.6</b><br>- 12-24 GB: <b>0.2</b> / <b>0.6</b><br>- 24 GB or more: <b>0.2</b> / <b>0.8</b><br>(or <b>none</b> if no GPU is detected)","reload":"model","ui":"settings_offload"},
{"id":"","label":"Model types not to offload","localized":"","hint":"Model architectures to skip when offloading, separated by spaces or commas.<br>Useful for model types that misbehave when offloaded.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default is empty.","ui":"settings_offload"},
{"id":"","label":"Modules to always offload","localized":"","hint":"Module names that are always offloaded in <b>balanced</b> mode, separated by spaces, commas, or semicolons, regardless of the watermarks.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default by GPU memory: the large text encoders (<i>T5</i>, <i>UMT5</i>) are added at roughly 4-12 GB and at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Modules to never offload","localized":"","hint":"Module names that are never offloaded in <b>balanced</b> mode, separated by spaces, commas, or semicolons, keeping them resident in VRAM.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default by GPU memory: the CLIP text encoders and the VAE are kept resident at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Modules to always offload","localized":"","hint":"Modules that are always offloaded in <b>balanced</b> mode, separated by spaces, commas, or semicolons, regardless of the watermarks.<br>Entries match either a class name (<i>T5EncoderModel</i>) or a pipeline component name (<i>text_encoder</i>, <i>text_encoder_2</i>).<br>A component entry covers every model architecture without naming each encoder class.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default by GPU memory: the large text encoders (<i>T5</i>, <i>UMT5</i>) are added at roughly 4-12 GB and at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Modules to never offload","localized":"","hint":"Modules that are never offloaded in <b>balanced</b> mode, separated by spaces, commas, or semicolons, keeping them resident in VRAM.<br>Entries match either a class name (<i>CLIPTextModel</i>) or a pipeline component name (<i>vae</i>).<br>This list takes precedence, so a class entry here exempts one model from a component entry in <b><i>Modules to always offload</i></b>.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default by GPU memory: the CLIP text encoders and the VAE are kept resident at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Model types not to quantize","localized":"","hint":"Model families that quantization always skips, even when it is otherwise enabled. Space or comma separated list of model type codes; when the loaded model matches, none of its components are quantized.<br><br>Codes are the short family names shown in the load log, such as <code>sd</code>, <code>sdxl</code>, <code>sd3</code>, <code>f1</code>.<br><br>Example: <code>sd sdxl</code> leaves <i>SD</i> and <i>SDXL</i> checkpoints in full precision while other families are still quantized.<br><br>Applies to all quantization backends.<br><br>Default is empty.","reload":"model","ui":"settings_quantization"},
{"id":"","label":"Modules to not convert","localized":"","hint":"Names of modules to leave unquantized (kept in original precision), separated by spaces, commas, or semicolons.<br>Useful for layers that are sensitive to quantization, such as gate or projection layers. Example: <code>proj_out, x_embedder</code>.<br><br>Some models already exclude sensitive modules by default; entries here extend that built-in list rather than replacing it.<br><br>Default is empty.","reload":"model","ui":"settings_quantization"},
{"id":"","label":"Modules dtype dict","localized":"","hint":"Advanced: JSON mapping a quantization type to a list of module names, to quantize specific modules at a different type than the global <b><i>Quantization type</i></b>. Example: <code>{\"uint4\": [\"proj_out\"]}</code>.<br><br>Some models already assign certain modules a specific type by default; entries here merge with those built-in mappings rather than replacing them.<br><br>Default is empty.","reload":"model","ui":"settings_quantization"},