From 5a7c1f50c100bc1fb2ebf382ffa288687645fc7e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 23 Jan 2025 12:07:08 -0500 Subject: [PATCH] add native torch fp8 storage dtype Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 16 +++++++--- modules/model_quant.py | 72 ++++++++++++++++++++++++++++++++++-------- modules/sd_models.py | 5 ++- modules/sd_offload.py | 17 ++++++---- modules/shared.py | 8 +++++ 5 files changed, 91 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cd1efa24..6b1a291a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-21 +## Update for 2025-01-23 - **Contributing**: - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines @@ -32,10 +32,14 @@ - adjust residual diff threshold to balance the speedup and the accuracy: higher values leads to more cache hits and speedups, but might also lead to a higher accuracy drop - **IPEX** - - enable force attention slicing for all GPUs - - enable fp64 emulation and sycl jit cache - - switch to pytorch test branch on windows - - update the supported python versions + - enable force attention slicing, fp64 emulation, jit cache + - use pytorch test branch on windows + - extend the supported python versions +- **Torch FP8** + - uses torch `float8_e4m3fn` or `float8_e5m2` as data storage and performs dynamic upcasting to compute `dtype` as needed + - compatible with most `unet` and `transformer` models: e.g. *sd15, sdxl, sd35, flux.1, hunyuan-video, ltx-video, etc.* + this is alternative to `bnb`/`quanto`/`torchao` quantization on models/platforms/gpus where those libraries are not available + - enable in *settings -> quantization -> layerwise casting* - [PerFlow](https://github.com/magic-research/piecewise-rectified-flow) - piecewise rectified flow as model acceleration - use `perflow` scheduler combined with one of the available pre-trained [models](https://huggingface.co/hansyan) @@ -58,6 +62,8 @@ - unique font family registration - mochi video number of frames - mark large models that should offload + - avoid repeated optimum-quanto installation + - avoid reinstalling bnb if not cuda ## Update for 2025-01-15 diff --git a/modules/model_quant.py b/modules/model_quant.py index 0665fc920..0a56805df 100644 --- a/modules/model_quant.py +++ b/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}') diff --git a/modules/sd_models.py b/modules/sd_models.py index 6234c2b0b..0b98bb504 100644 --- a/modules/sd_models.py +++ b/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) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 8a490e110..aecb572ba 100644 --- a/modules/sd_offload.py +++ b/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) diff --git a/modules/shared.py b/modules/shared.py index 2200813ea..7f406d950 100644 --- a/modules/shared.py +++ b/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("

Optimum Quanto

", "", 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("

TorchAO

", "", 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("

NNCF

", "", 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("

Layerwise Casting

", "", 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"), {