mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
Merge pull request #5092 from CalamitousFelicitousness/fix/offload-loaded-tensors
fix(offload): hand back the loaded cpu tensors on no-stream offload
This commit is contained in:
Vladimir Mandic
GitHub
@@ -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
|
||||
@@ -91,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
|
||||
@@ -99,6 +156,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 +212,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 +254,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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.<br>Lets a single component larger than the card run, at the cost of transferring weights throughout every step.<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.","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>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 <b>leaf_level</b>. The VAE is handled separately: it waits in system memory and loads as a whole when encoding or decoding.<br>Anything named in <b><i>Modules to never offload</i></b> or <b><i>Model types not to offload</i></b> stays in VRAM instead.<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>leaf_level</b>.","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>Group offload blocks</i></b>, one block when <b><i>Prefetch with streams</i></b> is enabled); faster with less savings<br>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 <b>leaf_level</b>. The VAE is handled separately: it waits in system memory and loads as a whole when encoding or decoding.<br>Anything named in <b><i>Modules to never offload</i></b> or <b><i>Model types not to offload</i></b> stays in VRAM instead.<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>leaf_level</b>.","ui":"settings_offload"},
|
||||
{"id":"","label":"Group offload blocks","localized":"","hint":"Number of transformer blocks moved together as one group on <b>block_level</b> group offload. Larger groups mean fewer, larger transfers and more weights resident in VRAM at once.<br>Ignored when <b><i>Prefetch with streams</i></b> is enabled, which runs one block per group, and on <b>leaf_level</b>, which has no blocks. Components used once per generation always offload at <b>leaf_level</b> and never read this value.<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>1</b>.","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"},
|
||||
|
||||
Reference in New Issue
Block a user