fp8_qat_train_config

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.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig

QuantizationConfig()

Base class for quantization configs.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_name

get_name() -> QuantizationMethods

Name of the quantization method.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_min_capability

get_min_capability() -> int

Minimum GPU capability to support the quantization method.

E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_config_filenames

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.from_config

from_config(config: dict[str, Any]) -> QuantizationConfig

Create a config class from the model's quantization config.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.override_quantization_method

override_quantization_method(hf_quant_cfg, user_quant) -> QuantizationMethods | None

Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_from_keys

get_from_keys(config: dict[str, Any], keys: list[str]) -> Any

Get a value from the model's quantization config.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_from_keys_or

get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any

Get a optional value from the model's quantization config.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_quant_method

get_quant_method(layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None

Get the quantize method to use for the quantized layer.

Parameters

layer: torch.nn.Module

The layer for the quant method.

prefix: str

The full name of the layer in the state dict

Returns: The quantize method. None if the given layer doesn't support quant method.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.validate_runtime

validate_runtime(device: torch.device) -> None

Fail fast on a device the scheme cannot execute on.

Called once by loaders that select a serialized checkpoint scheme, before any weight is read. The default accepts every device; schemes with kernel or capability requirements override it.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizationConfig.get_cache_scale

get_cache_scale(name: str) -> str | None

fastvideo.layers.quantization.fp8_qat_train_config.QuantizeMethodBase

QuantizeMethodBase()

Base class for different quantized methods.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizeMethodBase.create_weights

create_weights(layer: torch.nn.Module, *weight_args, **extra_weight_attrs)

Create weights for a layer.

The weights will be set as attributes of the layer.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizeMethodBase.apply

apply(layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor

Apply the weights in layer to the input tensor.

Expects create_weights to have been called before on the layer.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizeMethodBase.embedding

embedding(layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor

Gather embeddings in the layer based on indices in the input tensor.

Expects create_weights to have been called before on the layer.

fastvideo.layers.quantization.fp8_qat_train_config.QuantizeMethodBase.process_weights_after_loading

process_weights_after_loading(layer: nn.Module) -> None

Process the weight after loading.

This can be used for example, to transpose weights for computation.

fastvideo.layers.quantization.fp8_qat_train_config.set_weight_attrs

set_weight_attrs(weight: torch.Tensor, weight_attrs: dict[str, Any] | None)

Set attributes on a weight tensor.

This method is used to set attributes on a weight tensor. This method will not overwrite existing attributes.

Parameters

weight: torch.Tensor

The weight tensor.

weight_attrs: dict[str, Any] | None

A dictionary of attributes to set on the weight tensor.

fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainQuantizeMethod

FP8QATTrainQuantizeMethod()

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, &quot;tensor&quot;)</code></pre></details>
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainQuantizeMethod.create_weights
create_weights(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)
Source
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)
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainQuantizeMethod.apply
apply(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor
Source
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")
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig
FP8QATTrainConfig()
Source
class FP8QATTrainConfig(QuantizationConfig):
def __init__(self) -&gt; None:
    super().__init__()

def get_name(self):
    return &quot;fp8_qat_train&quot;

def get_supported_act_dtypes(self):
    return &#91;torch.bfloat16, torch.float16&#93;

@classmethod
def get_min_capability(cls):
    return 89

@staticmethod
def get_config_filenames():
    return &#91;&#93;

@classmethod
def from_config(cls, config: dict&#91;str, Any&#93;) -&gt; &quot;FP8QATTrainConfig&quot;:
    return cls()

def get_quant_method(self, layer: torch.nn.Module, prefix: str):
    from fastvideo.layers.linear import LinearBase
    fp8_layers = &#91;&quot;ffn.fc_in&quot;, &quot;ffn.fc_out&quot;, &quot;to_q&quot;, &quot;to_k&quot;, &quot;to_v&quot;, &quot;to_out&quot;&#93;
    if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp8_layers):
        return FP8QATTrainQuantizeMethod()
    return None</code></pre></details>
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig.get_name
get_name()
Source
def get_name(self):
    return "fp8_qat_train"
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig.get_supported_act_dtypes
get_supported_act_dtypes()
Source
def get_supported_act_dtypes(self):
    return [torch.bfloat16, torch.float16]
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig.get_min_capability
get_min_capability()
Source
@classmethod
def get_min_capability(cls):
    return 89
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig.get_config_filenames
get_config_filenames()
Source
@staticmethod
def get_config_filenames():
    return []
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig.from_config
from_config(config: dict[str, Any]) -> FP8QATTrainConfig
Source
@classmethod
def from_config(cls, config: dict[str, Any]) -> "FP8QATTrainConfig":
    return cls()
fastvideo.layers.quantization.fp8_qat_train_config.FP8QATTrainConfig.get_quant_method
get_quant_method(layer: torch.nn.Module, prefix: str)
Source
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