mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
feat(offload): per-component group offload engine
Group offload is applied per component through one engine shared by regular and modular pipelines. Each component carries a config signature: re-application with unchanged settings is a no-op instead of raising before the first forward or silently keeping a stale config, and changed settings remove the hooks and reapply. Switching offload modes cleans up the previous mode's hooks in both directions. - text encoders always offload at leaf level without streams, so their weights are never held in pinned host memory - the vae never takes group hooks and stays resident: the hooks are forward-scoped, while pipelines enter through encode/decode and tiled calls re-enter per tile - new pin offload memory option: disabled pins one group at a time instead of holding the whole module in non-pageable memory, and modules larger than half of system memory degrade automatically - record stream is clamped to stream mode; the standalone combination is rejected upstream
This commit is contained in:
CalamitousFelicitousness
+131
-16
@@ -76,30 +76,142 @@ def set_accelerate(sd_model):
|
||||
set_accelerate_to_module(sd_model.decoder_pipe)
|
||||
|
||||
|
||||
def apply_group_offload(sd_model, op:str='model'):
|
||||
offload_dct = {
|
||||
'onload_device': devices.device,
|
||||
'offload_device': devices.cpu,
|
||||
'offload_type': shared.opts.group_offload_type,
|
||||
'num_blocks_per_group': shared.opts.group_offload_blocks,
|
||||
def group_offload_config(main: bool) -> dict:
|
||||
"""Effective group offload settings for one component. Components that run once per
|
||||
generation take the leaf no-stream policy regardless of the main settings, so their
|
||||
weights are never held in pinned host memory."""
|
||||
stream = shared.opts.group_offload_stream if main else False
|
||||
blocks = max(1, int(shared.opts.group_offload_blocks))
|
||||
if stream and blocks != 1:
|
||||
blocks = 1 # streamed prefetch supports one block per group; upstream clamps with a warning otherwise
|
||||
return {
|
||||
'offload_type': shared.opts.group_offload_type if main else 'leaf_level',
|
||||
'num_blocks_per_group': blocks,
|
||||
'non_blocking': shared.opts.diffusers_offload_nonblocking,
|
||||
'use_stream': shared.opts.group_offload_stream,
|
||||
'record_stream': shared.opts.group_offload_record,
|
||||
'low_cpu_mem_usage': False,
|
||||
'use_stream': stream,
|
||||
'record_stream': shared.opts.group_offload_record and stream, # record without streams is rejected upstream
|
||||
'low_cpu_mem_usage': stream and not shared.opts.group_offload_pin,
|
||||
}
|
||||
if shared.opts.group_offload_type == 'block_level':
|
||||
offload_dct['exclude_modules'] = ['vae']
|
||||
log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} options={offload_dct}')
|
||||
if hasattr(sd_model, "enable_group_offload"):
|
||||
sd_model.enable_group_offload(**offload_dct)
|
||||
else:
|
||||
log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported')
|
||||
|
||||
|
||||
def remove_group_offload_component(module):
|
||||
if getattr(module, 'sdnext_group_offload_sig', None) is None:
|
||||
return
|
||||
from diffusers.hooks.group_offloading import _GROUP_OFFLOADING, _LAYER_EXECUTION_TRACKER, _LAZY_PREFETCH_GROUP_OFFLOADING
|
||||
from diffusers.hooks.hooks import HookRegistry
|
||||
registry = HookRegistry.check_if_exists_or_initialize(module)
|
||||
registry.remove_hook(_GROUP_OFFLOADING, recurse=True)
|
||||
registry.remove_hook(_LAYER_EXECUTION_TRACKER, recurse=True)
|
||||
registry.remove_hook(_LAZY_PREFETCH_GROUP_OFFLOADING, recurse=True)
|
||||
module.sdnext_group_offload_sig = None
|
||||
|
||||
|
||||
def remove_group_offload(sd_model):
|
||||
removed = []
|
||||
for module_name in get_module_names(sd_model):
|
||||
module = getattr(sd_model, module_name, None)
|
||||
if isinstance(module, torch.nn.Module) and getattr(module, 'sdnext_group_offload_sig', None) is not None:
|
||||
remove_group_offload_component(module)
|
||||
removed.append(module_name)
|
||||
if removed:
|
||||
log.debug(f'Offload: type=group op=remove modules={removed}')
|
||||
|
||||
|
||||
def apply_group_offload_component(module, module_name: str, main: bool, op: str = 'model') -> bool:
|
||||
"""Apply group offload to one component. Re-application with unchanged settings is a no-op:
|
||||
the hooks silently keep their original config when re-applied and raise before the first
|
||||
forward, so a changed config must remove the old hooks first."""
|
||||
from diffusers.hooks import apply_group_offloading
|
||||
cfg = group_offload_config(main)
|
||||
if cfg['use_stream'] and not cfg['low_cpu_mem_usage']:
|
||||
size_gb, _params = get_module_size(module)
|
||||
limit_gb = 0.5 * shared.cpu_memory # heuristic ceiling: pinned memory is non-pageable, and half of system memory must stay available to everything else
|
||||
if size_gb > limit_gb:
|
||||
cfg['low_cpu_mem_usage'] = True
|
||||
log.warning(f'Setting {op}: offload=group module={module_name} size={size_gb:.3f} limit={limit_gb:.3f} pin=dynamic memory guard')
|
||||
sig = f'{devices.device}:{main}:' + ':'.join(str(v) for v in cfg.values())
|
||||
if getattr(module, 'sdnext_group_offload_sig', None) == sig:
|
||||
return False
|
||||
if hasattr(module, '_hf_hook'): # leftover accelerate hooks from a previous offload mode abort the group apply upstream
|
||||
module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
|
||||
remove_group_offload_component(module)
|
||||
module.requires_grad_(False)
|
||||
apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg)
|
||||
module.sdnext_group_offload_sig = sig
|
||||
return True
|
||||
|
||||
|
||||
def set_group_resident(module):
|
||||
"""VAE-class components never take group hooks: the hooks are forward-scoped, while
|
||||
pipelines enter through encode/decode, and tiled calls re-enter per tile."""
|
||||
if hasattr(module, '_hf_hook'):
|
||||
module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
|
||||
remove_group_offload_component(module)
|
||||
module.requires_grad_(False)
|
||||
module.to(devices.device)
|
||||
|
||||
|
||||
def group_offload_role(module_name: str, module) -> str:
|
||||
cls = module.__class__.__name__
|
||||
if 'vae' in module_name.lower() or cls.startswith(('Autoencoder', 'VQModel', 'AsymmetricAutoencoder', 'ConsistencyDecoder')):
|
||||
return 'resident'
|
||||
if module_name.startswith(('text_encoder', 'image_encoder', 'safety_checker')):
|
||||
return 'aux'
|
||||
return 'main'
|
||||
|
||||
|
||||
def apply_modular_group_offload(sd_model, op:str='model'):
|
||||
"""Per-component group offload for modular pipelines, which lack the pipeline-level
|
||||
enable_*_offload entry points. The model and sequential modes also route here."""
|
||||
if shared.opts.diffusers_offload_mode != 'group' and not getattr(sd_model, 'sdnext_modular_offload_warned', False):
|
||||
sd_model.sdnext_modular_offload_warned = True
|
||||
log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported on modular pipelines: using group offload')
|
||||
applied = []
|
||||
for name in ('transformer', 'transformer_ref'):
|
||||
transformer = getattr(sd_model, name, None)
|
||||
if transformer is not None and apply_group_offload_component(transformer, name, main=True, op=op):
|
||||
applied.append(name)
|
||||
text_encoder = getattr(sd_model, 'text_encoder', None)
|
||||
if text_encoder is not None:
|
||||
# offload targets the inner model when present: conditioning may call it directly,
|
||||
# and hooks on the wrapper forward would never fire
|
||||
if apply_group_offload_component(getattr(text_encoder, 'model', text_encoder), 'text_encoder', main=False, op=op):
|
||||
applied.append('text_encoder')
|
||||
for name in ('vae', 'audio_vae'):
|
||||
component = getattr(sd_model, name, None)
|
||||
if component is not None:
|
||||
set_group_resident(component)
|
||||
applied.append(f'{name}:device')
|
||||
# has_accelerate stays unset: group hooks are not accelerate hooks, and the modular
|
||||
# pipeline's own to() skips group-offloaded components when move_model runs
|
||||
if any(':' not in name for name in applied):
|
||||
log.info(f'Setting {op}: offload=group type={shared.opts.group_offload_type} modules={applied}')
|
||||
|
||||
|
||||
def apply_group_offload(sd_model, op:str='model'):
|
||||
applied, resident = [], []
|
||||
for module_name in get_module_names(sd_model):
|
||||
module = getattr(sd_model, module_name, None)
|
||||
if not isinstance(module, torch.nn.Module):
|
||||
continue
|
||||
try:
|
||||
role = group_offload_role(module_name, module)
|
||||
if role == 'resident':
|
||||
set_group_resident(module)
|
||||
resident.append(module_name)
|
||||
elif apply_group_offload_component(module, module_name, main=role == 'main', op=op):
|
||||
applied.append(module_name)
|
||||
except Exception as e:
|
||||
log.error(f'Setting {op}: offload=group module={module_name} {e}')
|
||||
set_accelerate(sd_model)
|
||||
if applied:
|
||||
log.info(f'Setting {op}: offload=group type={shared.opts.group_offload_type} modules={applied} resident={resident}')
|
||||
return sd_model
|
||||
|
||||
|
||||
def apply_model_offload(sd_model, op:str='model', quiet:bool=False):
|
||||
try:
|
||||
remove_group_offload(sd_model)
|
||||
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
|
||||
if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
|
||||
shared.opts.diffusers_move_base = False
|
||||
@@ -117,6 +229,7 @@ def apply_model_offload(sd_model, op:str='model', quiet:bool=False):
|
||||
|
||||
def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False):
|
||||
try:
|
||||
remove_group_offload(sd_model)
|
||||
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
|
||||
if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
|
||||
shared.opts.diffusers_move_base = False
|
||||
@@ -144,6 +257,7 @@ def apply_none_offload(sd_model, op:str='model', quiet:bool=False):
|
||||
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
|
||||
try:
|
||||
sd_model.has_accelerate = False
|
||||
remove_group_offload(sd_model)
|
||||
if hasattr(sd_model, 'maybe_free_model_hooks'):
|
||||
sd_model.maybe_free_model_hooks()
|
||||
sd_model = accelerate.hooks.remove_hook_from_module(sd_model, recurse=True)
|
||||
@@ -535,6 +649,7 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc
|
||||
exclude = []
|
||||
if sd_model.__class__.__name__ in balanced_offload_exclude:
|
||||
return sd_model
|
||||
remove_group_offload(sd_model)
|
||||
|
||||
t0 = time.time()
|
||||
cached = True
|
||||
|
||||
@@ -153,6 +153,7 @@ def create_settings(cmd_opts):
|
||||
"group_offload_type": OptionInfo("leaf_level", "Group offload type", gr.Radio, {"choices": ['leaf_level', 'block_level']}),
|
||||
"group_offload_stream": OptionInfo(False, "Use torch streams", gr.Checkbox),
|
||||
'group_offload_record': OptionInfo(False, "Record torch streams", gr.Checkbox),
|
||||
'group_offload_pin': OptionInfo(True, "Pin offload memory", gr.Checkbox),
|
||||
'group_offload_blocks': OptionInfo(1, "Offload blocks", gr.Number),
|
||||
}))
|
||||
|
||||
|
||||
@@ -620,7 +620,7 @@
|
||||
{"id":"","label":"Generic","localized":"","hint":"","ui":"video"},
|
||||
{"id":"","label":"Google GenAI","localized":"","hint":"","ui":"settings_model_options"},
|
||||
{"id":"","label":"Group Offload","localized":"","hint":"","ui":"settings_offload"},
|
||||
{"id":"","label":"Group offload type","localized":"","hint":"Granularity used by <b>group</b> offload.<br>- <b>leaf_level</b>: offloads at the smallest module level; maximum memory savings, slower<br>- <b>block_level</b>: offloads groups of transformer blocks (size set by <b><i>Offload blocks</i></b>); faster with less savings, and keeps the VAE resident<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>leaf_level</b>.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"","label":"Group offload type","localized":"","hint":"Granularity used by <b>group</b> offload.<br>- <b>leaf_level</b>: offloads at the smallest module level; maximum memory savings, slower<br>- <b>block_level</b>: offloads groups of transformer blocks (size set by <b><i>Offload blocks</i></b>); faster with less savings<br>The VAE stays resident on the GPU in both modes, and text encoders always offload at leaf level.<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>leaf_level</b>.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"","label":"Grid Options","localized":"","hint":"","ui":"settings_saving-images"},
|
||||
{"id":"","label":"Grids","localized":"","hint":"","ui":"settings_saving-paths"},
|
||||
{"id":"","label":"Guider","localized":"","hint":"","ui":"txt2img"},
|
||||
@@ -1067,6 +1067,7 @@
|
||||
{"id":"","label":"OpenBody","localized":"","hint":"","ui":"control"}
|
||||
],
|
||||
"p": [
|
||||
{"id":"","label":"Pin offload memory","localized":"","hint":"Keeps the CPU copy of every stream-offloaded weight in pinned non-pageable memory for the fastest transfers, at a host memory cost equal to the full module size.<br>When disabled, memory is pinned one group at a time during transfer: slower, but the weights stay pageable and use no extra memory at rest.<br>Modules larger than half of system memory fall back to per-group pinning automatically.<br><br>Applies only to <b>group</b> offload with <b><i>Use torch streams</i></b> enabled.<br><br>Enabled by default.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"extras_nav","label":"Process","localized":"","hint":"Process existing image<br>Can be used to upscale images, remove backgrounds, obfuscate NSFW content, apply various filters and effects"},
|
||||
{"id":"txt2img_prompts","label":"Prompts","localized":"","hint":"Image prompt and negative prompt","ui":"txt2img"},
|
||||
{"id":"txt2img_pause","label":"Pause","localized":"","hint":"Pause processing","ui":"txt2img"},
|
||||
@@ -1237,7 +1238,7 @@
|
||||
{"id":"","label":"Rebase","localized":"","hint":"","ui":"tab_update"},
|
||||
{"id":"","label":"Repos","localized":"","hint":"","ui":"component-8779"},
|
||||
{"id":"","label":"Refiner model","localized":"","hint":"Refiner model used for second-pass operations","ui":"settings_sd"},
|
||||
{"id":"","label":"Record torch streams","localized":"","hint":"Records CUDA stream usage during <b>group</b> offload so reused buffers stay correct when <b><i>Use torch streams</i></b> is enabled.<br><br>Applies only to <b>group</b> offload with streams.<br><br>Disabled by default.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"","label":"Record torch streams","localized":"","hint":"Skips a stream synchronization each time a group offloads, letting transfers and compute overlap more tightly.<br>Slightly faster at the cost of slightly higher VRAM use; correctness is maintained either way.<br>Has no effect unless <b><i>Use torch streams</i></b> is enabled.<br><br>Applies only to <b>group</b> offload with streams.<br><br>Disabled by default.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"","label":"Remote VAE image type","localized":"","hint":"","ui":"settings_vae_encoder"},
|
||||
{"id":"","label":"Remote VAE for encode","localized":"","hint":"","ui":"settings_vae_encoder"},
|
||||
{"id":"","label":"RAS enabled","localized":"","hint":"","ui":"settings_advanced"},
|
||||
@@ -1562,7 +1563,7 @@
|
||||
{"id":"","label":"Username","localized":"","hint":"","ui":"component-8823"},
|
||||
{"id":"","label":"UNET model","localized":"","hint":"","ui":"settings_sd"},
|
||||
{"id":"","label":"UNET model secondary","localized":"","hint":"Override for the second transformer of dual-transformer architectures:<br>- <b>Ideogram 4</b>: the unconditional transformer<br>- <b>Wan</b> combined stage: the second expert (transformer_2)<br><br>Shown only when the loaded model has a second transformer.<br>Default loads the component from the base model.","ui":"settings_sd"},
|
||||
{"id":"","label":"Use torch streams","localized":"","hint":"In <b>group</b> offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.<br>Faster, but uses more VRAM.<br><br>Applies only to <b>group</b> offload.<br><br>Disabled by default.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"","label":"Use torch streams","localized":"","hint":"In <b>group</b> offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.<br>Faster, but offloaded weights are then staged in pinned non-pageable host memory: the whole module when <b><i>Pin offload memory</i></b> is enabled, one group at a time otherwise.<br>Text encoders are exempt and always offload without streams.<br><br>Applies only to <b>group</b> offload.<br><br>Disabled by default.","reload":"model","ui":"settings_offload"},
|
||||
{"id":"","label":"Use SVD quantization","localized":"","hint":"Adds a low-rank (SVDQuant) correction on top of SDNQ to recover accuracy lost at low bit widths, at the cost of extra size and compute.<br>Tuned by <b><i>SVD rank size</i></b> and <b><i>SVD steps</i></b>.<br><br>Disabled by default.","reload":"model","ui":"settings_quantization"},
|
||||
{"id":"","label":"Use Dynamic quantization","localized":"","hint":"Picks a per-layer weight type automatically instead of one type everywhere. Each layer starts at the <b><i>Quantization type</i></b> (the minimum) and steps up to higher precision until its error meets the <b><i>Dynamic loss threshold</i></b>.<br>Protects error-sensitive layers at the cost of a larger model and slower load.<br><br>Disabled by default.","reload":"model","ui":"settings_quantization"},
|
||||
{"id":"","label":"Use Hadamard rotations","localized":"","hint":"Applies a Hadamard rotation before quantizing to spread out weight outliers, which can improve accuracy at low bit widths.<br>Group size is set by <b><i>Hadamard group size</i></b>.<br><br>Disabled by default.","reload":"model","ui":"settings_quantization"},
|
||||
|
||||
Reference in New Issue
Block a user