From 1f24513507dc04da7d7a3eb59312e5f3416d3a6e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 7 Jun 2026 11:43:46 +0200 Subject: [PATCH] optimize sdnq quant-on-load and add platform stats Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 9 +- modules/model_quant.py | 6 +- modules/platform.py | 19 ++++ modules/platform_linux.py | 168 ++++++++++++++++++++++++++++++ modules/sd_models.py | 4 + modules/sd_offload.py | 47 +++++++-- modules/sdnq/loader.py | 2 +- modules/sdnq/quant_utils.py | 2 +- modules/sdnq/quantizer.py | 19 ++-- pipelines/generic_text_encoder.py | 8 +- pipelines/generic_transformer.py | 10 +- pipelines/model_ideogram4.py | 1 + wiki | 2 +- 13 files changed, 269 insertions(+), 28 deletions(-) create mode 100644 modules/platform.py create mode 100644 modules/platform_linux.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d9bc2df5..289a33395 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,8 @@ # Change Log for SD.Next -## Update for 2026-06-06 +## Update for 2026-06-07 -### Highlights for 2026-06-06 +### Highlights for 2026-06-07 *What's New?* - **Ideogram-4** released, Microsoft joins the game with **Lens** and **Anima** made it to release version @@ -19,7 +19,7 @@ And we have a new modular LoRA loader, new native Transformers loader and improv [ReadMe](https://github.com/vladmandic/automatic/blob/master/README.md) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) | [Sponsor](https://github.com/sponsors/vladmandic) -### Details for 2026-06-06 +### Details for 2026-06-07 - **Models** - [CircleStone Anima 1.0](https://huggingface.co/circlestone-labs/Anima) in *Base* and *Turbo* (distilled) variants @@ -137,6 +137,9 @@ And we have a new modular LoRA loader, new native Transformers loader and improv - `output path` use correct base folder for initial folders - `ltx` prompt embeds move to device, thanks @ryanmeador - `openpose` processor + - `img2img` api default sampler + - `sdnq` default dynamic loss value + - `samplers` ui sigma methods ## Update for 2026-05-13 diff --git a/modules/model_quant.py b/modules/model_quant.py index ab1aa355f..baacde86e 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -492,8 +492,7 @@ def get_dit_args(load_config: dict | None = None, module: str | None = None, dev config = {} if load_config is None else load_config.copy() if 'torch_dtype' not in config: config['torch_dtype'] = devices.dtype - if 'low_cpu_mem_usage' in config: - del config['low_cpu_mem_usage'] + low_cpu = config.get('low_cpu_mem_usage', False) if 'load_connected_pipeline' in config: del config['load_connected_pipeline'] if 'safety_checker' in config: @@ -509,6 +508,9 @@ def get_dit_args(load_config: dict | None = None, module: str | None = None, dev config['device_map'] = 'cpu' elif shared.opts.device_map == 'gpu': config['device_map'] = devices.device + elif low_cpu and module in {'Model', 'TE', 'LLM'}: + # Quantized transformer/text encoder loads should default to cpu device_map when low_cpu_mem_usage is requested to avoid full in-memory checkpoint expansion + config['device_map'] = 'cpu' if allow_quant: quant_args = create_config(module=module, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) else: diff --git a/modules/platform.py b/modules/platform.py new file mode 100644 index 000000000..83848ec40 --- /dev/null +++ b/modules/platform.py @@ -0,0 +1,19 @@ +import os +import platform +from modules.logger import log + + +def cleanup(): + if os.environ.get('SD_PLATFORM_DEBUG', None) is None: + return + if platform.system() == "Linux": + log.warning(f'Platform: {platform.system()} cleanup') + from modules.platform_linux import LinuxUtils + LinuxUtils.advise_mmap() + LinuxUtils.release_mmap() + LinuxUtils.advise_cache() + LinuxUtils.malloc_trim() + LinuxUtils.get_smaps() + LinuxUtils.get_status() + else: + log.warning(f'Platform: {platform.system()} not supported') diff --git a/modules/platform_linux.py b/modules/platform_linux.py new file mode 100644 index 000000000..793f40a58 --- /dev/null +++ b/modules/platform_linux.py @@ -0,0 +1,168 @@ +import os +import ctypes +from modules.logger import log + + +class LinuxUtils(): + @staticmethod + def get_status() -> dict[str, float] | None: + lines = [] + status = {} + try: + with open("/proc/self/status", encoding="utf-8") as handle: + lines = handle.readlines() + except OSError: + return status + for line in lines: + key, _sep, value = line.partition(":") + parts = value.strip().split() + if not parts: + continue + try: + status[key] = parts[0] + except ValueError: + continue + log.debug(f'Linux status: {status}') + return status + + @staticmethod + def get_smaps(limit: int = 8) -> list[dict[str, float | str]] | None: + try: + with open("/proc/self/smaps", encoding="utf-8") as handle: + lines = handle.readlines() + except OSError: + return None + entries = [] + current = None + for raw_line in lines: + line = raw_line.rstrip() + if not line: + continue + if "-" in line and line[:1].isalnum() and line.split(maxsplit=1)[0].count("-") == 1: + if current is not None: + entries.append(current) + parts = line.split(maxsplit=5) + current = { + "path": parts[5] if len(parts) > 5 else "[anonymous]", + "rss": 0, + "pss": 0, + "private": 0, + "shared": 0, + } + continue + if current is None or ":" not in line: + continue + key, value = line.split(":", maxsplit=1) + value = value.strip().split() + if not value: + continue + try: + amount = int(value[0]) + except ValueError: + continue + if key == "Rss": + current["rss"] += amount + elif key == "Pss": + current["pss"] += amount + elif key in {"Private_Clean", "Private_Dirty"}: + current["private"] += amount + elif key in {"Shared_Clean", "Shared_Dirty"}: + current["shared"] += amount + if current is not None: + entries.append(current) + merged = {} + for entry in entries: + path = entry["path"] + if path not in merged: + merged[path] = entry.copy() + else: + merged[path]["rss"] += entry["rss"] + merged[path]["pss"] += entry["pss"] + merged[path]["private"] += entry["private"] + merged[path]["shared"] += entry["shared"] + top = sorted(merged.values(), key=lambda item: item["rss"], reverse=True)[:limit] + for entry in top: + entry["rss"] = round(entry["rss"] / 1024 / 1024, 3) + entry["pss"] = round(entry["pss"] / 1024 / 1024, 3) + entry["private"] = round(entry["private"] / 1024 / 1024, 3) + entry["shared"] = round(entry["shared"] / 1024 / 1024, 3) + log.debug(f'Linux smaps: top={top}') + return top + + @staticmethod + def malloc_trim() -> bool | None: + try: + libc = ctypes.CDLL("libc.so.6") + libc.malloc_trim.argtypes = [ctypes.c_size_t] + libc.malloc_trim.restype = ctypes.c_int + status = bool(libc.malloc_trim(0)) + log.debug(f"Linux trim: status={status}") + except (AttributeError, OSError): + log.debug("Linux trim: not supported") + + @staticmethod + def advise_mmap(): + """Mark mmaps as temporary so OS prioritizes dropping them.""" + MADV_COLD = 5 # Linux 5.4+, mark as unlikely to be used + libc = ctypes.CDLL('libc.so.6') + advised = 0 + with open('/proc/self/maps', 'r', encoding='utf-8') as f: + for line in f: + if 'blobs' in line or '/dev/zero' in line: + try: + addr, size = line.split()[0].split('-') + addr = int(addr, 16) + size = int(size, 16) - addr + libc.madvise(ctypes.c_void_p(addr), size, MADV_COLD) + advised += 1 + except Exception: + log.error(f"Linux mmap advise: {line.strip()}") + log.debug(f"Linux mmap advise: num={advised}") + + @staticmethod + def release_mmap(): + """Use madvise to drop safetensors blob mmaps from page cache.""" + try: + libc = ctypes.CDLL('libc.so.6') + # Get all memory mappings for this process + dropped = [] + with open('/proc/self/maps', 'r', encoding='utf-8') as f: + for line in f: + parts = line.split() + if len(parts) >= 6: + path = parts[5] + if 'blobs' in path or '/dev/zero' in path: + if path in dropped: + continue + try: + addr, size = line.split()[0].split('-') + addr = int(addr, 16) + size = int(size, 16) - addr + if libc.madvise(ctypes.c_void_p(addr), size, 4) == 0: + dropped.append(path) + except Exception: + log.error(f"Linux mmap release: {line.strip()}") + log.debug(f"Linux mmap release: {dropped}") + except Exception as e: + log.error(f"Linux mmap release: {e}") + + + @staticmethod + def advise_cache(): + """Advise OS to drop cache for safetensors blobs.""" + from modules.shared import opts + try: + if hasattr(os, 'posix_fadvise') and hasattr(os, 'POSIX_FADV_DONTNEED'): + for root, _dirs, files in os.walk(opts.hfcache_dir, topdown=False): + for f in files: + if f.startswith(('blobs', 'snapshots')): + try: + path = os.path.join(root, f) + fd = os.open(path, os.O_RDONLY | os.O_NONBLOCK) + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + os.close(fd) + except Exception: + log.error(f"Linux cache: {path}") + log.debug("Linux cache: advised") + except Exception as e: + log.error(f"Linux cache: {e}") diff --git a/modules/sd_models.py b/modules/sd_models.py index 5c5444104..8771cb58e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1033,6 +1033,10 @@ def load_diffuser(checkpoint_info: CheckpointInfo | None = None, op='model', rev modelstats.analyze() log.info(f"Load {op}: family={shared.sd_model_type} time={timer.load.dct()} native={get_native(sd_model)} memory={memory_stats()}") + + from modules.platform import cleanup + cleanup() + shared.opts.save(silent=True) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 3dd931b1c..6fc50625f 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -3,6 +3,7 @@ import re import sys import time import inspect +import itertools import torch import accelerate.hooks import accelerate.utils.modeling @@ -365,6 +366,43 @@ def get_module_names(pipe=None, exclude=None): return modules_names +def get_module_memory(module: torch.nn.Module) -> dict[str, float]: + tensors = list(itertools.chain(module.parameters(), module.buffers())) + logical_gib = sum(tensor.numel() * tensor.element_size() for tensor in tensors) / 1024**3 + storages = {} + for tensor in tensors: + try: + storage = tensor.untyped_storage() + except (AttributeError, RuntimeError): + continue + storages[(storage.data_ptr(), storage.nbytes())] = storage.nbytes() + storage_gib = sum(storages.values()) / 1024**3 + return { + "logical": round(logical_gib, 3), + "storage": round(storage_gib, 3), + "overhead": round(storage_gib - logical_gib, 3), + "tensors": len(tensors), + "storages": len(storages), + } + + +def get_module_size(module: torch.nn.Module) -> tuple[float, float]: + module_size = 0 + param_num = 0 + if not isinstance(module, torch.nn.Module): + return 0, 0 + try: + # module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + tensors = set(itertools.chain(module.parameters(recurse=True), module.buffers(recurse=True))) + module_size = sum(t.numel() * t.element_size() for t in tensors) / 1024**3 + param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + except Exception as e: + log.error(f'Offload: type=balanced op=calc module={module.__class__.__name__} {e}') + module_size = 0 + param_num = 0 + return module_size, param_num + + def get_module_sizes(pipe=None, exclude=None): if exclude is None: exclude = [] @@ -373,14 +411,7 @@ def get_module_sizes(pipe=None, exclude=None): module_size = offload_hook_instance.offload_map.get(module_name, None) if module_size is None: module = getattr(pipe, module_name, None) - if not isinstance(module, torch.nn.Module): - continue - try: - module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 - param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 - except Exception as e: - log.error(f'Offload: type=balanced op=calc module={module_name} {e}') - module_size = 0 + module_size, param_num = get_module_size(module) offload_hook_instance.offload_map[module_name] = module_size offload_hook_instance.param_map[module_name] = param_num modules[module_name] = module_size diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 3ccb40ad5..90a7a5363 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -146,7 +146,7 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin | None = None, file_n files.append(os.path.join(model_path, file_name)) else: all_files = os.listdir(model_path) - files = sorted([os.path.join(model_path, f) for f in all_files if f.endswith(".safetensors")]) + files = sorted([os.path.join(model_path, f) for f in all_files if f.endswith(".safetensors")]) # pylint: disable=not-an-iterable state_dict = load_files(files, key_mapping=key_mapping, device=device, method=load_method) diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index 60d91707d..c702831f4 100644 --- a/modules/sdnq/quant_utils.py +++ b/modules/sdnq/quant_utils.py @@ -24,7 +24,7 @@ def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: int | list[in @devices.inference_context() def quantize_weight(weight: torch.FloatTensor, reduction_axes: int | list[int], weights_dtype: str, dtype: torch.dtype = None, use_stochastic_rounding: bool = False) -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: if weight.dtype != torch.float64: - weight = weight.to(dtype=torch.float32) + weight = weight.to(dtype=torch.float32, copy=False) if dtype_dict[weights_dtype]["is_unsigned"]: scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index fc223c39a..3392331a7 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -235,8 +235,9 @@ def sdnq_quantize_layer_weight_dynamic(weight, layer_class_name=None, weights_dt dynamic_loss_threshold = 10 ** -(dtype_dict[weights_dtype]["num_bits"] / 2) if weight.dtype != torch.float64: - weight = weight.to(dtype=torch.float32) - original_weight_fp32 = weight.clone() + weight = weight.to(dtype=torch.float32, copy=False) + weight = weight.detach() + original_weight_fp32 = weight.clone() if use_svd else weight weight_std = original_weight_fp32.std().square_().clamp_(min=1e-8) if use_hadamard: @@ -337,7 +338,7 @@ def sdnq_quantize_layer(layer, quantization_config: "SDNQConfig", torch_dtype: t if return_device is None: return_device = layer.weight.device if quantization_device is not None: - layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking) + layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking, copy=False) if use_dynamic_quantization: weight_data = sdnq_quantize_layer_weight_dynamic(layer.weight, **quant_kwargs) @@ -350,7 +351,7 @@ def sdnq_quantize_layer(layer, quantization_config: "SDNQConfig", torch_dtype: t for key, value in weight_data.items(): if isinstance(value, (torch.Tensor, torch.nn.Parameter)): - setattr(layer, key, torch.nn.Parameter(value.to(return_device, non_blocking=non_blocking), requires_grad=False)) + setattr(layer, key, torch.nn.Parameter(value.to(return_device, non_blocking=non_blocking, copy=False), requires_grad=False)) setattr(getattr(layer, key), "_is_hf_initialized", True) # noqa: B010 else: setattr(layer, key, value) @@ -365,7 +366,7 @@ def sdnq_quantize_layer(layer, quantization_config: "SDNQConfig", torch_dtype: t if quant_kwargs["use_quantized_matmul"] and not layer.sdnq_dequantizer.use_quantized_matmul: quantization_config.modules_to_not_use_matmul.append(param_name) else: - layer.weight = torch.nn.Parameter(layer.weight.to(return_device, dtype=torch_dtype, non_blocking=non_blocking), requires_grad=False) + layer.weight = torch.nn.Parameter(layer.weight.to(return_device, dtype=torch_dtype, non_blocking=non_blocking, copy=False), requires_grad=False) if use_dynamic_quantization: quantization_config.modules_to_not_convert.append(param_name) @@ -575,9 +576,9 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): return_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype) if param_value.dtype == return_dtype and devices.same_device(param_value.device, target_device): - param_value = param_value.clone() + param_value = param_value.detach() else: - param_value = param_value.to(target_device, dtype=return_dtype) + param_value = param_value.to(target_device, dtype=return_dtype, copy=False) if tensor_name == "weight" and layer.sdnq_dequantizer.use_quantized_matmul and not layer.sdnq_dequantizer.re_quantize_for_matmul: param_value = prepare_weight_for_matmul(param_value) @@ -600,9 +601,9 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): quant_kwargs["quantization_device"] = None if param_value.dtype in {torch.float32, torch.float64} and devices.same_device(param_value.device, target_device): - param_value = param_value.clone() + param_value = param_value.detach() else: - param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking).to(dtype=torch.float32 if param_value.dtype != torch.float64 else torch.float64) + param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking, copy=False).to(dtype=torch.float32 if param_value.dtype != torch.float64 else torch.float64) layer.weight = torch.nn.Parameter(param_value, requires_grad=False) layer, self.quantization_config = sdnq_quantize_layer(layer, self.quantization_config, torch_dtype=torch_dtype, param_name=param_name, quant_kwargs=quant_kwargs) # pylint: disable=attribute-defined-outside-init diff --git a/pipelines/generic_text_encoder.py b/pipelines/generic_text_encoder.py index 1687969be..2bb94e58d 100644 --- a/pipelines/generic_text_encoder.py +++ b/pipelines/generic_text_encoder.py @@ -1,7 +1,7 @@ import os import json import transformers -from modules import shared, devices, errors, sd_models, model_quant +from modules import shared, devices, errors, sd_models, sd_offload, model_quant from modules.logger import log from pipelines.generic_util import get_loader from pipelines.generic_shared import shared_te_map @@ -141,4 +141,10 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod devices.torch_gc() shared.state.end(jobid) + + if text_encoder is not None: + module_size, param_num = sd_offload.get_module_size(text_encoder) + module_memory = sd_offload.get_module_memory(text_encoder) + log.debug(f'Load model: text_encoder="{repo_id}" quant="{quant_type}" size={module_size:.3f} params={param_num:.3f} memory={module_memory}') + return text_encoder diff --git a/pipelines/generic_transformer.py b/pipelines/generic_transformer.py index eec297b7f..f7ee42c8d 100644 --- a/pipelines/generic_transformer.py +++ b/pipelines/generic_transformer.py @@ -1,5 +1,5 @@ import os -from modules import shared, devices, errors, sd_models, model_quant +from modules import shared, devices, errors, sd_models, sd_offload, model_quant from modules.logger import log from pipelines.generic_util import get_loader @@ -34,7 +34,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer def load_from_repo(): nonlocal quant_args - log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}') + log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} loader={get_loader("diffusers")} args={load_args}') if 'sdnq-' in repo_id.lower(): quant_args = {} if dtype is not None: @@ -127,4 +127,10 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer devices.torch_gc() shared.state.end(jobid) + + if transformer is not None: + module_size, param_num = sd_offload.get_module_size(transformer) + module_memory = sd_offload.get_module_memory(transformer) + log.debug(f'Load model: transformer="{repo_id}" quant="{quant_type}" size={module_size:.3f} params={param_num:.3f} memory={module_memory}') + return transformer diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py index ad4c102f0..5c57e61b6 100644 --- a/pipelines/model_ideogram4.py +++ b/pipelines/model_ideogram4.py @@ -40,6 +40,7 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): return None transformer_cls = diffusers.Ideogram4Transformer2DModel + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="transformer", load_config=diffusers_load_config) if shared.opts.model_ideogram4_enable_cg: unconditional_transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="unconditional_transformer", load_config=diffusers_load_config) diff --git a/wiki b/wiki index 974703655..7526ee878 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 974703655102502cb5260d292916fbf36043add8 +Subproject commit 7526ee87821e8cf69e3a1e4e26f5ecc7d3508cee