From 46c30a00f24c32f11e7823cea511c5382827b4fa Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 7 Aug 2026 22:45:16 +0100 Subject: [PATCH] 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 --- modules/sd_offload.py | 147 +++++++++++++++++++++++++++++++++----- modules/ui_definitions.py | 1 + ui/locale/locale_en.json | 7 +- 3 files changed, 136 insertions(+), 19 deletions(-) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index ca6c66570..58699d8b0 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -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 diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 5944dc426..4aa269a11 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -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), })) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 149d6f381..710940442 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -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 group offload.
- leaf_level: offloads at the smallest module level; maximum memory savings, slower
- block_level: offloads groups of transformer blocks (size set by Offload blocks); faster with less savings, and keeps the VAE resident

Applies only when Model offload mode is group.

Default is leaf_level.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Group offload type","localized":"","hint":"Granularity used by group offload.
- leaf_level: offloads at the smallest module level; maximum memory savings, slower
- block_level: offloads groups of transformer blocks (size set by Offload blocks); faster with less savings
The VAE stays resident on the GPU in both modes, and text encoders always offload at leaf level.

Applies only when Model offload mode is group.

Default is leaf_level.","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.
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.
Modules larger than half of system memory fall back to per-group pinning automatically.

Applies only to group offload with Use torch streams enabled.

Enabled by default.","reload":"model","ui":"settings_offload"}, {"id":"extras_nav","label":"Process","localized":"","hint":"Process existing image
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 group offload so reused buffers stay correct when Use torch streams is enabled.

Applies only to group offload with streams.

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.
Slightly faster at the cost of slightly higher VRAM use; correctness is maintained either way.
Has no effect unless Use torch streams is enabled.

Applies only to group offload with streams.

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:
- Ideogram 4: the unconditional transformer
- Wan combined stage: the second expert (transformer_2)

Shown only when the loaded model has a second transformer.
Default loads the component from the base model.","ui":"settings_sd"}, - {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but uses more VRAM.

Applies only to group offload.

Disabled by default.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but offloaded weights are then staged in pinned non-pageable host memory: the whole module when Pin offload memory is enabled, one group at a time otherwise.
Text encoders are exempt and always offload without streams.

Applies only to group offload.

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.
Tuned by SVD rank size and SVD steps.

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 Quantization type (the minimum) and steps up to higher precision until its error meets the Dynamic loss threshold.
Protects error-sensitive layers at the cost of a larger model and slower load.

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.
Group size is set by Hadamard group size.

Disabled by default.","reload":"model","ui":"settings_quantization"},