From 09e861157f9c8ee7cbd25b81a424df8a07a20128 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 16 Sep 2026 01:57:23 +0100 Subject: [PATCH 1/3] fix(offload): hand back the loaded cpu tensors on no-stream offload Group hooks on the no-stream path and the on-demand hook return a component to cpu through a device copy, so a memory-mapped text encoder sits in memory twice, as the mapped file its never-run vision tower keeps alive and as the copies, and every encode pays a device-to-host transfer of unchanged weights. The engine now records the cpu tensors at onload and hands them back at offload; a component moved by any other path still takes the copy. --- modules/sd_offload_group.py | 64 ++++++++++++++++++++++++++++++++++++- test/test-offload-roles.py | 55 +++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 1 deletion(-) diff --git a/modules/sd_offload_group.py b/modules/sd_offload_group.py index 014456be7..51967b3b4 100644 --- a/modules/sd_offload_group.py +++ b/modules/sd_offload_group.py @@ -1,4 +1,5 @@ import time +import types import itertools import torch import accelerate.hooks @@ -28,6 +29,59 @@ def group_offload_config(main: bool) -> dict: } +def group_tensors(group) -> list: + """Every parameter and buffer a diffusers group moves, read at call time so a tensor replaced while offloaded is seen.""" + tensors = [] + for module in group.modules: + tensors.extend(module.parameters()) + tensors.extend(module.buffers()) + tensors.extend(group.parameters) + tensors.extend(group.buffers) + return tensors + + +def loaded_tensors(tensors) -> dict: + """The cpu tensors held right now, keyed by parameter, to hand back on offload in place of a fresh copy.""" + return {t: t.data for t in tensors if t.data.device.type == 'cpu'} + + +def restore_tensors(tensors, loaded: dict | None, non_blocking: bool = False): + """Return tensors to cpu: the tensor each was onloaded from where known, a copy otherwise.""" + for t in tensors: + if t.data.device.type == 'cpu': + continue + source = loaded.get(t) if loaded else None + t.data = source if source is not None else t.data.to(devices.cpu, non_blocking=non_blocking) + + +def onload_remember(group): + group.sdnext_loaded = loaded_tensors(group_tensors(group)) + group.sdnext_onload() + + +def offload_restore(group): + restore_tensors(group_tensors(group), getattr(group, 'sdnext_loaded', None)) + group.sdnext_loaded = None # a record lives from one onload to its offload + + +def keep_loaded_tensors(module) -> int: + """Groups on the no-stream path copy their weights to fresh cpu memory on every offload; record the cpu + tensors at onload and hand them back at offload instead. Returns the number of groups patched.""" + from diffusers.hooks.group_offloading import _GROUP_OFFLOADING + count = 0 + for sub in module.modules(): + registry = getattr(sub, '_diffusers_hook', None) + hook = registry.get_hook(_GROUP_OFFLOADING) if registry is not None else None + group = getattr(hook, 'group', None) + if group is None or group.stream is not None or getattr(group, 'offload_to_disk_path', None) or hasattr(group, 'sdnext_onload'): + continue + group.sdnext_onload = group._onload_from_memory # pylint: disable=protected-access + group._onload_from_memory = types.MethodType(onload_remember, group) # pylint: disable=protected-access + group._offload_to_memory = types.MethodType(offload_restore, group) # pylint: disable=protected-access + count += 1 + return count + + def remove_group_offload_component(module) -> bool: if getattr(module, 'sdnext_group_offload_sig', None) is None: module = getattr(module, 'model', None) # wrapper components carry the hooks on the inner model @@ -99,6 +153,8 @@ def apply_group_offload_component(module, module_name: str, main: bool) -> bool: s.debug_move(f'Offload: type=group op=apply type={shared.opts.group_offload_type} module={module_name} pin={cfg["use_stream"] and not cfg["low_cpu_mem_usage"]}') # before the apply: pinning large components takes a while and would otherwise run silently module.sdnext_group_offload_sig = 'partial' # a raise below leaves hooks that only a non-empty signature will remove apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg) + if not cfg['use_stream']: + s.debug_move(f'Offload: type=group op=keep module={module_name} groups={keep_loaded_tensors(module)}') module.sdnext_group_offload_sig = sig return True @@ -153,6 +209,7 @@ class OnDemandHook(accelerate.hooks.ModelHook): param = next(module.parameters(), None) if param is not None and not devices.same_device(param.device, devices.device): t0 = time.time() + module.sdnext_loaded = loaded_tensors(list(module.parameters()) + list(module.buffers())) module.to(devices.device, non_blocking=shared.opts.diffusers_offload_nonblocking) t1 = time.time() process_timer.add('onload', t1 - t0) @@ -194,7 +251,12 @@ def offload_ondemand(sd_model, include=[], exclude=[], reason='', force=False): continue try: t0 = time.time() - module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking) + loaded = getattr(module, 'sdnext_loaded', None) + if loaded: + restore_tensors(list(module.parameters()) + list(module.buffers()), loaded, non_blocking=shared.opts.diffusers_offload_nonblocking) + module.sdnext_loaded = None # a record lives from one onload to its offload + else: + module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking) dt = time.time() - t0 process_timer.add('offload', dt) moved.append(module_name) diff --git a/test/test-offload-roles.py b/test/test-offload-roles.py index 533e4a8d7..8af4bfe56 100644 --- a/test/test-offload-roles.py +++ b/test/test-offload-roles.py @@ -466,6 +466,59 @@ def test_resident_placement_clears_the_ondemand_stamp(): assert not hasattr(module, '_hf_hook'), 'the on-demand hook must be removed' +def accelerator(): + """The restore round trip needs a device that swaps tensor data with cpu both ways; meta cannot, so these two tests take a real accelerator or skip.""" + if torch.cuda.is_available(): + return torch.device('cuda') + log.warning(' SKIP: no accelerator for the round trip') + return None + + +def test_ondemand_offload_hands_back_the_loaded_tensors(): + device = accelerator() + if device is None: + return True + module = BridgeModule() + sd_offload_group.apply_group_offload_ondemand(module) + loaded = {name: p.data for name, p in module.named_parameters()} + pipe = FakePipe({'vae': module}) + orig_device = sd_offload_group.devices.device + sd_offload_group.devices.device = device + try: + module._hf_hook.pre_forward(module, torch.zeros(1, 4)) # pylint: disable=protected-access + assert next(module.parameters()).device.type == device.type, 'the entry hook must onload the whole module' + sd_offload_group.offload_ondemand(pipe, force=True) + finally: + sd_offload_group.devices.device = orig_device + for name, param in module.named_parameters(): + assert param.data.data_ptr() == loaded[name].data_ptr(), f'{name} came back as a copy rather than the loaded tensor' + return True + + +def test_group_offload_hands_back_the_loaded_tensors(): + from diffusers.hooks.group_offloading import _GROUP_OFFLOADING + device = accelerator() + if device is None: + return True + module = PlainModule() + orig_device = sd_offload_group.devices.device + sd_offload_group.devices.device = device + try: + assert sd_offload_group.apply_group_offload_component(module, 'text_encoder', main=False) is True + loaded = {name: p.data for name, p in module.named_parameters()} + group = module.proj._diffusers_hook.get_hook(_GROUP_OFFLOADING).group # pylint: disable=protected-access + assert group.stream is None and hasattr(group, 'sdnext_onload'), 'aux components take the no-stream path and must carry the restore patch' + group.onload_() + assert next(module.parameters()).device.type == device.type, 'onload must still move the group' + group.offload_() + finally: + sd_offload_group.devices.device = orig_device + for name, param in module.named_parameters(): + assert param.data.data_ptr() == loaded[name].data_ptr(), f'{name} came back as a copy rather than the loaded tensor' + assert sd_offload_group.keep_loaded_tensors(module) == 0, 'a second pass must not patch the same groups again' + return True + + # ============================================================ # get_module_names # ============================================================ @@ -706,6 +759,8 @@ def run_all(): test_ondemand_apply_returns_bool_and_is_idempotent, test_ondemand_apply_leaves_weights_on_cpu, test_resident_placement_clears_the_ondemand_stamp, + test_ondemand_offload_hands_back_the_loaded_tensors, + test_group_offload_hands_back_the_loaded_tensors, ]: run_test(cat, fn) From c3fc8560cb1fe26c598d2e0a301096970720dec3 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 16 Sep 2026 02:13:28 +0100 Subject: [PATCH 2/3] docs(offload): describe the group offload blocks setting The blocks count is clamped to one under streams and leaf level never reads it; the hint says so, and the type hint names the setting it refers to. --- ui/locale/locale_en.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 47bb4dd2d..9b736af0d 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -624,7 +624,8 @@ {"id":"","label":"Generic","localized":"","hint":"","ui":"video"}, {"id":"","label":"Google GenAI","localized":"","hint":"","ui":"settings_model_options"}, {"id":"","label":"Group Offload","localized":"","hint":"Offloads components in groups of layers rather than as a whole, so only the layers in use occupy VRAM.
Lets a single component larger than the card run, at the cost of transferring weights throughout every step.

