Blame for toolkit/util/quantize.py

import os
import time
from fnmatch import fnmatch
from typing import List, Optional, Union, TYPE_CHECKING
import torch

from optimum.quanto.quantize import _quantize_submodule
from optimum.quanto.tensor import Optimizer, qtype, qtypes
from torchao.quantization.quant_api import (
    quantize_ as torchao_quantize_,
    Float8WeightOnlyConfig,
    Int8WeightOnlyConfig
)
from optimum.quanto import freeze
from tqdm import tqdm
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

from toolkit.print import print_acc
from toolkit.util.ostris_quant import (
    OstrisLinear,
    OstrisQuantizer,
    convert_linear_to_ostris,
    get_ostris_quantizer,
)
import os

if TYPE_CHECKING:
    from toolkit.models.base_model import BaseModel

# the quantize function in quanto had a bug where it was using exclude instead of include

Q_MODULES = [
    "QLinear",
    "QConv2d",
    "QEmbedding",
    "QBatchNorm2d",
    "QLayerNorm",
    "QConvTranspose2d",
    "QEmbeddingBag",
    "OstrisLinear",
]

torchao_qtypes = {
    # "int4": Int4WeightOnlyConfig(),
    # uint2..uint8 are handled by the UIntXQuantizer ostris backend
    # (toolkit/util/uintx_quant.py), a bit-exact reproduction of torchao 0.10.0's
    # UIntXWeightOnlyConfig, so ARAs stay byte-identical after torchao upgrades
    "int8": Int8WeightOnlyConfig(),
    "float8": Float8WeightOnlyConfig(),
}


class aotype:
    def __init__(self, name: str):
        self.name = name
        self.config = torchao_qtypes[name]


class ostristype:
    # custom quantization backend (see toolkit/util/ostris_quant.py), e.g. orbit2/3/4
    def __init__(self, name: str, quantizer: OstrisQuantizer):
        self.name = name
        self.quantizer = quantizer


def get_qtype(qtype: Union[str, qtype]) -> qtype:
    if qtype in torchao_qtypes:
        return aotype(qtype)
    if isinstance(qtype, str):
        ostris_quantizer = get_ostris_quantizer(qtype)
        if ostris_quantizer is not None:
            return ostristype(qtype, ostris_quantizer)
        return qtypes[qtype]
    else:
        return qtype


def is_quantized_tensor(t) -> bool:
    # torchao stores quantized weights as tensor subclasses (e.g. AffineQuantizedTensor) under torchao.*
    # that still report as nn.Parameter and expose .dequantize(). (quanto is handled separately.)
    # _is_ostris_weight tags two OstrisLinear tensors: the .weight property's eager tensor
    # (already dequantized; .dequantize() is a no-op) so the merge paths route through
    # requantize_module_weight, and the lazy OstrisLazyWeight emitted by state_dict()
    # (holds no data; .dequantize() materializes) so save loops dequantize it per key.
    if getattr(t, '_is_ostris_weight', False):
        return True
    return 'torchao' in type(t).__module__ and hasattr(t, 'dequantize')


def dequantize_if_quantized(t):
    return t.dequantize() if is_quantized_tensor(t) else t


def get_torchao_config(qtype):
    # returns the requantization config for a given qtype string (a torchao config, or the
    # ostristype for custom backends), or None if the qtype supports neither
    if qtype is None:
        return None
    try:
        q = get_qtype(qtype)
    except Exception:
        return None
    if isinstance(q, aotype):
        return q.config
    if isinstance(q, ostristype):
        return q
    return None


def requantize_module_weight(module, fp_weight, orig_dtype, config) -> None:
    """Write a full precision weight back into module.weight, re-quantizing in place if a
    requantization config is provided so the module stays quantized (used by the continuous
    merge/reset method). If config is None the weight is left in full precision."""
    if isinstance(module, OstrisLinear):
        # the module's backend reuses its existing quantization state; config is not needed
        module.requantize_(fp_weight)
        return
    if isinstance(config, ostristype):
        # custom backend config but the module was never converted (e.g. skipped at
        # quantize time); leave it in full precision
        config = None
    module.weight = torch.nn.Parameter(fp_weight.to(orig_dtype), requires_grad=False)
    if config is not None:
        torchao_quantize_(module, config)


