From 09e861157f9c8ee7cbd25b81a424df8a07a20128 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 16 Sep 2026 01:57:23 +0100 Subject: [PATCH] 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)