optimize balanced offload

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -50,8 +50,10 @@
 - **Memory** improvements:  
   - faster and more compatible *balanced offload* mode  
   - balanced offload: units are now in percentage instead of bytes  
-  - balanced offload: add both high and low watermark  
-    default is 25% for low-watermark (skip offload if memory usage is below 25%) and 70% high-watermark (must offload if memory usage is above 70%)  
+  - balanced offload: add both high and low watermark and pinned threshold, defaults as below  
+    25% for low-watermark: skip offload if memory usage is below 25%  
+    70% high-watermark: must offload if memory usage is above 70%  
+    15% pin-watermark: any model component smaller than 15% of total memory is pinned and not offloaded  
   - change-in-behavior:  
     low-end systems, triggered by either `lowvrwam` or by detection of <=4GB will use *sequential offload*  
     all other systems use *balanced offload* by default (can be changed in settings)  

modules/memstats.py

@@ -4,6 +4,8 @@ import torch
 from modules import shared, errors
 
 fail_once = False
+mem = {}
+
 
 def gb(val: float):
     return round(val / 1024 / 1024 / 1024, 2)
@@ -11,7 +13,7 @@ def gb(val: float):
 
 def memory_stats():
     global fail_once # pylint: disable=global-statement
-    mem = {}
+    mem.clear()
     try:
         process = psutil.Process(os.getpid())
         res = process.memory_info()
@@ -41,6 +43,10 @@ def memory_stats():
     return mem
 
 
+def memory_cache():
+    return mem
+
+
 def ram_stats():
     try:
         process = psutil.Process(os.getpid())

modules/processing_diffusers.py

@@ -349,7 +349,7 @@ def process_refine(p: processing.StableDiffusionProcessing, output):
 
 
 def process_decode(p: processing.StableDiffusionProcessing, output):
-    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae'])
     if output is not None:
         if not hasattr(output, 'images') and hasattr(output, 'frames'):
             shared.log.debug(f'Generated: frames={len(output.frames[0])}')
@@ -463,7 +463,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
     timer.process.record('decode')
 
     shared.sd_model = orig_pipeline
-    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+    # shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
 
     if p.state == '':
         global last_p # pylint: disable=global-statement

modules/processing_vae.py

@@ -104,8 +104,6 @@ def full_vae_decode(latents, model):
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False):
         base_device = sd_models.move_base(model, devices.cpu)
 
-    if shared.opts.diffusers_offload_mode == "balanced":
-        shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     elif shared.opts.diffusers_offload_mode != "sequential":
         sd_models.move_model(model.vae, devices.device)
 
@@ -159,8 +157,8 @@ def full_vae_decode(latents, model):
             model.vae.apply(sd_models.convert_to_faketensors)
             devices.torch_gc(force=True)
 
-    if shared.opts.diffusers_offload_mode == "balanced":
-        shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+    # if shared.opts.diffusers_offload_mode == "balanced":
+    #    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     elif shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None:
         sd_models.move_base(model, base_device)
     t1 = time.time()

modules/sd_models.py

@@ -18,7 +18,7 @@ from omegaconf import OmegaConf
 from ldm.util import instantiate_from_config
 from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect
 from modules.timer import Timer, process as process_timer
-from modules.memstats import memory_stats
+from modules.memstats import memory_stats, memory_cache
 from modules.modeldata import model_data
 from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import
 
@@ -416,9 +416,10 @@ class OffloadHook(accelerate.hooks.ModelHook):
 
 
 offload_hook_instance = None
+offload_component_map = {}
 
 
-def apply_balanced_offload(sd_model):
+def apply_balanced_offload(sd_model, exclude=[]):
     global offload_hook_instance # pylint: disable=global-statement
     if shared.opts.diffusers_offload_mode != "balanced":
         return sd_model
@@ -428,8 +429,6 @@ def apply_balanced_offload(sd_model):
     excluded = ['OmniGenPipeline']
     if sd_model.__class__.__name__ in excluded:
         return sd_model
-    fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
-    debug_move(f'Apply offload: type=balanced fn={fn}')
     checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else None
     if checkpoint_name is None:
         checkpoint_name = sd_model.__class__.__name__
@@ -442,32 +441,38 @@ def apply_balanced_offload(sd_model):
             keys = pipe._internal_dict.keys() # pylint: disable=protected-access
         else:
             keys = get_signature(pipe).keys()
