diff --git a/fastvideo/layers/fp8linear.py b/fastvideo/layers/fp8linear.py new file mode 100644 index 000000000..f9e141f76 --- /dev/null +++ b/fastvideo/layers/fp8linear.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP8 quantization-aware training for linear layers. + +Mirror of ``fp4linear.py`` but for FP8 (e4m3). The forward pass quantizes both +activations and weights to FP8 and runs ``torch._scaled_mm``; the backward pass +is a bf16 straight-through estimator so the high-precision master weights stay +trainable. Falls back to a bf16 fake-quant forward on GPUs older than sm89. +""" +import torch + +FP8_DTYPE = torch.float8_e4m3fn +FP8_MAX = float(torch.finfo(FP8_DTYPE).max) # 448.0 +FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0) + + +def _supports_fp8_compute() -> bool: + """Whether the active device supports FP8 ``_scaled_mm`` (sm89+).""" + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9) + + +def _quantize_tensorwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(x_fp8 [M, K], x_scale [1] float32)``.""" + x_absmax = x_2d.abs().amax().float() + x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE) + x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) + return x_fp8, x_scale.view(1) + + +def _quantize_rowwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``.""" + x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float() + x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE) + x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) + return x_fp8, x_scale + + +def _fake_quant(x_2d: torch.Tensor, granularity: str) -> torch.Tensor: + """bf16 fake-quant (quantize then dequantize) for pre-sm89 fallback.""" + if granularity == "channel": + x_fp8, x_scale = _quantize_rowwise(x_2d) + else: + x_fp8, x_scale = _quantize_tensorwise(x_2d) + return x_fp8.to(x_2d.dtype) * x_scale.to(x_2d.dtype) + + +class _LinearFWD8BWD16Fn(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, weight, bias, granularity="tensor"): + # assert/normalize activation dtype + if x.dtype not in (torch.float16, torch.bfloat16): + x = x.to(dtype=torch.bfloat16) + + # cast params (can be fp32) to activation dtype for quantization + weight_cast = weight.to(dtype=x.dtype) + bias_cast = bias.to(dtype=x.dtype) if bias is not None else None + + orig_shape = x.shape + k = weight_cast.shape[1] + n = weight_cast.shape[0] + x2d = x.reshape(-1, k).contiguous() + + if not _supports_fp8_compute(): + # bf16 fake-quant fallback: simulate the FP8 rounding error but + # compute the matmul in bf16. + x_fq = _fake_quant(x2d, granularity) + w_fq = _fake_quant(weight_cast, granularity) + out2d = x_fq.matmul(w_fq.t()) + if bias_cast is not None: + out2d = out2d + bias_cast + ctx.save_for_backward(x2d, weight, bias) + ctx.n = n + ctx.orig_shape = orig_shape + return out2d.reshape(*orig_shape[:-1], n) + + if granularity == "channel": + x_fp8, x_scale = _quantize_rowwise(x2d) + w_fp8, w_scale = _quantize_rowwise(weight_cast) + scale_b = w_scale.view(1, -1) + else: + x_fp8, x_scale = _quantize_tensorwise(x2d) + w_fp8, w_scale = _quantize_tensorwise(weight_cast) + scale_b = w_scale + + out2d = torch._scaled_mm( + x_fp8, + w_fp8.t(), + scale_a=x_scale, + scale_b=scale_b, + out_dtype=x.dtype, + ) + if isinstance(out2d, tuple): + out2d = out2d[0] + + if bias_cast is not None: + out2d = out2d + bias_cast + + # save tensors for backward (keep original dtypes) + ctx.save_for_backward(x2d, weight, bias) + ctx.n = n + ctx.orig_shape = orig_shape + return out2d.reshape(*orig_shape[:-1], n) + + @staticmethod + def backward(ctx, grad_out): + x2d, weight, bias = ctx.saved_tensors + M = x2d.shape[0] + n = ctx.n + + grad_out_2d = grad_out.reshape(M, n).contiguous() + + # bf16 straight-through estimator: gradients flow through the + # full-precision master weights, not the FP8 quantized values. + weight_cast = weight.to(dtype=grad_out.dtype) + x_cast = x2d.to(dtype=grad_out.dtype) + + grad_x = grad_out_2d.matmul(weight_cast).reshape(*ctx.orig_shape) + grad_w = grad_out_2d.t().matmul(x_cast) + grad_b = grad_out_2d.sum(dim=0) if bias is not None else None + + # None for the extra forward arg (granularity) + return grad_x, grad_w, grad_b, None + + +def fp8_linear_forward(self, x: torch.Tensor) -> torch.Tensor: + # pass config **positionally**; autograd.Function.apply ignores kwargs + return _LinearFWD8BWD16Fn.apply(x, self.weight, self.bias, "tensor"), None diff --git a/fastvideo/layers/quantization/__init__.py b/fastvideo/layers/quantization/__init__.py index e1d45bc22..8f8d8e131 100644 --- a/fastvideo/layers/quantization/__init__.py +++ b/fastvideo/layers/quantization/__init__.py @@ -2,7 +2,7 @@ from typing import Literal, get_args from fastvideo.layers.quantization.base_config import QuantizationConfig -QuantizationMethods = Literal[None, "AbsMaxFP8", "NVFP4", "nvfp4_qat", "nvfp4_qat_train"] +QuantizationMethods = Literal[None, "AbsMaxFP8", "NVFP4", "nvfp4_qat", "nvfp4_qat_train", "fp8_qat_train"] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) @@ -54,12 +54,14 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .nvfp4_config import NVFP4Config from .nvfp4_qat_config import NVFP4QATConfig from .nvfp4_qat_train_config import NVFP4QATTrainConfig + from .fp8_qat_train_config import FP8QATTrainConfig method_to_config: dict[str, type[QuantizationConfig]] = { "AbsMaxFP8": AbsMaxFP8Config, "NVFP4": NVFP4Config, "nvfp4_qat": NVFP4QATConfig, "nvfp4_qat_train": NVFP4QATTrainConfig, + "fp8_qat_train": FP8QATTrainConfig, } # Update the `method_to_config` with customized quantization methods. method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG) diff --git a/fastvideo/layers/quantization/fp8_qat_train_config.py b/fastvideo/layers/quantization/fp8_qat_train_config.py new file mode 100644 index 000000000..516112ebb --- /dev/null +++ b/fastvideo/layers/quantization/fp8_qat_train_config.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP8 (e4m3) quantization-aware *training* linear method (straight-through estimator). + +Mirror of ``nvfp4_qat_train_config.py`` but for FP8. The weight stays a trainable +bf16/fp32 master that is fake-quantized to FP8 on every forward, with a +full-precision backward (STE), so the model learns to absorb FP8 linear error. + +The STE lives in ``fastvideo.layers.fp8linear._LinearFWD8BWD16Fn`` (FP8 forward +via ``torch._scaled_mm`` on sm89+, with a bf16 fake-quant fallback on older GPUs; +full-precision backward). This method bridges it into the standard +``quant_config`` path, so it activates via ``transformer_quant="fp8_qat_train"`` +on the same Wan-2.1 layers as the FP4 path (to_q/k/v/out + ffn). No conversion is +needed: the weight is kept in full precision and quantized on the fly each step. + +Unlike the FP4 path this needs no flashinfer and runs on any sm89+ GPU (and even +older ones via the bf16 fallback), not just Blackwell. +""" +import logging +from typing import Any + +import torch +from torch.nn.parameter import Parameter + +from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase +from fastvideo.models.utils import set_weight_attrs + +logger = logging.getLogger(__name__) + + +class FP8QATTrainQuantizeMethod(QuantizeMethodBase): + + def create_weights(self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], + input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs): + # Trainable master weight, fake-quantized to FP8 on each forward. + weight = Parameter(torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=True) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + # FP8 forward + full-precision backward (STE). + from fastvideo.layers.fp8linear import _LinearFWD8BWD16Fn + return _LinearFWD8BWD16Fn.apply(x, layer.weight, bias, "tensor") + + +class FP8QATTrainConfig(QuantizationConfig): + + def __init__(self) -> None: + super().__init__() + + def get_name(self): + return "fp8_qat_train" + + def get_supported_act_dtypes(self): + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls): + return 89 + + @staticmethod + def get_config_filenames(): + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "FP8QATTrainConfig": + return cls() + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + from fastvideo.layers.linear import LinearBase + fp8_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"] + if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp8_layers): + return FP8QATTrainQuantizeMethod() + return None