add native torch fp8 storage dtype

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

modules/model_quant.py

@@ -8,6 +8,24 @@ quanto = None
 ao = None
 
 
+def get_quant(name):
+    if "qint8" in name.lower():
+        return 'qint8'
+    if "qint4" in name.lower():
+        return 'qint4'
+    if "fp8" in name.lower():
+        return 'fp8'
+    if "fp4" in name.lower():
+        return 'fp4'
+    if "nf4" in name.lower():
+        return 'nf4'
+    if name.endswith('.gguf'):
+        return 'gguf'
+    return 'none'
+
+
+
+
 def create_bnb_config(kwargs = None, allow_bnb: bool = True):
     from modules import shared, devices
     if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
@@ -117,17 +135,43 @@ def load_quanto(msg='', silent=False):
     return None
 
 
-def get_quant(name):
-    if "qint8" in name.lower():
-        return 'qint8'
-    if "qint4" in name.lower():
-        return 'qint4'
-    if "fp8" in name.lower():
-        return 'fp8'
-    if "fp4" in name.lower():
-        return 'fp4'
-    if "nf4" in name.lower():
-        return 'nf4'
-    if name.endswith('.gguf'):
-        return 'gguf'
-    return 'none'
+def apply_layerwise(sd_model, quiet:bool=False):
+    import torch
+    from diffusers.quantizers import quantization_config
+    from modules import shared, devices, sd_models
+    if shared.opts.layerwise_quantization_storage == 'float8_e4m3fn' and hasattr(torch, 'float8_e4m3fn'):
+        storage_dtype = torch.float8_e4m3fn
+    elif shared.opts.layerwise_quantization_storage == 'float8_e5m2' and hasattr(torch, 'float8_e5m2'):
+        storage_dtype = torch.float8_e5m2
+    else:
+        storage_dtype = None
+        shared.log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported')
+        return
+    non_blocking = False
+    if not hasattr(quantization_config.QuantizationMethod, 'LAYERWISE'):
+        setattr(quantization_config.QuantizationMethod, 'LAYERWISE', 'layerwise') # noqa: B010
+    for module in sd_models.get_signature(sd_model).keys():
+        if not hasattr(sd_model, module):
+            continue
+        try:
+            cls = getattr(sd_model, module).__class__.__name__
+            if module.startswith('unet') and ('Model' in shared.opts.layerwise_quantization):
+                m = getattr(sd_model, module)
+                if hasattr(m, 'enable_layerwise_casting'):
+                    m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
+                    m.quantization_method = 'LayerWise'
+                    log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
+            if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization or 'Transformer' in shared.opts.layerwise_quantization):
+                m = getattr(sd_model, module)
+                if hasattr(m, 'enable_layerwise_casting'):
+                    m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
+                    m.quantization_method = 'LayerWise'
+                    log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
+            if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'Text Encoder' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
+                m = getattr(sd_model, module)
+                if hasattr(m, 'enable_layerwise_casting'):
+                    m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
+                    m.quantization_method = quantization_config.QuantizationMethod.LAYERWISE # pylint: disable=no-member
+                    log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
+        except Exception as e:
+            shared.log.error(f'Quantization: type=layerwise {e}')

modules/sd_models.py

@@ -9,7 +9,7 @@ import diffusers
 import diffusers.loaders.single_file_utils
 import torch
 
-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 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, model_quant
 from modules.timer import Timer, process as process_timer
 from modules.memstats import memory_stats
 from modules.modeldata import model_data
@@ -570,9 +570,12 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             sd_model = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU
         if shared.opts.optimum_quanto_weights:
             sd_model = sd_models_compile.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU
+        if shared.opts.layerwise_quantization:
+            model_quant.apply_layerwise(sd_model)
         timer.record("options")
 
         set_diffuser_offload(sd_model, op)
+
         if op == 'model' and not (os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface'):
             if getattr(shared.sd_model, 'sd_checkpoint_info', None) is not None and vae_file is not None:
                 sd_vae.apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model)

modules/sd_offload.py