+        keys = [k for k in keys if k not in exclude and not k.startswith('_')]
         for module_name in keys: # pylint: disable=protected-access
             module = getattr(pipe, module_name, None)
-            if isinstance(module, torch.nn.Module):
-                network_layer_name = getattr(module, "network_layer_name", None)
-                device_map = getattr(module, "balanced_offload_device_map", None)
-                max_memory = getattr(module, "balanced_offload_max_memory", None)
-                module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
-                try:
-                    do_offload = used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory
-                    debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} dtype={module.dtype} op={"move" if do_offload else "skip"} component={module.__class__.__name__}')
-                    if do_offload:
-                        module = module.to(devices.cpu)
-                        used_gpu, used_ram = devices.torch_gc(fast=True, force=True)
-                except Exception as e:
-                    if 'bitsandbytes' not in str(e):
-                        shared.log.error(f'Balanced offload: module={module_name} {e}')
-                    if os.environ.get('SD_MOVE_DEBUG', None):
-                        errors.display(e, f'Balanced offload: module={module_name}')
-                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
-                if network_layer_name:
-                    module.network_layer_name = network_layer_name
-                if device_map and max_memory:
-                    module.balanced_offload_device_map = device_map
-                    module.balanced_offload_max_memory = max_memory
+            if not isinstance(module, torch.nn.Module):
+                continue
+            network_layer_name = getattr(module, "network_layer_name", None)
+            device_map = getattr(module, "balanced_offload_device_map", None)
+            max_memory = getattr(module, "balanced_offload_max_memory", None)
+            module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
+            module_size = offload_component_map.get(module_name, None)
+            if module_size is None:
+                module_size = sum(p.numel()*p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024
+                offload_component_map[module_name] = module_size
+            do_offload = (used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory) and (module_size > shared.gpu_memory * shared.opts.diffusers_offload_pin_gpu_memory)
+            try:
+                debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} dtype={module.dtype} op={"move" if do_offload else "skip"} component={module.__class__.__name__} size={module_size:.3f}')
+                if do_offload and module.device != devices.cpu:
+                    module = module.to(devices.cpu)
+                    used_gpu, used_ram = devices.torch_gc(fast=True, force=True)
+            except Exception as e:
+                if 'bitsandbytes' not in str(e):
+                    shared.log.error(f'Balanced offload: module={module_name} {e}')
+                if os.environ.get('SD_MOVE_DEBUG', None):
+                    errors.display(e, f'Balanced offload: module={module_name}')
+            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
+            if network_layer_name:
+                module.network_layer_name = network_layer_name
+            if device_map and max_memory:
+                module.balanced_offload_device_map = device_map
+                module.balanced_offload_max_memory = max_memory
 
     apply_balanced_offload_to_module(sd_model)
     if hasattr(sd_model, "pipe"):
@@ -478,7 +483,10 @@ def apply_balanced_offload(sd_model):
         apply_balanced_offload_to_module(sd_model.decoder_pipe)
     set_accelerate(sd_model)
     devices.torch_gc(fast=True)
-    process_timer.add('offload', time.time() - t0)
+    t = time.time() - t0
+    process_timer.add('offload', t)
+    fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+    debug_move(f'Apply offload: time={t:.2f} type=balanced fn={fn}')
     return sd_model
 
 

modules/shared.py

@@ -483,6 +483,7 @@ options_templates.update(options_section(('sd', "Models & Loading"), {
     "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'model', 'sequential']}),
     "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.70, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
+    "diffusers_offload_pin_gpu_memory": OptionInfo(0.15, "Balanced offload GPU pin watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
     "diffusers_offload_max_cpu_memory": OptionInfo(0.90, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
 
     "advanced_sep": OptionInfo("<h2>Advanced Options</h2>", "", gr.HTML),

scripts/cogvideo.py

@@ -180,7 +180,7 @@ class Script(scripts.Script):
             else:
                 shared.sd_model = sd_models.switch_pipe(diffusers.CogVideoXPipeline, shared.sd_model)
             args['num_frames'] = p.frames # only txt2vid has num_frames
-            shared.log.info(f'CogVideoX: class={shared.sd_model.__class__.__name__} frames={p.frames} input={args.get('video', None) or args.get('image', None)}')
+            shared.log.info(f"CogVideoX: class={shared.sd_model.__class__.__name__} frames={p.frames} input={args.get('video', None) or args.get('image', None)}")
             if debug:
                 shared.log.debug(f'CogVideoX args: {args}')
             frames = shared.sd_model(**args).frames[0]