File size: 6,843 Bytes
a3c20e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | # Adapted from: https://github.com/bghira/SimpleTuner
# With improvements from: https://github.com/ostris/ai-toolkit
from typing import Literal
import torch
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
from ltx_trainer import logger
QuantizationOptions = Literal[
"int8-quanto",
"int4-quanto",
"int2-quanto",
"fp8-quanto",
"fp8uz-quanto",
]
# Modules to exclude from quantization.
# These are glob patterns passed to quanto's `exclude` parameter.
# When quantizing the full model at once, these patterns match against full module paths.
# When quantizing block-by-block, we also use SKIP_ROOT_MODULES for top-level modules.
EXCLUDE_PATTERNS = [
# Input/output projection layers
"patchify_proj",
"audio_patchify_proj",
"proj_out",
"audio_proj_out",
# Timestep embedding layers - int4 tinygemm requires strict bfloat16 input
# and these receive float32 sinusoidal embeddings that are cast to bfloat16
"*adaln*",
"time_proj",
"timestep_embedder*",
# Caption/text projection layers
"caption_projection*",
"audio_caption_projection*",
# Normalization layers (usually excluded from quantization)
"*norm*",
]
# Top-level modules to skip entirely during block-by-block quantization.
# These are exact matches against model.named_children() names.
# (Needed because quanto's exclude patterns don't work when calling quantize() directly on a module)
SKIP_ROOT_MODULES = {
"patchify_proj",
"audio_patchify_proj",
"proj_out",
"audio_proj_out",
"audio_caption_projection",
}
def quantize_model(
model: torch.nn.Module,
precision: QuantizationOptions,
quantize_activations: bool = False,
device: torch.device | str | None = None,
) -> torch.nn.Module:
"""
Quantize a model using optimum-quanto.
For large models with transformer_blocks, this function quantizes block-by-block
on GPU then moves back to CPU, which is much faster than quantizing on CPU and
uses less peak VRAM than loading the entire model to GPU at once.
Args:
model: The model to quantize.
precision: The quantization precision (e.g. "int8-quanto", "fp8-quanto").
quantize_activations: Whether to quantize activations in addition to weights.
device: Device to use for quantization. If None, uses CUDA if available, else CPU.
Returns:
The quantized model.
"""
from optimum.quanto import freeze, quantize # noqa: PLC0415
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
elif isinstance(device, str):
device = torch.device(device)
weight_quant = _get_quanto_dtype(precision)
if quantize_activations:
logger.debug("Quantizing model weights and activations")
activations_quant = weight_quant
else:
activations_quant = None
# Remember original device to restore after quantization
original_device = next(model.parameters()).device
# Check if model has transformer_blocks for block-by-block quantization
if hasattr(model, "transformer_blocks"):
logger.debug("Quantizing model using block-by-block approach for memory efficiency")
_quantize_blockwise(
model,
weight_quant=weight_quant,
activations_quant=activations_quant,
device=device,
)
else:
# Fallback: quantize entire model at once
model.to(device)
quantize(model, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
freeze(model)
# Restore model to original device
model.to(original_device)
return model
def _quantize_blockwise(
model: torch.nn.Module,
weight_quant: torch.dtype,
activations_quant: torch.dtype | None,
device: torch.device,
) -> None:
"""Quantize a model block-by-block using optimum-quanto.
This approach:
1. Moves each transformer block to GPU
2. Quantizes on GPU (fast!)
3. Freezes the quantized weights
4. Moves back to CPU
This is much faster than quantizing on CPU and uses less peak VRAM
than loading the entire model to GPU.
"""
from optimum.quanto import freeze, quantize # noqa: PLC0415
original_dtype = next(model.parameters()).dtype
transformer_blocks = list(model.transformer_blocks)
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
transient=True,
) as progress:
task = progress.add_task("Quantizing transformer blocks", total=len(transformer_blocks))
for block in transformer_blocks:
# Move block to GPU
block.to(device, dtype=original_dtype, non_blocking=True)
# Quantize on GPU
quantize(block, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
freeze(block)
# Move back to CPU to free up VRAM for next block
block.to("cpu", non_blocking=True)
progress.advance(task)
# Quantize remaining non-transformer-block modules (e.g., embeddings, timestep projections)
# Skip modules that should not be quantized (patchify_proj, proj_out, etc.)
logger.debug("Quantizing remaining model components")
for name, module in model.named_children():
if name == "transformer_blocks":
continue # Already quantized
if name in SKIP_ROOT_MODULES:
logger.debug(f"Skipping quantization for module: {name}")
continue # Don't quantize these modules
# Move to device, quantize, freeze, move back
module.to(device, dtype=original_dtype, non_blocking=True)
quantize(module, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
freeze(module)
module.to("cpu", non_blocking=True)
def _get_quanto_dtype(precision: QuantizationOptions) -> torch.dtype:
"""Map precision string to quanto dtype."""
from optimum.quanto import ( # noqa: PLC0415
qfloat8,
qfloat8_e4m3fnuz,
qint2,
qint4,
qint8,
)
if precision == "int2-quanto":
return qint2
elif precision == "int4-quanto":
return qint4
elif precision == "int8-quanto":
return qint8
elif precision in ("fp8-quanto", "fp8uz-quanto"):
if torch.backends.mps.is_available():
raise ValueError("FP8 quantization is not supported on MPS devices. Use int2, int4, or int8 instead.")
if precision == "fp8-quanto":
return qfloat8
elif precision == "fp8uz-quanto":
return qfloat8_e4m3fnuz
raise ValueError(f"Invalid quantization precision: {precision}")
|