@@ -3,9 +3,9 @@ import sys
 import time
 import inspect
 import torch
-import accelerate
+import accelerate.hooks
 
-from modules import shared, devices, errors
+from modules import shared, devices, errors, model_quant
 from modules.timer import process as process_timer
 
 
@@ -20,7 +20,6 @@ def get_signature(cls):
 
 
 def disable_offload(sd_model):
-    from accelerate.hooks import remove_hook_from_module
     if not getattr(sd_model, 'has_accelerate', False):
         return
     if hasattr(sd_model, "_internal_dict"):
@@ -31,7 +30,7 @@ def disable_offload(sd_model):
         module = getattr(sd_model, module_name, None)
         if isinstance(module, torch.nn.Module):
             network_layer_name = getattr(module, "network_layer_name", None)
-            module = remove_hook_from_module(module, recurse=True)
+            module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
             if network_layer_name:
                 module.network_layer_name = network_layer_name
     sd_model.has_accelerate = False
@@ -188,7 +187,7 @@ def apply_balanced_offload(sd_model, exclude=[]):
     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__
-    if offload_hook_instance is None or offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory or offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory or checkpoint_name != offload_hook_instance.checkpoint_name:
+    if (offload_hook_instance is None) or (offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory) or (offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory) or (checkpoint_name != offload_hook_instance.checkpoint_name):
         cached = False
         offload_hook_instance = OffloadHook(checkpoint_name)
 
@@ -241,9 +240,11 @@ def apply_balanced_offload(sd_model, exclude=[]):
                 if do_offload:
                     module = module.to(devices.cpu, non_blocking=True)
                     used_gpu -= module_size
+                cls = module.__class__.__name__
+                quant = getattr(module, "quantization_method", None)
                 if not cached:
-                    shared.log.debug(f'Model module={module_name} type={module.__class__.__name__} dtype={module.dtype} quant={getattr(module, "quantization_method", None)} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}')
-                debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={getattr(module, "quantization_method", None)} module={module.__class__.__name__} size={module_size:.3f}')
+                    shared.log.debug(f'Model module={module_name} type={cls} dtype={module.dtype} quant={quant} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}')
+                debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
             except Exception as e:
                 if 'out of memory' in str(e):
                     devices.torch_gc(fast=True, force=True, reason='oom')
@@ -270,6 +271,8 @@ def apply_balanced_offload(sd_model, exclude=[]):
         apply_balanced_offload_to_module(sd_model.prior_pipe)
     if hasattr(sd_model, "decoder_pipe"):
         apply_balanced_offload_to_module(sd_model.decoder_pipe)
+    if shared.opts.layerwise_quantization:
+        model_quant.apply_layerwise(sd_model, quiet=True) # need to reapply since hooks were removed/readded
     set_accelerate(sd_model)
     t = time.time() - t0
     process_timer.add('offload', t)

modules/shared.py

@@ -564,14 +564,17 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}),
     "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}),
     "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}),
+
     "optimum_quanto_sep": OptionInfo("<h2>Optimum Quanto</h2>", "", gr.HTML),
     "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
     "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}),
     "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}),
+
     "torchao_sep": OptionInfo("<h2>TorchAO</h2>", "", gr.HTML),
     "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}),
     "torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}),
     "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ['int4_weight_only', 'int8_dynamic_activation_int4_weight', 'int8_weight_only', 'int8_dynamic_activation_int8_weight', 'float8_weight_only', 'float8_dynamic_activation_float8_weight', 'float8_static_activation_float8_weight'], "visible": native}),
+
     "nncf_sep": OptionInfo("<h2>NNCF</h2>", "", gr.HTML),
     "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
     "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}),
@@ -580,6 +583,11 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": cmd_opts.use_openvino}),
     "nncf_quant_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}),
     "quant_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}),
+
+    "layerwise_sep": OptionInfo("<h2>Layerwise Casting</h2>", "", gr.HTML),
+    "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "Text Encoder"], "visible": native}),
+    "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}),
+    "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}),
 }))
 
 options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {