| 1 | from typing import Callable |
| 2 | |
| 3 | import torch |
| 4 | from torch import nn |
| 5 | |
| 6 | from ltx_core.loader.module_ops import ModuleOps |
| 7 | from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps |
| 8 | from ltx_core.model.transformer import LTXModel |
| 9 | |
| 10 | |
| 11 | class FP8Linear(nn.Module): |
| 12 | """Linear layer with FP8 weight storage for scaled matrix multiplication.""" |
| 13 | |
| 14 | in_features: int |
| 15 | out_features: int |
| 16 | |
| 17 | def __init__( |
| 18 | self, |
| 19 | in_features: int, |
| 20 | out_features: int, |
| 21 | bias: bool = True, |
| 22 | device: torch.device | str | None = None, |
| 23 | ): |
| 24 | super().__init__() |
| 25 | self.in_features = in_features |
| 26 | self.out_features = out_features |
| 27 | |
| 28 | fp8_shape = (in_features, out_features) |
| 29 | self.weight = nn.Parameter(torch.empty(fp8_shape, dtype=torch.float8_e4m3fn, device=device)) |
| 30 | # Weight scale for FP8 dequantization (shape matches checkpoint format) |
| 31 | self.weight_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device)) |
| 32 | # Input scale for static quantization (pre-quantized checkpoints) |
| 33 | self.input_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device)) |
| 34 | |
| 35 | if bias: |
| 36 | self.bias = nn.Parameter(torch.empty(out_features, device=device)) |
| 37 | else: |
| 38 | self.register_parameter("bias", None) |
| 39 | |
| 40 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 41 | origin_shape = x.shape |
| 42 | |
| 43 | # Static quantization: use pre-computed scale |
| 44 | qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale) |
| 45 | |
| 46 | # Flatten to 2D for matmul |
| 47 | if qinput.dim() == 3: |
| 48 | qinput = qinput.reshape(-1, qinput.shape[-1]) |
| 49 | |
| 50 | # FP8 scaled matmul |
| 51 | output = torch.ops.trtllm.cublas_scaled_mm( |
| 52 | qinput, |
| 53 | self.weight, |
| 54 | scale_a=cur_input_scale, |
| 55 | scale_b=self.weight_scale, |
| 56 | bias=None, |
| 57 | out_dtype=x.dtype, |
| 58 | ) |
| 59 | |
| 60 | # Add bias |
| 61 | if self.bias is not None: |
| 62 | bias = self.bias |
| 63 | if bias.dtype != output.dtype: |
| 64 | bias = bias.to(output.dtype) |
| 65 | output = output + bias |
| 66 | |
| 67 | # Restore original shape |
| 68 | if output.dim() != len(origin_shape): |
| 69 | output_shape = list(origin_shape) |
| 70 | output_shape[-1] = output.shape[-1] |
| 71 | output = output.reshape(output_shape) |
| 72 | |
| 73 | return output |
| 74 | |
| 75 | |
| 76 | def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 77 | """ |
| 78 | Quantize a weight tensor to FP8 (float8_e4m3fn) using per-tensor scaling. |
| 79 | Args: |
| 80 | weight: The weight tensor to quantize (any dtype, will be cast to float32) |
| 81 | Returns: |
| 82 | Tuple of (quantized_weight, weight_scale): |
| 83 | - quantized_weight: FP8 tensor, transposed for cublas_scaled_mm |
| 84 | - weight_scale: Per-tensor scale factor (reciprocal of quantization scale) |
| 85 | """ |
| 86 | weight_fp32 = weight.to(torch.float32) |
| 87 | |
| 88 | fp8_min = torch.finfo(torch.float8_e4m3fn).min |
| 89 | fp8_max = torch.finfo(torch.float8_e4m3fn).max |
| 90 | |
| 91 | max_abs = torch.amax(torch.abs(weight_fp32)) |
| 92 | scale = fp8_max / max_abs |
| 93 | |
| 94 | @torch.compiler.disable |
| 95 | def _quantize( |
| 96 | weight_fp32: torch.Tensor, scale: torch.Tensor, fp8_min: torch.Tensor, fp8_max: torch.Tensor |
| 97 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 98 | quantized_weight = torch.clamp(weight_fp32 * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) |
| 99 | quantized_weight = quantized_weight.t() |
| 100 | weight_scale = scale.reciprocal() |
| 101 | return quantized_weight, weight_scale |
| 102 | |
| 103 | quantized_weight, weight_scale = _quantize(weight_fp32, scale, fp8_min, fp8_max) |
| 104 | return quantized_weight, weight_scale |
| 105 | |
| 106 | |
| 107 | def _should_skip_layer(layer_name: str, excluded_layer_substrings: tuple[str, ...]) -> bool: |
| 108 | return any(substring in layer_name for substring in excluded_layer_substrings) |
| 109 | |
| 110 | |
| 111 | EXCLUDED_LAYER_SUBSTRINGS = ( |
| 112 | "patchify_proj", |
| 113 | "adaln_single", |
| 114 | "av_ca_video_scale_shift_adaln_single", |
| 115 | "av_ca_a2v_gate_adaln_single", |
| 116 | "caption_projection", |
| 117 | "proj_out", |
| 118 | "audio_patchify_proj", |
| 119 | "audio_adaln_single", |
| 120 | "av_ca_audio_scale_shift_adaln_single", |
| 121 | "av_ca_v2a_gate_adaln_single", |
| 122 | "audio_caption_projection", |
| 123 | "audio_proj_out", |
| 124 | "transformer_blocks.0.", |
| 125 | *[f"transformer_blocks.{i}." for i in range(43, 48)], |
| 126 | ) |
| 127 | |
| 128 | |
| 129 | def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear: |
| 130 | """ |
| 131 | Create an FP8Linear layer from an nn.Linear layer. |
| 132 | Args: |
| 133 | layer: The nn.Linear layer to convert (typically on meta device) |
| 134 | Returns: |
| 135 | A new FP8Linear with the same configuration |
| 136 | """ |
| 137 | return FP8Linear( |
| 138 | in_features=layer.in_features, |
| 139 | out_features=layer.out_features, |
| 140 | bias=layer.bias is not None, |
| 141 | device=layer.weight.device, |
| 142 | ) |
| 143 | |
| 144 | |
| 145 | def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tuple[str, ...]) -> nn.Module: |
| 146 | """Replace nn.Linear layers with FP8Linear in the module tree.""" |
| 147 | replacements: list[tuple[nn.Module, str, nn.Linear]] = [] |
| 148 | |
| 149 | for name, module in model.named_modules(): |
| 150 | if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear): |
| 151 | continue |
| 152 | |
| 153 | if _should_skip_layer(name, excluded_layer_substrings): |
| 154 | continue |
| 155 | |
| 156 | if "." in name: |
| 157 | parent_name, attr_name = name.rsplit(".", 1) |
| 158 | parent = model.get_submodule(parent_name) |
| 159 | else: |
| 160 | parent = model |
| 161 | attr_name = name |
| 162 | |
| 163 | replacements.append((parent, attr_name, module)) |
| 164 | |
| 165 | for parent, attr_name, linear in replacements: |
| 166 | setattr(parent, attr_name, _linear_to_fp8linear(linear)) |
| 167 | |
| 168 | return model |
| 169 | |
| 170 | |
| 171 | def _create_transpose_kv_operation( |
| 172 | excluded_layer_substrings: tuple[str, ...], |
| 173 | ) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]: |
| 174 | def transpose_if_matches(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: |
| 175 | # Only process .weight keys |
| 176 | if not key.endswith(".weight"): |
| 177 | return [KeyValueOperationResult(key, value)] |
| 178 | |
| 179 | # Only transpose 2D FP8 tensors (Linear weights) |
| 180 | if value.dim() != 2 or value.dtype != torch.float8_e4m3fn: |
| 181 | return [KeyValueOperationResult(key, value)] |
| 182 | |
| 183 | # Check if the layer is excluded |
| 184 | layer_name = key.rsplit(".weight", 1)[0] |
| 185 | if _should_skip_layer(layer_name, excluded_layer_substrings): |
| 186 | return [KeyValueOperationResult(key, value)] |
| 187 | |
| 188 | # Transpose to cuBLAS layout (in, out) |
| 189 | transposed_weight = value.t() |
| 190 | |
| 191 | return [KeyValueOperationResult(key, transposed_weight)] |
| 192 | |
| 193 | return transpose_if_matches |
| 194 | |
| 195 | |
| 196 | FP8_TRANSPOSE_SD_OPS = SDOps("fp8_transpose_weights").with_kv_operation( |
| 197 | _create_transpose_kv_operation(EXCLUDED_LAYER_SUBSTRINGS), |
| 198 | key_prefix="transformer_blocks.", |
| 199 | key_suffix=".weight", |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | FP8_PREPARE_MODULE_OPS = ModuleOps( |
| 204 | name="fp8_prepare_for_loading", |
| 205 | matcher=lambda model: isinstance(model, LTXModel), |
| 206 | mutator=lambda model: _apply_fp8_prepare_to_model(model, EXCLUDED_LAYER_SUBSTRINGS), |
| 207 | ) |
| 208 |