From b7aff134a2f50f41b0371489506408a7ca600e85 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 30 Nov 2024 19:03:51 -0500 Subject: [PATCH] add low/high threshold to balanced offload Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ modules/devices.py | 3 ++- modules/lora/networks.py | 1 + modules/processing_helpers.py | 5 ++--- modules/sd_models.py | 27 ++++++++++++++++++++------- modules/shared.py | 5 +++-- 6 files changed, 30 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ab3bcbd1..c62b6b917 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,8 @@ - Flux: do not recast quants - **Offload** improvements: - faster and more compatible *balanced* mode + - balanced offload: units are now in percentage instead of bytes + - balanced offload: add both high and low watermark - **UI**: - improved stats on generate completion - improved live preview display and performance diff --git a/modules/devices.py b/modules/devices.py index 9ca1863a5..64968a30c 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -224,7 +224,7 @@ def torch_gc(force=False, fast=False): timer.process.records['gc'] = 0 timer.process.records['gc'] += t1 - t0 if not force or collected == 0: - return + return used_gpu mem = memstats.memory_stats() saved = round(gpu.get('used', 0) - mem.get('gpu', {}).get('used', 0), 2) before = { 'gpu': gpu.get('used', 0), 'ram': ram.get('used', 0) } @@ -233,6 +233,7 @@ def torch_gc(force=False, fast=False): results = { 'collected': collected, 'saved': saved } fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access log.debug(f'GC: utilization={utilization} gc={results} before={before} after={after} device={torch.device(get_optimal_device_name())} fn={fn} time={round(t1 - t0, 2)}') # pylint: disable=protected-access + return used_gpu def set_cuda_sync_mode(mode): diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 604e591a9..69db5fce3 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -218,6 +218,7 @@ def maybe_recompile_model(names, te_multipliers): def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None): + timer['list'] = 0 global backup_size # pylint: disable=global-statement networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names] if any(x is None for x in networks_on_disk): diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index ab08d4cc8..5d2661cc2 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -368,14 +368,13 @@ def validate_sample(tensor): sample = 255.0 * np.moveaxis(sample, 0, 2) if not shared.native else 255.0 * sample with warnings.catch_warnings(record=True) as w: cast = sample.astype(np.uint8) - minimum, maximum, mean = np.min(cast), np.max(cast), np.mean(cast) - if len(w) > 0 or minimum == maximum: + if len(w) > 0: nans = np.isnan(sample).sum() cast = np.nan_to_num(sample) cast = cast.astype(np.uint8) vae = shared.sd_model.vae.dtype if hasattr(shared.sd_model, 'vae') else None upcast = getattr(shared.sd_model.vae.config, 'force_upcast', None) if hasattr(shared.sd_model, 'vae') and hasattr(shared.sd_model.vae, 'config') else None - shared.log.error(f'Decode: sample={sample.shape} invalid={nans} mean={mean} dtype={dtype} vae={vae} upcast={upcast} failed to validate') + shared.log.error(f'Decode: sample={sample.shape} invalid={nans} dtype={dtype} vae={vae} upcast={upcast} failed to validate') if upcast is not None and not upcast: setattr(shared.sd_model.vae.config, 'force_upcast', True) # noqa: B010 shared.log.warning('Decode: upcast=True set, retry operation') diff --git a/modules/sd_models.py b/modules/sd_models.py index 6c3ddc6b5..42bd33d82 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -361,7 +361,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'): shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') if shared.opts.diffusers_offload_mode == "balanced": try: - shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} threshold={shared.opts.diffusers_offload_max_gpu_memory} limit={shared.opts.cuda_mem_fraction}') + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} watermarks low={shared.opts.diffusers_offload_min_gpu_memory} high={shared.opts.diffusers_offload_max_gpu_memory} limit={shared.opts.cuda_mem_fraction:.2f}') sd_model = apply_balanced_offload(sd_model) except Exception as e: shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') @@ -369,6 +369,16 @@ def set_diffuser_offload(sd_model, op: str = 'model'): class OffloadHook(accelerate.hooks.ModelHook): + def __init__(self): + if shared.opts.diffusers_offload_max_gpu_memory > 1: + shared.opts.diffusers_offload_max_gpu_memory = 0.75 + if shared.opts.diffusers_offload_max_cpu_memory > 1: + shared.opts.diffusers_offload_max_cpu_memory = 0.75 + self.gpu = int(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024) + self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024) + shared.log.info(f'Init offload: type=balanced gpu={self.gpu} cpu={self.cpu}') + super().__init__() + def init_hook(self, module): return module @@ -377,10 +387,7 @@ class OffloadHook(accelerate.hooks.ModelHook): device_index = torch.device(devices.device).index if device_index is None: device_index = 0 - max_memory = { - device_index: int(shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024), - "cpu": int(shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024), - } + max_memory = { device_index: self.gpu, "cpu": self.cpu } device_map = getattr(module, "balanced_offload_device_map", None) if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None): device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory) @@ -399,10 +406,13 @@ class OffloadHook(accelerate.hooks.ModelHook): return module -offload_hook_instance = OffloadHook() +offload_hook_instance = None def apply_balanced_offload(sd_model): + global offload_hook_instance # pylint: disable=global-statement + if offload_hook_instance is None: + offload_hook_instance = OffloadHook() t0 = time.time() excluded = ['OmniGenPipeline'] if sd_model.__class__.__name__ in excluded: @@ -414,6 +424,7 @@ def apply_balanced_offload(sd_model): checkpoint_name = sd_model.__class__.__name__ def apply_balanced_offload_to_module(pipe): + used_gpu = devices.torch_gc(fast=True) if hasattr(pipe, "pipe"): apply_balanced_offload_to_module(pipe.pipe) if hasattr(pipe, "_internal_dict"): @@ -429,7 +440,9 @@ def apply_balanced_offload(sd_model): max_memory = getattr(module, "balanced_offload_max_memory", None) module = accelerate.hooks.remove_hook_from_module(module, recurse=True) try: - module = module.to(devices.cpu, non_blocking=True) + if used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory: + module = module.to(devices.cpu, non_blocking=True) + used_gpu = devices.torch_gc(fast=True) module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access diff --git a/modules/shared.py b/modules/shared.py index bcb506cee..21a70fea1 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -560,8 +560,9 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_extract_ema": OptionInfo(False, "Use model EMA weights when possible"), "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}), "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'model', 'sequential']}), - "diffusers_offload_max_gpu_memory": OptionInfo(round(gpu_memory * 0.75, 1), "Max GPU memory before balanced offload", gr.Slider, {"minimum": 0, "maximum": gpu_memory, "step": 0.01, "visible": True }), - "diffusers_offload_max_cpu_memory": OptionInfo(round(cpu_memory * 0.75, 1), "Max CPU memory before balanced offload", gr.Slider, {"minimum": 0, "maximum": cpu_memory, "step": 0.01, "visible": False }), + "diffusers_offload_min_gpu_memory": OptionInfo(0.25, "Balanced offload GPU low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), + "diffusers_offload_max_gpu_memory": OptionInfo(0.75, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), + "diffusers_offload_max_cpu_memory": OptionInfo(0.75, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, {"choices": ['default', 'true', 'false']}), "diffusers_vae_slicing": OptionInfo(True, "VAE slicing"), "diffusers_vae_tiling": OptionInfo(cmd_opts.lowvram or cmd_opts.medvram, "VAE tiling"),