def _wrap_qlinear_ndim(qlinear: torch.nn.Module) -> None:
    """quanto QLinear forward that tolerates >3D activations by flattening
    the leading dims for the mm and restoring them after (no-op otherwise)."""
    orig_forward = qlinear.forward

    def forward(x):
        if x.ndim > 3:
            lead = x.shape[:-1]
            out = orig_forward(x.reshape(-1, x.shape[-1]))
            return out.reshape(*lead, out.shape[-1])
        return orig_forward(x)

    qlinear.forward = forward
    qlinear._aitk_ndim_wrapped = True


def quantize(
    model: torch.nn.Module,
    weights: Optional[Union[str, qtype, aotype]] = None,
    activations: Optional[Union[str, qtype]] = None,
    optimizer: Optional[Optimizer] = None,
    include: Optional[Union[str, List[str]]] = None,
    exclude: Optional[Union[str, List[str]]] = None,
    quantize_device: Optional[torch.device] = None,
    keep_on_quantize_device: bool = False,
):
    """Quantize the specified model submodules

    Recursively quantize the submodules of the specified parent model.

    Only modules that have quantized counterparts will be quantized.

    If include patterns are specified, the submodule name must match one of them.

    If exclude patterns are specified, the submodule must not match one of them.

    Include or exclude patterns are Unix shell-style wildcards which are NOT regular expressions. See
    https://docs.python.org/3/library/fnmatch.html for more details.

    Note: quantization happens in-place and modifies the original model and its descendants.

    Args:
        model (`torch.nn.Module`): the model whose submodules will be quantized.
        weights (`Optional[Union[str, qtype]]`): the qtype for weights quantization.
        activations (`Optional[Union[str, qtype]]`): the qtype for activations quantization.
        include (`Optional[Union[str, List[str]]]`):
            Patterns constituting the allowlist. If provided, module names must match at
            least one pattern from the allowlist.
        exclude (`Optional[Union[str, List[str]]]`):
            Patterns constituting the denylist. If provided, module names must not match
            any patterns from the denylist.
        quantize_device (`Optional[torch.device]`):
            If provided, each module is moved to this device to quantize, then moved
            back to the device its weights were on initially. Lets a CPU-resident
            model (low vram) quantize layer-by-layer on the GPU.
        keep_on_quantize_device (`bool`):
            With quantize_device set: leave each layer on the quantize device after
            quantizing instead of moving it back. The gpu transient is then one
            bf16 layer at a time — never the whole unquantized remainder at once —
            for models whose final home is that device.
    """
    if include is not None:
        include = [include] if isinstance(include, str) else include
    if exclude is not None:
        exclude = [exclude] if isinstance(exclude, str) else exclude
    for name, m in model.named_modules():
        if include is not None and not any(
            fnmatch(name, pattern) for pattern in include
        ):
            continue
        if exclude is not None and any(fnmatch(name, pattern) for pattern in exclude):
            continue
        try:
            # check if m is QLinear or QConv2d
            if m.__class__.__name__ in Q_MODULES:
                # OstrisLinear may still be RE-quantized into a different
                # ostris qtype (same qtype is a per-layer no-op); every other
                # already-quantized module type is always left alone, which
                # also keeps quanto/torchao from double-quantizing
                # pre-quantized checkpoints
                if not (
                    isinstance(weights, ostristype)
                    and m.__class__.__name__ == "OstrisLinear"
                ):
                    continue
                if getattr(m.ostris_quantizer, "qtype", None) == weights.quantizer.qtype:
                    # guaranteed per-layer no-op: skip before any
                    # quantize_device round-trip
                    continue
            if isinstance(weights, ostristype) and not isinstance(m, torch.nn.Linear):
                # ostris backends only quantize nn.Linear; don't ferry norms/
                # embeddings across the bus for nothing when quantize_device
                # is set (containers fall through so children are visited)
                if quantize_device is not None and next(m.children(), None) is None:
                    continue
            if (
                isinstance(weights, aotype)
                and not isinstance(m, torch.nn.Linear)
                and (
                    quantize_device is not None
                    or include is not None
                    or exclude is not None
                )
            ):
                # torchao only quantizes nn.Linear; when a device round-trip or
                # include/exclude filtering is in play, skip containers so each
                # linear is handled individually (a container-level torchao call
                # would quantize excluded children too)
                continue
            orig_device = None
            if quantize_device is not None and next(m.children(), None) is None:
                # OstrisLinear layers being re-quantized hold buffers, not params
                param = next(m.parameters(recurse=False), None)
                if param is None:
                    param = next(m.buffers(recurse=False), None)
                if param is not None:
                    orig_device = param.device
                    m.to(quantize_device)
            try:
                if isinstance(weights, ostristype):
                    if isinstance(m, torch.nn.Linear):
                        convert_linear_to_ostris(m, weights.quantizer)
                elif isinstance(weights, aotype):
                    torchao_quantize_(m, weights.config)
                else:
                    _quantize_submodule(
                        model,
                        name,
                        m,
                        weights=weights,
                        activations=activations,
                        optimizer=optimizer,
                    )
                    # quanto's qbytes_mm only takes 2D/3D activations; video
                    # patch embeds feed their linears >3D tensors (convrot/
                    # torchao reshape internally). Flatten around the QLinear.
                    replaced = model.get_submodule(name)
                    if replaced.__class__.__name__ == "QLinear" and not getattr(
                        replaced, "_aitk_ndim_wrapped", False
                    ):
                        _wrap_qlinear_ndim(replaced)
            finally:
                if orig_device is not None and not keep_on_quantize_device:
                    # quanto replaces the module in its parent, so re-fetch by name
                    model.get_submodule(name).to(orig_device)
        except Exception as e:
            print(f"Failed to quantize {name}: {e}")
            # raise e