Applies only when Model offload mode is group.","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
This setting applies to the parts of the model that run at every step. Components used once per generation, such as text encoders, always offload at leaf_level. The VAE is handled separately: it waits in system memory and loads as a whole when encoding or decoding.
Anything named in Modules to never offload or Model types not to offload stays in VRAM instead.

Applies only when Model offload mode is group.

Default is leaf_level.","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 Group offload blocks, one block when Prefetch with streams is enabled); faster with less savings
This setting applies to the parts of the model that run at every step. Components used once per generation, such as text encoders, always offload at leaf_level. The VAE is handled separately: it waits in system memory and loads as a whole when encoding or decoding.
Anything named in Modules to never offload or Model types not to offload stays in VRAM instead.

Applies only when Model offload mode is group.

Default is leaf_level.","ui":"settings_offload"}, + {"id":"","label":"Group offload blocks","localized":"","hint":"Number of transformer blocks moved together as one group on block_level group offload. Larger groups mean fewer, larger transfers and more weights resident in VRAM at once.
Ignored when Prefetch with streams is enabled, which runs one block per group, and on leaf_level, which has no blocks. Components used once per generation always offload at leaf_level and never read this value.

Applies only when Model offload mode is group.

Default is 1.","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"}, From a7d430a8bac5232eb2d7dd7584a4b3541e777f7d Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 16 Sep 2026 02:18:07 +0100 Subject: [PATCH 3/3] feat(offload): warn when streams clamp the group offload blocks Streamed prefetch runs one block per group, so a larger blocks setting is reduced to one at apply time. The clamp is now logged once per component. --- modules/sd_offload_group.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modules/sd_offload_group.py b/modules/sd_offload_group.py index 51967b3b4..96bfeffa6 100644 --- a/modules/sd_offload_group.py +++ b/modules/sd_offload_group.py @@ -145,6 +145,9 @@ def apply_group_offload_component(module, module_name: str, main: bool) -> bool: sig = f'{devices.device}:{main}:' + ':'.join(str(v) for v in cfg.values()) if getattr(module, 'sdnext_group_offload_sig', None) == sig: return False + requested_blocks = int(shared.opts.group_offload_blocks) + if cfg['use_stream'] and requested_blocks > 1: + log.warning(f'Offload: type=group module={module_name} blocks={requested_blocks} streams=True clamped=1') 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) module.sdnext_ondemand = False # group placement replaces any on-demand hook