--- title: nvfp4_qat_train_config ---

NVFP4 quantization-aware training linear method (straight-through estimator).

The inference nvfp4_qat config quantizes each weight to FP4 once at load time (convert_model_to_fp4) and has no gradient path — it is inference only. For QAT finetuning the weight must stay a trainable bf16/fp32 master that is fake-quantized to FP4 on every forward, with a full-precision backward (a straight-through estimator), so the model learns to absorb FP4 linear error.

That STE already exists in fastvideo.layers.fp4linear._LinearFWD4BWD16Fn (FP4 forward, full-precision backward) but is otherwise unwired. This method bridges it into the standard quant_config path. Wan-style prefixes retain the broad attention/FFN projection profile used by nvfp4_qat; LTX-2 prefixes use the exact curated target set from its NVFP4 deployment config. No convert_model_to_fp4 is needed: the weight is kept in full precision and quantized on the fly each step.

fastvideo.layers.quantization.nvfp4_qat_train_config.QuantizationConfig

QuantizationConfig()

Base class for quantization configs.

fastvideo.layers.quantization.nvfp4_qat_train_config.QuantizationConfig.get_name

get_name() -> QuantizationMethods

Name of the quantization method.

fastvideo.layers.quantization.nvfp4_qat_train_config.QuantizationConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

fastvideo.layers.quantization.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_qat_train_config.QuantizationConfig.get_cache_scale

get_cache_scale(name: str) -> str | None

fastvideo.layers.quantization.nvfp4_qat_train_config.QuantizeMethodBase

QuantizeMethodBase()

Base class for different quantized methods.

fastvideo.layers.quantization.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_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.nvfp4_qat_train_config.is_ltx2_nvfp4_linear_prefix

is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool

Return whether prefix belongs to the LTX-2 NVFP4 deployment set.

fastvideo.layers.quantization.nvfp4_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.nvfp4_qat_train_config.NVFP4QATTrainQuantizeMethod

NVFP4QATTrainQuantizeMethod()
Source
class NVFP4QATTrainQuantizeMethod(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 FP4 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:
        # FP4 forward + full-precision backward (STE).
        from fastvideo.layers.fp4linear import _LinearFWD4BWD16Fn
        return _LinearFWD4BWD16Fn.apply(x, layer.weight, bias, "cutlass", 16, True)

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainQuantizeMethod.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 FP4 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.nvfp4_qat_train_config.NVFP4QATTrainQuantizeMethod.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:
    # FP4 forward + full-precision backward (STE).
    from fastvideo.layers.fp4linear import _LinearFWD4BWD16Fn
    return _LinearFWD4BWD16Fn.apply(x, layer.weight, bias, "cutlass", 16, True)

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig

NVFP4QATTrainConfig()
Source
class NVFP4QATTrainConfig(QuantizationConfig):

    def __init__(self) -> None:
        super().__init__()

    def get_name(self):
        return "nvfp4_qat_train"

    def get_supported_act_dtypes(self):
        return [torch.bfloat16, torch.float16]

    @classmethod
    def get_min_capability(cls):
        return 100

    @staticmethod
    def get_config_filenames():
        return []

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "NVFP4QATTrainConfig":
        return cls()

    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        from fastvideo.layers.linear import LinearBase
        if not isinstance(layer, LinearBase):
            return None
        matches = (is_ltx2_nvfp4_linear_prefix(prefix) if prefix.startswith("ltx2.") else any(
            layer_name in prefix for layer_name in DEFAULT_FP4_LAYERS))
        if matches:
            return NVFP4QATTrainQuantizeMethod()
        return None

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig.get_name

get_name()
Source
def get_name(self):
    return "nvfp4_qat_train"

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig.get_supported_act_dtypes

get_supported_act_dtypes()
Source
def get_supported_act_dtypes(self):
    return [torch.bfloat16, torch.float16]

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig.get_min_capability

get_min_capability()
Source
@classmethod
def get_min_capability(cls):
    return 100

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig.get_config_filenames

get_config_filenames()
Source
@staticmethod
def get_config_filenames():
    return []

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig.from_config

from_config(config: dict[str, Any]) -> NVFP4QATTrainConfig
Source
@classmethod
def from_config(cls, config: dict[str, Any]) -> "NVFP4QATTrainConfig":
    return cls()

fastvideo.layers.quantization.nvfp4_qat_train_config.NVFP4QATTrainConfig.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
    if not isinstance(layer, LinearBase):
        return None
    matches = (is_ltx2_nvfp4_linear_prefix(prefix) if prefix.startswith("ltx2.") else any(
        layer_name in prefix for layer_name in DEFAULT_FP4_LAYERS))
    if matches:
        return NVFP4QATTrainQuantizeMethod()
    return None