def _has_quantizable_linear(module: torch.nn.Module, weights, exclude=None) -> bool:
    """Whether quantizing ``module`` with ``weights`` would change anything.

    False when every non-excluded linear is already quantized with the same
    ostris qtype (or cannot be quantized at all) — the pre-quantized-checkpoint
    case, where the whole block can be skipped without the device round-trip.
    ``exclude`` patterns are matched against module names relative to
    ``module`` (use leading wildcards for patterns aimed at inner layers)."""
    if not isinstance(weights, ostristype):
        return True
    for name, m in module.named_modules():
        if not isinstance(m, torch.nn.Linear):
            continue
        if exclude is not None and any(fnmatch(name, pattern) for pattern in exclude):
            continue
        if isinstance(m, OstrisLinear):
            if getattr(
                m.ostris_quantizer, "qtype", None
            ) != weights.quantizer.qtype and weights.quantizer.can_quantize(m):
                return True
            continue
        if m.__class__.__name__ in Q_MODULES:
            continue
        if weights.quantizer.can_quantize(m):
            return True
    return False


@torch.no_grad()
def dequantize_ostris_to_linear(module: torch.nn.Module) -> int:
    """Replace every OstrisLinear with a plain nn.Linear holding the full-
    precision weight (activation-side transforms folded), in place, layer by
    layer — the full-precision transient never exceeds one layer. Used when a
    pre-quantized checkpoint is loaded but a DIFFERENT quantization (or none,
    e.g. full finetuning) was requested. Returns the number of layers
    restored."""
    replaced = 0
    for parent in module.modules():
        for child_name, child in list(parent.named_children()):
            if not isinstance(child, OstrisLinear):
                continue
            weight = child.ostris_quantizer.dequantize_folded(child).to(
                child.ostris_orig_dtype
            )
            new = torch.nn.Linear(
                child.in_features,
                child.out_features,
                bias=child.bias is not None,
                device="meta",
                dtype=weight.dtype,
            )
            new.weight = torch.nn.Parameter(weight)
            if child.bias is not None:
                new.bias = torch.nn.Parameter(
                    child.bias.data.to(weight.device, weight.dtype)
                )
            setattr(parent, child_name, new)
            replaced += 1
    return replaced


