ComputeNerd/ltx-2

1# Adapted from: https://github.com/bghira/SimpleTuner/blob/main/helpers/training/quantisation/__init__.py2from typing import Literal3 4import torch5from optimum.quanto import qtype6 7from ltx_trainer import logger8 9QuantizationOptions = Literal[10    "no_change",11    "int8-quanto",12    "int4-quanto",13    "int2-quanto",14    "fp8-quanto",15    "fp8uz-quanto",16]17 18 19def quantize_model(20    model: torch.nn.Module,21    precision: QuantizationOptions,22    quantize_activations: bool = False,23) -> torch.nn.Module:24    """25    Quantize a model using the specified precision settings.26 27    Args:28        model: The model to quantize.29        precision: The precision level to quantize to (e.g. "int8-quanto", "fp8-quanto").30        quantize_activations: Whether to quantize activations in addition to weights.31 32    Returns:33        The quantized model, or the original model if no quantization is performed.34    """35    if precision is None or precision == "no_change":36        return model37 38    from optimum.quanto import freeze, quantize  # noqa: PLC041539 40    weight_quant = _quanto_type_map(precision)41    extra_quanto_args = {42        "exclude": [43            "proj_in",44            "time_embed.*",45            "caption_projection.*",46            "rope",47            "*norm*",48            "proj_out",49        ]50    }51    if quantize_activations:52        logger.info("Freezing model weights and activations")53        extra_quanto_args["activations"] = weight_quant54    else:55        logger.info("Freezing model weights only")56 57    quantize(model, weights=weight_quant, **extra_quanto_args)58    freeze(model)59    return model60 61 62def _quanto_type_map(precision: QuantizationOptions) -> torch.dtype | qtype | None:  # noqa: PLR091163    if precision == "no_change":64        return None65 66    from optimum.quanto import (  # noqa: PLC041567        qfloat8,68        qfloat8_e4m3fnuz,69        qint2,70        qint4,71        qint8,72    )73 74    if precision == "int2-quanto":75        return qint276    elif precision == "int4-quanto":77        return qint478    elif precision == "int8-quanto":79        return qint880    elif precision in ("fp8-quanto", "fp8uz-quanto"):81        if torch.backends.mps.is_available():82            logger.warning(83                "MPS doesn't support dtype float8. "84                "you must select another precision level such as int2, int8, or int8.",85            )86            return None87        if precision == "fp8-quanto":88            return qfloat889        elif precision == "fp8uz-quanto":90            return qfloat8_e4m3fnuz91 92    raise ValueError(f"Invalid quantisation level: {precision}")93