@torch.no_grad()
def quantize_module(
    module: torch.nn.Module,
    qtype: str,
    device=None,
    dtype: torch.dtype = torch.bfloat16,
    block_names: Optional[List[str]] = None,
    exclude: Optional[List[str]] = None,
    quantize_kwargs: Optional[dict] = None,
    status_fn=print_acc,
    keep_on_device: bool = False,
):
    """Module-centric quantization: block-streamed (each repeated block moves
    to ``device`` for the math) with a whole-module pass for the extras. This
    is the core the per-model loaders call; holders' quantize_model wraps it
    with model_config plumbing.

    ``keep_on_device``: when the model's final home IS ``device`` (no layer
    offloading / low_vram), quantized blocks stay there instead of round-
    tripping back to cpu — the weights then cross the bus exactly once
    (mmap/page-cache -> gpu) instead of three times (up, back into a fresh
    host copy, up again at final placement), and the model-sized host RAM
    spike of that intermediate copy never happens. The extras pass runs on
    the gpu too. With it off (offload paths), blocks return to cpu as before
    and the extras quantize layer-by-layer on ``device`` via quantize_device
    rather than burning every cpu core."""
    from toolkit.dequantize import patch_dequantization_on_save

    patch_dequantization_on_save(module)
    quantization_type = get_qtype(qtype)
    exclude = list(exclude or [])
    quantize_kwargs = quantize_kwargs or {}
    keep_on_device = keep_on_device and device is not None

    all_blocks: List[torch.nn.Module] = []
    for name in block_names or []:
        # name may be a dotted path for models that nest their blocks
        block_list = module
        for part in name.split("."):
            block_list = getattr(block_list, part, None)
            if block_list is None:
                break
        if block_list is not None:
            all_blocks += list(block_list)
    if all_blocks:
        status_fn(f" - quantizing {len(all_blocks)} blocks")
    already_quantized = 0
    debug_phases = os.environ.get("AITK_QUANT_DEBUG") == "1"
    t_check = t_h2d = t_quant = t_d2h = 0.0
    for block in tqdm(all_blocks):
        t = time.perf_counter()
        skip = not _has_quantizable_linear(block, quantization_type, exclude)
        t_check += time.perf_counter() - t
        if skip:
            # pre-quantized checkpoint with a matching qtype: nothing in this
            # block would change — skip the dtype cast entirely so the load
            # stays byte-identical (placement still honors keep_on_device)
            already_quantized += 1
            if keep_on_device:
                block.to(device)
            continue
        t = time.perf_counter()
        if device is not None:
            block.to(device, dtype=dtype, non_blocking=True)
        t_h2d += time.perf_counter() - t
        t = time.perf_counter()
        quantize(block, weights=quantization_type, exclude=exclude, **quantize_kwargs)
        freeze(block)
        t_quant += time.perf_counter() - t
        # NOT non_blocking: an async D2H allocates the cpu destination in pinned
        # memory, which the caching host allocator keeps forever — that silently
        # retained a model-sized chunk of host ram
        t = time.perf_counter()
        if device is not None and not keep_on_device:
            block.to("cpu")
        t_d2h += time.perf_counter() - t
    if debug_phases and all_blocks:
        status_fn(
            f" - [debug] check {t_check:.1f}s h2d {t_h2d:.1f}s "
            f"quant {t_quant:.1f}s d2h {t_d2h:.1f}s"
        )
    if already_quantized:
        status_fn(
            f" - {already_quantized} blocks already quantized with a matching qtype; left untouched"
        )

    status_fn(" - quantizing extras")
    if keep_on_device:
        # blocks already live on the gpu; quantize each remaining layer there
        # one at a time and leave it — the transient is one bf16 layer, never
        # the whole unquantized remainder at once (that transient bump could
        # exceed the model's final footprint by GBs on big-extras models)
        quantize(
            module,
            weights=quantization_type,
            exclude=exclude,
            quantize_device=device,
            keep_on_quantize_device=True,
            **quantize_kwargs,
        )
        # non-quantized leftovers (norms, embeddings, buffers) follow — final
        # residency, not a transient
        module.to(device)
    else:
        # cpu-resident model: quantize each extra layer with a gpu round-trip
        quantize(
            module,
            weights=quantization_type,
            exclude=exclude,
            quantize_device=device,
            **quantize_kwargs,
        )
    freeze(module)
    return module


def quantize_model(
    base_model: "BaseModel",
    model_to_quantize: torch.nn.Module,
):
    from toolkit.dequantize import patch_dequantization_on_save

    if not hasattr(base_model, "get_transformer_block_names"):
        raise ValueError(
            "The model to quantize must have a method `get_transformer_block_names`."
        )

    # patch the state dict method
    patch_dequantization_on_save(model_to_quantize)

    # sensitive modules to keep in full precision (fnmatch patterns)
    exclude_modules = base_model.get_quantization_exclude_modules() or []

    mc = base_model.model_config
    device = base_model.device_torch
    keep_on_device = (
        not mc.low_vram
        and not (mc.layer_offloading and mc.layer_offloading_transformer_percent > 0)
        and torch.device(device).type != "cpu"
    )

    if mc.accuracy_recovery_adapter is not None:
        attach_ara_and_quantize(
            base_model,
            model_to_quantize,
            ara_path=mc.accuracy_recovery_adapter,
            exclude=exclude_modules,
            device=device,
            keep_on_device=keep_on_device,
        )
    else:
        quantize_module(
            model_to_quantize,
            mc.qtype,
            device=device,
            dtype=base_model.torch_dtype,
            block_names=base_model.get_transformer_block_names(),
            exclude=exclude_modules,
            quantize_kwargs=mc.quantize_kwargs,
            status_fn=base_model.print_and_status_update,
            keep_on_device=keep_on_device,
        )


@torch.no_grad()
def attach_ara_and_quantize(
    base_model: "BaseModel",
    model_to_quantize: torch.nn.Module,
    ara_path: str,
    exclude: Optional[List[str]] = None,
    device=None,
    keep_on_device: bool = False,
):
    """Load an accuracy recovery adapter as a live network on the module and
    quantize around it (adapter-hijacked linears at the configured qtype,
    everything else uint8). The network lands on
    base_model.accuracy_recovery_adapter.

    ``device``: quantize each hijacked linear there (gpu kernels) instead of
    wherever it happens to live (historically the cpu — slow). With
    ``keep_on_device`` the quantized linears stay there (final home is that
    gpu); otherwise each returns to the device it came from."""
    from toolkit.config_modules import NetworkConfig
    from toolkit.lora_special import LoRASpecialNetwork

    exclude_modules = list(exclude or [])
    load_lora_path = ara_path

    if not os.path.exists(load_lora_path):
        # not local file, grab from the hub

        path_split = load_lora_path.split("/")
        if len(path_split) > 3:
            raise ValueError(
                "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'."
            )
        repo_id = f"{path_split[0]}/{path_split[1]}"
        print_acc(f"Grabbing lora from the hub: {load_lora_path}")
        new_lora_path = hf_hub_download(
            repo_id,
            filename=path_split[-1],
        )
        # replace the path
        load_lora_path = new_lora_path

    # build the lora config based on the lora weights
    lora_state_dict = load_file(load_lora_path)
        
    if hasattr(base_model, "convert_lora_weights_before_load"):
        lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict)
        
    network_config = {
        "type": "lora",
        "network_kwargs": {"only_if_contains": []},
        "transformer_only": False,
    }
    first_key = list(lora_state_dict.keys())[0]
    first_weight = lora_state_dict[first_key]
    # if it starts with lycoris and includes lokr
    if first_key.startswith("lycoris") and any(
        "lokr" in key for key in lora_state_dict.keys()
    ):
        network_config["type"] = "lokr"
        
    network_kwargs = {}

    # find firse loraA weight
    if network_config["type"] == "lora":
        linear_dim = None
        for key, value in lora_state_dict.items():
            if "lora_A" in key:
                linear_dim = int(value.shape[0])
                break
        linear_alpha = linear_dim
        network_config["linear"] = linear_dim
        network_config["linear_alpha"] = linear_alpha

        # we build the keys to match every key
        only_if_contains = []
        for key in lora_state_dict.keys():
            contains_key = key.split(".lora_")[0]
            if contains_key not in only_if_contains:
                only_if_contains.append(contains_key)

        network_kwargs["only_if_contains"] = only_if_contains
    elif network_config["type"] == "lokr":
        # find the factor
        largest_factor = 0
        for key, value in lora_state_dict.items():
            if "lokr_w1" in key:
                factor = int(value.shape[0])
                if factor > largest_factor:
                    largest_factor = factor
        network_config["lokr_full_rank"] = True
        network_config["lokr_factor"] = largest_factor

        only_if_contains = []
        for key in lora_state_dict.keys():
            if "lokr_w1" in key:
                contains_key = key.split(".lokr_w1")[0]
                contains_key = contains_key.replace("lycoris_", "")
                if contains_key not in only_if_contains:
                    only_if_contains.append(contains_key)
        network_kwargs["only_if_contains"] = only_if_contains
        
    if hasattr(base_model, 'target_lora_modules'):
        network_kwargs['target_lin_modules'] = base_model.target_lora_modules

    # todo auto grab these
    # get dim and scale
    network_config = NetworkConfig(**network_config)

    network = LoRASpecialNetwork(
        text_encoder=None,
        unet=model_to_quantize,
        lora_dim=network_config.linear,
        multiplier=1.0,
        alpha=network_config.linear_alpha,
        # conv_lora_dim=self.network_config.conv,
        # conv_alpha=self.network_config.conv_alpha,
        train_unet=True,
        train_text_encoder=False,
        network_config=network_config,
        network_type=network_config.type,
        transformer_only=network_config.transformer_only,
        is_transformer=base_model.is_transformer,
        base_model=base_model,
        is_ara=True,
        **network_kwargs
    )
    network.apply_to(
        None, model_to_quantize, apply_text_encoder=False, apply_unet=True
    )
    network.force_to(base_model.device_torch, dtype=base_model.torch_dtype)
    network._update_torch_multiplier()
    network.load_weights(lora_state_dict)
    network.eval()
    network.is_active = True
    network.can_merge_in = False
    base_model.accuracy_recovery_adapter = network

    # quantize it
    keep_on_device = keep_on_device and device is not None
    lora_exclude_modules = []
    quantization_type = get_qtype(base_model.model_config.qtype)
    for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"):
        # the lora has already hijacked the original module
        orig_module = lora_module.org_module[0]
        orig_device = None
        if device is not None:
            param = next(orig_module.parameters(), None)
            orig_device = param.device if param is not None else None
            orig_module.to(device, dtype=base_model.torch_dtype)
        else:
            orig_module.to(base_model.torch_dtype)
        # make the params not require gradients
        for param in orig_module.parameters():
            param.requires_grad = False
        quantize(orig_module, weights=quantization_type)
        freeze(orig_module)
        module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '')
        lora_exclude_modules.append(module_name)
        if not keep_on_device and orig_device is not None:
            orig_module.to(orig_device)
        elif base_model.model_config.low_vram and device is None:
            # legacy behavior when no quantize device was given
            orig_module.to("cpu")
    # quantize additional layers
    print_acc(" - quantizing additional layers")
    quantization_type = get_qtype('uint8')
    quantize(
        model_to_quantize,
        weights=quantization_type,
        exclude=lora_exclude_modules + exclude_modules,
        quantize_device=device,
        keep_on_quantize_device=keep_on_device,
    )
    if keep_on_device:
        # non-quantized leftovers follow — final residency, not a transient
        model_to_quantize.to(device)