nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat]: NVFP4 text encoder for MiniMax H3 (#1837)

Browse Source
This commit is contained in:
Kyle Hu
2026-09-12 18:31:59 -07:00
committed by GitHub
parent aef4e9b3b1
commit 3a3ad3d209
5 changed files with 946 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/layers/quantization/base_config.py
+9
View File
@@ -139,5 +139,14 @@ class QuantizationConfig(ABC):
"""
raise NotImplementedError
def validate_runtime(self, 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.
"""
return None
def get_cache_scale(self, name: str) -> str | None:
return None
fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py
+425
View File
@@ -0,0 +1,425 @@
# SPDX-License-Identifier: Apache-2.0
"""Serialized NVFP4 execution for the MiniMax-H3 Qwen3-VL encoder.
The bf16 conditioner is 48 GB resident and sets the single-GB10 peak once the
DiT runs in FP8. A checkpoint written by
``scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py``
stores every ``language_model.layers.*`` linear as three tensors and no
``weight``::
<prefix>.weight_packed uint8 [out, in // 2] two E2M1 values per byte
<prefix>.weight_scale uint8 [out, in // 16] one E4M3 scale per 16 values,
FlashInfer 128x4 swizzled layout
<prefix>.weight_global_scale float32 [1] (448 * 6) / amax(|W|)
The bytes are exactly what ``flashinfer.nvfp4_quantize(W, global_scale,
sfLayout=SfLayout.layout_128x4)`` returns, so loading them reproduces the state
``nvfp4_config.convert_model_to_nvfp4`` builds at runtime without ever
materializing the bf16 weight. Activations are quantized per call with a unit
global scale and multiplied with ``flashinfer.mm_fp4``.
The checkpoint selects this path through ``config.json``; every field is
required so a checkpoint from another exporter cannot pass by omission::
"quantization_config": {"quant_method": "nvfp4", "activation_scheme": "dynamic",
"fmt": "e2m1", "group_size": 16, "scale_fmt": "e4m3",
"scale_layout": "128x4",
"modules_to_not_convert": ["model.visual", "lm_head"]}
Whole projection kinds may stay bf16 by listing their suffix in
``modules_to_not_convert`` (for example ``"mlp.down_proj"``); those linears
keep a plain ``weight`` and the unquantized method. Only the seven language
projection names are accepted there, so a typo fails at config time.
Single GPU only: packed columns and swizzled scale rows cannot be narrowed per
tensor-parallel rank without repacking. Missing tensors are reported by the
loader's strict check before the post-load hook runs; the hook adds only the
content checks a copied tensor can still fail.
"""
from typing import Any
import torch
from torch import nn
from torch.nn.parameter import Parameter
from fastvideo.distributed import get_tp_world_size
from fastvideo.layers.linear import LinearBase, LinearMethodBase
from fastvideo.layers.quantization.base_config import QuantizationConfig
from fastvideo.layers.quantization.nvfp4_config import (
_coerce_fp4_input_dtype,
_mm_fp4,
_nvfp4_quantize,
_register_ops_once,
_require_flashinfer,
)
from fastvideo.models.utils import set_weight_attrs
NVFP4_GROUP_SIZE = 16
NVFP4_SCALE_LAYOUT = "128x4"
# FlashInfer's 128x4 layout tiles scales in 128 rows by 4 scale columns. Keeping
# every weight a whole number of tiles means the serialized scale tensor is
# exactly [out, in // 16] with no padding to describe. Every Qwen3-VL language
# linear satisfies this (out in {1024, 5120, 8192, 25600}, in in {5120, 8192, 25600}).
NVFP4_ROW_TILE = 128
NVFP4_COLUMN_MULTIPLE = 4 * NVFP4_GROUP_SIZE
# The linears of one Qwen3-VL language layer, as the checkpoint and the model name them.
LANGUAGE_PROJECTIONS = (
"self_attn.q_proj",
"self_attn.k_proj",
"self_attn.v_proj",
"self_attn.o_proj",
"mlp.gate_proj",
"mlp.up_proj",
"mlp.down_proj",
)
# The three tensors a quantized linear stores, as parameter names on the layer
# and as checkpoint key suffixes. The converter imports these so both sides
# spell them once.
NVFP4_TENSOR_SUFFIXES = ("weight_packed", "weight_scale", "weight_global_scale")
_REQUIRED_METADATA_KEYS = ("activation_scheme", "fmt", "group_size", "scale_fmt", "scale_layout",
"modules_to_not_convert")
_E2M1_MAX = 6.0
_E4M3_MAX = 448.0
# An E4M3 byte of 0xFF is NaN, so a scale tile that is still all 0xFF after
# loading is one the checkpoint never filled. The loader copies whole tensors
# or asserts, so the first 128x4 tile speaks for the tensor.
_UNLOADED_SCALE_BYTE = 0xFF
_SCALE_TILE_BYTES = NVFP4_ROW_TILE * 4
def validate_nvfp4_geometry(output_size: int, input_size: int) -> None:
if output_size % NVFP4_ROW_TILE:
raise ValueError(f"MiniMax-H3 serialized NVFP4 requires output_size divisible by {NVFP4_ROW_TILE}, "
f"got {output_size}")
if input_size % NVFP4_COLUMN_MULTIPLE:
raise ValueError(f"MiniMax-H3 serialized NVFP4 requires input_size divisible by {NVFP4_COLUMN_MULTIPLE}, "
f"got {input_size}")
def nvfp4_packed_weight_shape(output_size: int, input_size: int) -> tuple[int, int]:
return output_size, input_size // 2
def nvfp4_scale_shape(output_size: int, input_size: int) -> tuple[int, int]:
return output_size, input_size // NVFP4_GROUP_SIZE
def nvfp4_weight_global_scale(weight: torch.Tensor) -> torch.Tensor:
"""The per-tensor scale ``convert_model_to_nvfp4`` uses: E4M3 max times E2M1 max over amax.
Unlike the runtime converter this refuses NaN or infinite weights instead of
mapping them to zero, because a checkpoint written from them would be wrong
forever.
"""
amax = weight.float().abs().max()
if not torch.isfinite(amax) or amax <= 0:
raise ValueError("MiniMax-H3 NVFP4 global scale needs a finite, non-zero weight amax")
scale = ((_E4M3_MAX * _E2M1_MAX) / amax).to(torch.float32)
if not torch.isfinite(scale):
raise ValueError(f"MiniMax-H3 NVFP4 global scale overflowed float32 for weight amax {amax.item():.3e}")
return scale
def serialized_nvfp4_quantization_config(
*,
keep_bf16: tuple[str, ...] | list[str] = (),
producer: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""The ``config.json`` ``quantization_config`` the converter writes and ``from_config`` accepts.
``keep_bf16`` lists projection kinds from ``LANGUAGE_PROJECTIONS`` that stay
bf16 in every language layer; they are appended to ``modules_to_not_convert``.
"""
unknown = sorted(set(keep_bf16) - set(LANGUAGE_PROJECTIONS))
if unknown:
raise ValueError(f"keep_bf16 names unknown projection kinds {unknown}; choose from {LANGUAGE_PROJECTIONS}")
config: dict[str, Any] = {
"quant_method": "nvfp4",
"activation_scheme": "dynamic",
"fmt": "e2m1",
"group_size": NVFP4_GROUP_SIZE,
"scale_fmt": "e4m3",
"scale_layout": NVFP4_SCALE_LAYOUT,
"modules_to_not_convert": ["model.visual", "lm_head", *keep_bf16],
}
if producer:
config["producer"] = dict(producer)
return config
def _quantize_activation_nvfp4(x_2d: torch.Tensor, global_scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize one bf16 activation the way ``NVFP4QuantizeMethod`` does: unit global scale, 128x4 layout.
The scale tensor is returned as uint8 bytes whatever dtype FlashInfer labels
it with, so it always matches the uint8 weight scales the checkpoint stores.
"""
sf_layout, _, _ = _require_flashinfer()
x_fp4, x_scale = _nvfp4_quantize(x_2d, global_scale, sfLayout=sf_layout.layout_128x4, do_shuffle=False)
if x_scale.dtype != torch.uint8:
x_scale = x_scale.view(torch.uint8)
return x_fp4, x_scale
def _nvfp4_linear(
x_fp4: torch.Tensor,
x_scale: torch.Tensor,
weight_packed: torch.Tensor,
weight_scale: torch.Tensor,
alpha: torch.Tensor,
) -> torch.Tensor:
"""``x @ W.T`` on packed operands. Mirrors ``NVFP4QuantizeMethod.apply``: the weight is
handed to ``mm_fp4`` transposed, the output is bf16, the backend is FlashInfer's choice."""
return _mm_fp4(
x_fp4,
weight_packed.t(),
x_scale,
weight_scale.t(),
alpha,
torch.bfloat16,
None,
backend="auto",
)
class MiniMaxH3SerializedNVFP4Config(QuantizationConfig):
"""Serialized 16-group NVFP4 contract for the H3 text encoder.
The group size and scale layout are fixed by the loader's parameter shapes,
so the only state is which projection kinds the checkpoint kept in bf16.
"""
def __init__(self, bf16_projections: tuple[str, ...] = ()) -> None:
super().__init__()
unknown = sorted(set(bf16_projections) - set(LANGUAGE_PROJECTIONS))
if unknown:
raise ValueError(f"MiniMax-H3 serialized NVFP4 cannot keep unknown projection kinds {unknown} in bf16; "
f"choose from {LANGUAGE_PROJECTIONS}")
# Projection kinds the checkpoint kept in bf16 in every language layer,
# e.g. ("mlp.down_proj",). Those linears load a plain weight.
self.bf16_projections = tuple(bf16_projections)
@classmethod
def get_name(cls) -> str:
return "nvfp4"
@classmethod
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
return [torch.bfloat16]
@classmethod
def get_min_capability(cls) -> int:
return 100
@staticmethod
def get_config_filenames() -> list[str]:
return []
@classmethod
def from_config(cls, config: dict[str, Any]) -> "MiniMaxH3SerializedNVFP4Config":
quant_method = str(config.get("quant_method", "")).lower()
if quant_method != "nvfp4":
raise ValueError(f"MiniMax-H3 serialized NVFP4 config got quant_method {quant_method!r}")
missing = [key for key in _REQUIRED_METADATA_KEYS if key not in config]
if missing:
raise ValueError("MiniMax-H3 serialized NVFP4 requires explicit quantization_config fields; "
f"missing {missing}")
if str(config["activation_scheme"]).lower() != "dynamic":
raise ValueError("MiniMax-H3 serialized NVFP4 requires dynamic activation quantization")
if str(config["fmt"]).lower() not in ("e2m1", "float4_e2m1fn", "nvfp4"):
raise ValueError(f"MiniMax-H3 serialized NVFP4 requires E2M1 weights, got {config['fmt']!r}")
if str(config["scale_fmt"]).lower() not in ("e4m3", "float8_e4m3fn"):
raise ValueError(f"MiniMax-H3 serialized NVFP4 requires E4M3 block scales, got {config['scale_fmt']!r}")
group_size = config["group_size"]
if isinstance(group_size, bool) or group_size != NVFP4_GROUP_SIZE:
raise ValueError(f"MiniMax-H3 serialized NVFP4 requires group_size={NVFP4_GROUP_SIZE}, got {group_size!r}")
if config["scale_layout"] != NVFP4_SCALE_LAYOUT:
raise ValueError(f"MiniMax-H3 serialized NVFP4 requires scale_layout={NVFP4_SCALE_LAYOUT!r}, "
f"got {config['scale_layout']!r}")
ignored_layers = config["modules_to_not_convert"]
if not isinstance(ignored_layers, list | tuple):
raise ValueError("MiniMax-H3 serialized NVFP4 modules_to_not_convert must be a sequence")
bf16_projections: list[str] = []
saw_visual = False
for name in ignored_layers:
if not isinstance(name, str) or not name:
raise ValueError("MiniMax-H3 serialized NVFP4 modules_to_not_convert entries must be non-empty "
f"strings, got {name!r}")
if "visual" in name:
saw_visual = True
elif name == "lm_head" or name.endswith(".lm_head"):
continue
elif name in LANGUAGE_PROJECTIONS:
bf16_projections.append(name)
else:
# A whole projection kind may stay bf16 across every language
# layer; a single layer, a typo, or an HF-style full path is not
# a contract this loader supports.
raise ValueError("MiniMax-H3 serialized NVFP4 keeps bf16 per projection kind only; "
f"modules_to_not_convert entry {name!r} is not one of {LANGUAGE_PROJECTIONS}")
if not saw_visual:
raise ValueError("MiniMax-H3 serialized NVFP4 requires the vision stack to be listed in "
"modules_to_not_convert")
return cls(tuple(bf16_projections))
def validate_runtime(self, device: torch.device) -> None:
if device.type != "cuda":
raise RuntimeError(f"MiniMax-H3 serialized NVFP4 requires a CUDA device; got {device.type!r}")
capability = torch.cuda.get_device_capability(device)
capability_number = capability[0] * 10 + capability[1]
if capability_number < self.get_min_capability():
raise RuntimeError("MiniMax-H3 serialized NVFP4 requires GPU capability "
f"sm{self.get_min_capability()} or newer, got sm{capability_number}")
if capability[0] not in (10, 12):
raise RuntimeError("MiniMax-H3 serialized NVFP4 runs FlashInfer's Blackwell FP4 GEMM; "
f"got unsupported sm{capability_number}")
if get_tp_world_size() > 1:
raise NotImplementedError("MiniMax-H3 serialized NVFP4 supports a single GPU: packed FP4 columns and "
"128x4 swizzled scale rows cannot be narrowed per tensor-parallel rank")
sf_layout, _, _ = _require_flashinfer()
if not hasattr(sf_layout, "layout_128x4"):
raise RuntimeError("The installed flashinfer has no SfLayout.layout_128x4; MiniMax-H3 serialized NVFP4 "
"checkpoints store their scales in that layout")
_register_ops_once()
def is_kept_bf16(self, prefix: str) -> bool:
return any(prefix == name or prefix.endswith("." + name) for name in self.bf16_projections)
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
if not isinstance(layer, LinearBase) or ".language_model.layers." not in prefix:
return None
if self.is_kept_bf16(prefix):
return None
return MiniMaxH3SerializedNVFP4LinearMethod()
def _strict_dtype_loader(base_loader):
"""Wrap a layer's weight loader so a checkpoint tensor must carry the parameter's exact dtype."""
if base_loader is None:
raise ValueError("MiniMax-H3 serialized NVFP4 linears need the layer's weight_loader")
def load(param: torch.Tensor, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any):
if loaded_weight.dtype != param.dtype:
raise ValueError("Serialized MiniMax-H3 NVFP4 tensors must be stored as their declared dtype; "
f"got {loaded_weight.dtype} for a {param.dtype} parameter of shape {tuple(param.shape)}")
return base_loader(param, loaded_weight, *args, **kwargs)
return load
class MiniMaxH3SerializedNVFP4LinearMethod(LinearMethodBase):
"""Execute serialized NVFP4 weights without re-quantizing them."""
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,
) -> None:
output_size_per_partition = sum(output_partition_sizes)
validate_nvfp4_geometry(output_size_per_partition, input_size_per_partition)
layer.logical_widths = output_partition_sizes
layer.input_size_per_partition = input_size_per_partition
layer.output_size_per_partition = output_size_per_partition
layer.orig_dtype = params_dtype
# No input_dim/output_dim: these tensors are not shardable, and without
# the attributes the weight loaders copy them whole. The copy is a numeric
# cast, so a checkpoint that stored scales as float8 bytes would be
# silently converted value by value; refuse any dtype but the declared one.
loader_attrs = {"weight_loader": _strict_dtype_loader(extra_weight_attrs.get("weight_loader"))}
packed_name, scale_name, global_scale_name = NVFP4_TENSOR_SUFFIXES
weight_packed = Parameter(
torch.zeros(nvfp4_packed_weight_shape(output_size_per_partition, input_size_per_partition),
dtype=torch.uint8),
requires_grad=False,
)
set_weight_attrs(weight_packed, loader_attrs)
layer.register_parameter(packed_name, weight_packed)
weight_scale = Parameter(
torch.full(nvfp4_scale_shape(output_size_per_partition, input_size_per_partition),
_UNLOADED_SCALE_BYTE,
dtype=torch.uint8),
requires_grad=False,
)
set_weight_attrs(weight_scale, loader_attrs)
layer.register_parameter(scale_name, weight_scale)
weight_global_scale = Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False)
set_weight_attrs(weight_global_scale, loader_attrs)
layer.register_parameter(global_scale_name, weight_global_scale)
# No bf16 weight ever exists on this layer; ``None`` keeps ``layer.weight``
# readable for code that inspects it, matching the purged NVFP4 path.
layer.register_parameter("weight", None)
def process_weights_after_loading(self, layer: nn.Module) -> None:
"""Reject a layer the checkpoint did not fill and derive the GEMM multiplier.
Shapes and dtypes are fixed by ``create_weights`` and enforced by the
weight loader's copy, and the loader reports missing tensors by name
before this runs; what remains is content a copied tensor can still get
wrong: a non-positive global scale, or a scale tile left at its 0xFF
initializer by a direct caller that bypassed the loader.
"""
weight_scale = layer.weight_scale
global_scale = float(layer.weight_global_scale.item())
if not (global_scale > 0) or global_scale == float("inf"):
raise ValueError("Serialized MiniMax-H3 NVFP4 weight_global_scale was not loaded: "
f"it must be a finite positive value, got {global_scale}")
if bool((weight_scale.view(-1)[:_SCALE_TILE_BYTES] == _UNLOADED_SCALE_BYTE).all()):
raise ValueError("Serialized MiniMax-H3 NVFP4 weight_scale was not loaded: its first tile is still 0xFF")
# ``mm_fp4`` folds both global scales into one multiplier. Activations use a
# unit global scale, so the multiplier is the inverse weight global scale.
device = weight_scale.device
layer.register_buffer("_nvfp4_alpha", torch.tensor(1.0 / global_scale, dtype=torch.float32, device=device),
persistent=False)
layer.register_buffer("_nvfp4_x_global_scale", torch.ones((), dtype=torch.float32, device=device),
persistent=False)
@staticmethod
def _apply_finalized(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
x = _coerce_fp4_input_dtype(x)
original_shape = x.shape
if x.numel() == 0:
# An empty prompt has nothing to quantize; the FP4 kernels are not defined for zero rows.
return x.new_zeros(*original_shape[:-1], layer.output_size_per_partition, dtype=torch.bfloat16)
x_fp4, x_scale = _quantize_activation_nvfp4(x.reshape(-1, original_shape[-1]), layer._nvfp4_x_global_scale)
output = _nvfp4_linear(x_fp4, x_scale, layer.weight_packed, layer.weight_scale, layer._nvfp4_alpha)
if bias is not None:
# The GEMM emits bf16; keep it that way whatever dtype the bias was built in.
output = output + bias.to(output.dtype)
return output.view(*original_shape[:-1], output.shape[-1])
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if x.device.type != "cuda":
raise RuntimeError("MiniMax-H3 serialized NVFP4 execution requires CUDA")
if getattr(layer, "_nvfp4_alpha", None) is None:
raise RuntimeError("MiniMax-H3 serialized NVFP4 linear was not finalized: "
"process_weights_after_loading has not run")
return self._apply_finalized(layer, x, bias)
__all__ = [
"LANGUAGE_PROJECTIONS",
"NVFP4_TENSOR_SUFFIXES",
"MiniMaxH3SerializedNVFP4Config",
"MiniMaxH3SerializedNVFP4LinearMethod",
"NVFP4_GROUP_SIZE",
"NVFP4_SCALE_LAYOUT",
"nvfp4_packed_weight_shape",
"nvfp4_scale_shape",
"nvfp4_weight_global_scale",
"serialized_nvfp4_quantization_config",
"validate_nvfp4_geometry",
]
fastvideo/models/encoders/minimax_h3_qwen3_vl.py
+18 -3
View File
@@ -12,9 +12,11 @@ from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3
from fastvideo.distributed import get_tp_world_size
from fastvideo.layers.layernorm import RMSNorm
from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear
from fastvideo.layers.quantization.base_config import QuantizationConfig
from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
from fastvideo.models.encoders.base import TextEncoder
from fastvideo.models.encoders.minimax_h3_checkpoint_fp8 import MiniMaxH3SerializedFP8Config
from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4Config
from fastvideo.models.loader.weight_utils import default_weight_loader
@@ -503,14 +505,26 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder[torch.Tensor]):
"""H3 conditioner returning the unnormalized layer-50 hidden tensor."""
supports_hf_from_pretrained = False
supported_checkpoint_quantization_methods = frozenset({"fp8"})
# Serialized quantization contracts this conditioner can load, keyed by the
# ``quant_method`` a checkpoint declares. The loader gates on the key set
# before calling the factory, so both stay in step by construction.
_checkpoint_quantization_configs: dict[str, type[QuantizationConfig]] = {
MiniMaxH3SerializedFP8Config.get_name(): MiniMaxH3SerializedFP8Config,
MiniMaxH3SerializedNVFP4Config.get_name(): MiniMaxH3SerializedNVFP4Config,
}
supported_checkpoint_quantization_methods = frozenset(_checkpoint_quantization_configs)
@classmethod
def checkpoint_quantization_config_from_metadata(
cls,
metadata: dict[str, Any],
) -> MiniMaxH3SerializedFP8Config:
return MiniMaxH3SerializedFP8Config.from_config(metadata)
) -> QuantizationConfig:
quant_method = str(metadata.get("quant_method", "")).lower()
config_cls = cls._checkpoint_quantization_configs.get(quant_method)
if config_cls is None:
raise ValueError(f"MiniMax-H3 has no serialized {quant_method!r} text-encoder contract; "
f"supported: {sorted(cls._checkpoint_quantization_configs)}")
return config_cls.from_config(metadata)
def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None:
super().__init__(config)
@@ -756,4 +770,5 @@ EntryClass = MiniMaxH3Qwen3VLConditioner
__all__ = [
"MiniMaxH3Qwen3VLConditioner",
"MiniMaxH3SerializedFP8Config",
"MiniMaxH3SerializedNVFP4Config",
]
fastvideo/models/loader/component_loader.py
+11 -9
View File
@@ -450,9 +450,19 @@ class TextEncoderLoader(ComponentLoader):
self.counter_after_loading_weights - self.counter_before_loading_weights,
)
# Strict check for unquantized models and for serialized checkpoint
# schemes, which track every loaded tensor. It runs before the
# post-load hooks so a missing tensor is reported by name here
# instead of as one anonymous "not loaded" failure inside a hook.
weights_not_loaded = weights_to_load - loaded_weights
if weights_not_loaded and (model_config.quant_config is None or checkpoint_quant_config is not None):
raise ValueError("Following weights were not initialized from "
f"checkpoint: {weights_not_loaded}")
if checkpoint_quant_config is not None:
processed_linears = _process_quantized_text_encoder_weights(model, runtime_device)
logger.info("Validated %d serialized blockwise FP8 text-encoder linears", processed_linears)
logger.info("Validated %d serialized %s text-encoder linears", processed_linears,
checkpoint_quant_config.get_name())
# Explicitly move model to target device after loading weights
model = model.to(target_device)
@@ -492,14 +502,6 @@ class TextEncoderLoader(ComponentLoader):
fsdp_shard_conditions=model._fsdp_shard_conditions,
pin_cpu_memory=pin_cpu_memory,
)
# We only enable strict check for non-quantized models
# that have loaded weights tracking currently.
# if loaded_weights is not None:
weights_not_loaded = weights_to_load - loaded_weights
if weights_not_loaded and (model_config.quant_config is None or checkpoint_quant_config is not None):
raise ValueError("Following weights were not initialized from "
f"checkpoint: {weights_not_loaded}")
return model.eval()
fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py
+483
View File
@@ -0,0 +1,483 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import json
import os
import pytest
import torch
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29515")
import fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 as h3_nvfp4
from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (
MiniMaxH3Qwen3VLArchConfig,
MiniMaxH3Qwen3VLConfig,
)
from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear, UnquantizedLinearMethod
from fastvideo.layers.vocab_parallel_embedding import UnquantizedEmbeddingMethod, VocabParallelEmbedding
from fastvideo.models.encoders.base import TextEncoder
from fastvideo.models.encoders.minimax_h3_checkpoint_fp8 import MiniMaxH3SerializedFP8Config
from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import (
LANGUAGE_PROJECTIONS,
NVFP4_TENSOR_SUFFIXES,
MiniMaxH3SerializedNVFP4Config,
MiniMaxH3SerializedNVFP4LinearMethod,
serialized_nvfp4_quantization_config,
)
from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
from fastvideo.models.loader.text_encoder_quantization import (
_configure_text_encoder_quantization,
_process_quantized_text_encoder_weights,
_read_text_encoder_checkpoint_quantization_config,
)
LAYER_PREFIX = "minimax_h3_qwen3_vl.language_model.layers.0"
E4M3_ONE = 0x38
def _checkpoint_quantization_config(**overrides) -> dict:
config = serialized_nvfp4_quantization_config()
config.update(overrides)
return config
def _language_linear(config: MiniMaxH3SerializedNVFP4Config,
input_size: int = 128,
output_size: int = 256,
proj: str = "self_attn.q_proj",
bias: bool = False) -> ColumnParallelLinear:
return ColumnParallelLinear(
input_size=input_size,
output_size=output_size,
bias=bias,
quant_config=config,
prefix=f"{LAYER_PREFIX}.{proj}",
)
def _fill_loaded(layer: torch.nn.Module, global_scale: float = 2.0) -> None:
layer.weight_packed.data.fill_(0x11)
layer.weight_scale.data.fill_(E4M3_ONE)
layer.weight_global_scale.data.fill_(global_scale)
def test_h3_accepts_only_the_serialized_nvfp4_contract() -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
assert config.bf16_projections == ()
assert config.get_name() == "nvfp4"
assert config.get_supported_act_dtypes() == [torch.bfloat16]
with pytest.raises(ValueError, match="group_size=16"):
MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(group_size=32))
with pytest.raises(ValueError, match="scale_layout='128x4'"):
MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(scale_layout="linear"))
with pytest.raises(ValueError, match="dynamic activation"):
MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(activation_scheme="static"))
with pytest.raises(ValueError, match="E2M1"):
MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(fmt="e4m3"))
with pytest.raises(ValueError, match="vision stack"):
MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(modules_to_not_convert=["lm_head"]))
with pytest.raises(ValueError, match="quant_method 'fp8'"):
MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(quant_method="fp8"))
def test_h3_requires_every_layout_field_to_be_explicit() -> None:
bare = {"quant_method": "nvfp4", "modules_to_not_convert": ["model.visual", "lm_head"]}
with pytest.raises(ValueError, match="explicit quantization_config fields"):
MiniMaxH3SerializedNVFP4Config.from_config(bare)
def test_h3_keeps_bf16_per_projection_kind_only() -> None:
kept = MiniMaxH3SerializedNVFP4Config.from_config(
_checkpoint_quantization_config(modules_to_not_convert=["model.visual", "lm_head", "mlp.down_proj"]))
assert kept.bf16_projections == ("mlp.down_proj", )
assert kept.is_kept_bf16(f"{LAYER_PREFIX}.mlp.down_proj")
assert not kept.is_kept_bf16(f"{LAYER_PREFIX}.mlp.up_proj")
with pytest.raises(ValueError, match="per projection kind only"):
MiniMaxH3SerializedNVFP4Config.from_config(
_checkpoint_quantization_config(modules_to_not_convert=["model.visual", "language_model.layers.3"]))
with pytest.raises(ValueError, match="mlp.dwn_proj"):
MiniMaxH3SerializedNVFP4Config.from_config(
_checkpoint_quantization_config(modules_to_not_convert=["model.visual", "mlp.dwn_proj"]))
with pytest.raises(ValueError, match="non-empty strings"):
MiniMaxH3SerializedNVFP4Config.from_config(
_checkpoint_quantization_config(modules_to_not_convert=["model.visual", ""]))
def test_converter_metadata_round_trips_and_tolerates_producer_notes() -> None:
metadata = serialized_nvfp4_quantization_config(producer={"converter": "x.py", "flashinfer": "0.6.13rc2"})
assert MiniMaxH3SerializedNVFP4Config.from_config(metadata).bf16_projections == ()
assert metadata["producer"]["flashinfer"] == "0.6.13rc2"
mixed = serialized_nvfp4_quantization_config(keep_bf16=("mlp.down_proj", ))
assert mixed["modules_to_not_convert"] == ["model.visual", "lm_head", "mlp.down_proj"]
assert MiniMaxH3SerializedNVFP4Config.from_config(mixed).bf16_projections == ("mlp.down_proj", )
with pytest.raises(ValueError, match="unknown projection kinds"):
serialized_nvfp4_quantization_config(keep_bf16=("mlp.dwn_proj", ))
assert set(LANGUAGE_PROJECTIONS) == {
"self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj",
"mlp.up_proj", "mlp.down_proj"
}
def test_serialized_nvfp4_allocates_packed_weight_and_scales_without_a_bf16_weight(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
layer = _language_linear(config)
assert isinstance(layer.quant_method, MiniMaxH3SerializedNVFP4LinearMethod)
assert layer.weight is None
assert layer.weight_packed.dtype == torch.uint8
assert layer.weight_packed.shape == (256, 64)
assert layer.weight_scale.dtype == torch.uint8
assert layer.weight_scale.shape == (256, 8)
assert layer.weight_global_scale.dtype == torch.float32
assert layer.weight_global_scale.shape == (1, )
for name in ("weight_packed", "weight_scale", "weight_global_scale"):
param = getattr(layer, name)
assert callable(getattr(param, "weight_loader", None))
assert not hasattr(param, "output_dim") and not hasattr(param, "input_dim")
_fill_loaded(layer, global_scale=4.0)
packed_pointer = layer.weight_packed.data_ptr()
scale_pointer = layer.weight_scale.data_ptr()
layer.quant_method.process_weights_after_loading(layer)
assert layer.weight_packed.data_ptr() == packed_pointer
assert layer.weight_scale.data_ptr() == scale_pointer
assert layer.weight is None
assert layer._nvfp4_alpha.dtype == torch.float32
assert layer._nvfp4_alpha.item() == pytest.approx(0.25)
assert layer._nvfp4_x_global_scale.item() == 1.0
def test_serialized_nvfp4_finalization_rejects_tensors_the_checkpoint_never_filled(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
layer = _language_linear(config)
with pytest.raises(ValueError, match="weight_global_scale was not loaded"):
layer.quant_method.process_weights_after_loading(layer)
layer.weight_global_scale.data.fill_(2.0)
with pytest.raises(ValueError, match="weight_scale was not loaded"):
layer.quant_method.process_weights_after_loading(layer)
layer.weight_scale.data.fill_(E4M3_ONE)
layer.quant_method.process_weights_after_loading(layer)
assert layer._nvfp4_alpha.item() == pytest.approx(0.5)
def test_serialized_nvfp4_tensors_arrive_through_the_layer_weight_loader(distributed_setup) -> None:
"""Drive the three tensors the way ``load_weights`` does, through the layer's own loader."""
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
column = _language_linear(config, proj="self_attn.q_proj")
row = RowParallelLinear(
input_size=128,
output_size=256,
bias=False,
quant_config=config,
prefix=f"{LAYER_PREFIX}.self_attn.o_proj",
)
for layer in (column, row):
assert isinstance(layer.quant_method, MiniMaxH3SerializedNVFP4LinearMethod)
packed = torch.full((256, 64), 0x22, dtype=torch.uint8)
scale = torch.full((256, 8), E4M3_ONE, dtype=torch.uint8)
layer.weight_packed.weight_loader(layer.weight_packed, packed)
layer.weight_scale.weight_loader(layer.weight_scale, scale)
layer.weight_global_scale.weight_loader(layer.weight_global_scale, torch.tensor(8.0))
layer.quant_method.process_weights_after_loading(layer)
assert torch.equal(layer.weight_packed.data, packed)
assert layer._nvfp4_alpha.item() == pytest.approx(0.125)
with pytest.raises(AssertionError):
column.weight_packed.weight_loader(column.weight_packed, torch.zeros(256, 128, dtype=torch.uint8))
# A float8 scale tensor would be value-cast into uint8 by a plain copy; the loader refuses it.
with pytest.raises(ValueError, match="declared dtype"):
column.weight_scale.weight_loader(column.weight_scale, torch.zeros(256, 8, dtype=torch.float8_e4m3fn))
with pytest.raises(ValueError, match="declared dtype"):
column.weight_global_scale.weight_loader(column.weight_global_scale, torch.tensor(8.0, dtype=torch.float16))
def test_serialized_nvfp4_quantizes_only_language_linears(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
visual_linear = ColumnParallelLinear(
input_size=128,
output_size=128,
bias=False,
quant_config=config,
prefix="minimax_h3_qwen3_vl.visual.blocks.0.attn.proj",
)
embedding = VocabParallelEmbedding(
num_embeddings=128,
embedding_dim=128,
org_num_embeddings=128,
quant_config=config,
prefix="minimax_h3_qwen3_vl.language_model.embed_tokens",
)
assert isinstance(visual_linear.quant_method, UnquantizedLinearMethod)
assert visual_linear.weight.dtype == torch.get_default_dtype()
assert isinstance(embedding.quant_method, UnquantizedEmbeddingMethod)
assert embedding.weight.dtype == torch.get_default_dtype()
def test_kept_bf16_projection_builds_a_plain_linear(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(
_checkpoint_quantization_config(modules_to_not_convert=["model.visual", "lm_head", "mlp.down_proj"]))
down = _language_linear(config, input_size=128, output_size=128, proj="mlp.down_proj")
up = _language_linear(config, input_size=128, output_size=128, proj="mlp.up_proj")
assert isinstance(down.quant_method, UnquantizedLinearMethod)
assert down.weight is not None and down.weight.shape == (128, 128)
assert not hasattr(down, "weight_packed")
assert isinstance(up.quant_method, MiniMaxH3SerializedNVFP4LinearMethod)
assert up.weight is None
def test_serialized_nvfp4_rejects_geometry_outside_flashinfer_tiles(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
with pytest.raises(ValueError, match="input_size divisible by 64"):
_language_linear(config, input_size=96, output_size=256)
with pytest.raises(ValueError, match="output_size divisible by 128"):
_language_linear(config, input_size=128, output_size=200)
def test_serialized_nvfp4_cpu_execution_fails_closed(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
layer = _language_linear(config, input_size=128, output_size=128)
_fill_loaded(layer)
layer.quant_method.process_weights_after_loading(layer)
with pytest.raises(RuntimeError, match="requires CUDA"):
layer(torch.zeros(2, 128, dtype=torch.bfloat16))
def test_runtime_preflight_gates_device_capability_parallelism_and_flashinfer(monkeypatch) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
with pytest.raises(RuntimeError, match="requires a CUDA device"):
config.validate_runtime(torch.device("cpu"))
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9))
with pytest.raises(RuntimeError, match="sm100 or newer"):
config.validate_runtime(torch.device("cuda"))
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (12, 1))
monkeypatch.setattr(h3_nvfp4, "get_tp_world_size", lambda: 2)
with pytest.raises(NotImplementedError, match="single GPU"):
config.validate_runtime(torch.device("cuda"))
monkeypatch.setattr(h3_nvfp4, "get_tp_world_size", lambda: 1)
def missing_flashinfer():
raise ImportError("NVFP4 quantization requires flashinfer")
monkeypatch.setattr(h3_nvfp4, "_require_flashinfer", missing_flashinfer)
with pytest.raises(ImportError, match="requires flashinfer"):
config.validate_runtime(torch.device("cuda"))
class LayoutWithoutSwizzle:
pass
monkeypatch.setattr(h3_nvfp4, "_require_flashinfer", lambda: (LayoutWithoutSwizzle, None, None))
with pytest.raises(RuntimeError, match="layout_128x4"):
config.validate_runtime(torch.device("cuda"))
class Layout:
layout_128x4 = 1
registered = []
monkeypatch.setattr(h3_nvfp4, "_require_flashinfer", lambda: (Layout, None, None))
monkeypatch.setattr(h3_nvfp4, "_register_ops_once", lambda: registered.append(True))
config.validate_runtime(torch.device("cuda"))
assert registered == [True]
def test_loader_selects_nvfp4_from_checkpoint_metadata(tmp_path) -> None:
checkpoint_config = _checkpoint_quantization_config()
(tmp_path / "config.json").write_text(json.dumps({"quantization_config": checkpoint_config}), encoding="utf-8")
assert _read_text_encoder_checkpoint_quantization_config(str(tmp_path)) == checkpoint_config
model_config = MiniMaxH3Qwen3VLConfig()
quant_config = _configure_text_encoder_quantization(model_config, MiniMaxH3Qwen3VLConditioner, str(tmp_path))
assert isinstance(quant_config, MiniMaxH3SerializedNVFP4Config)
assert model_config.quant_config is quant_config
with pytest.raises(ValueError, match="does not support serialized 'nvfp4'"):
_configure_text_encoder_quantization(MiniMaxH3Qwen3VLConfig(), TextEncoder, str(tmp_path))
def test_conditioner_routes_each_serialized_scheme_to_its_config() -> None:
fp8_metadata = {
"quant_method": "fp8",
"activation_scheme": "dynamic",
"fmt": "e4m3",
"weight_block_size": [128, 128],
"modules_to_not_convert": ["model.visual", "lm_head"],
}
assert MiniMaxH3Qwen3VLConditioner.supported_checkpoint_quantization_methods == frozenset({"fp8", "nvfp4"})
assert isinstance(MiniMaxH3Qwen3VLConditioner.checkpoint_quantization_config_from_metadata(fp8_metadata),
MiniMaxH3SerializedFP8Config)
assert isinstance(
MiniMaxH3Qwen3VLConditioner.checkpoint_quantization_config_from_metadata(_checkpoint_quantization_config()),
MiniMaxH3SerializedNVFP4Config)
with pytest.raises(ValueError, match="no serialized 'int4'"):
MiniMaxH3Qwen3VLConditioner.checkpoint_quantization_config_from_metadata({"quant_method": "int4"})
def test_post_load_processing_visits_only_serialized_nvfp4_linears(distributed_setup) -> None:
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
quantized = _language_linear(config, input_size=128, output_size=128)
plain = ColumnParallelLinear(input_size=128, output_size=128, bias=False, prefix="plain")
_fill_loaded(quantized)
model = torch.nn.ModuleList([quantized, plain])
assert _process_quantized_text_encoder_weights(model, torch.device("cpu")) == 1
assert quantized.weight_packed.device.type == "cpu"
assert quantized._nvfp4_alpha.device.type == "cpu"
assert plain.weight.device.type == "cpu"
def test_nvfp4_linear_hands_mm_fp4_transposed_operands_and_a_bf16_output(monkeypatch) -> None:
x_fp4 = torch.zeros(4, 64, dtype=torch.uint8)
x_scale = torch.zeros(128, 8, dtype=torch.uint8)
weight_packed = torch.zeros(256, 64, dtype=torch.uint8)
weight_scale = torch.zeros(256, 8, dtype=torch.uint8)
alpha = torch.tensor(0.5, dtype=torch.float32)
receipt: dict[str, object] = {}
def fake_mm_fp4(a, b, a_scale, b_scale, alpha_arg, out_dtype, out, **kwargs):
receipt.update(a=a, b=b, a_scale=a_scale, b_scale=b_scale, alpha=alpha_arg, out_dtype=out_dtype, out=out,
kwargs=kwargs)
return torch.zeros(a.shape[0], b.shape[1], dtype=out_dtype)
monkeypatch.setattr(h3_nvfp4, "_mm_fp4", fake_mm_fp4)
output = h3_nvfp4._nvfp4_linear(x_fp4, x_scale, weight_packed, weight_scale, alpha)
assert output.shape == (4, 256)
assert output.dtype == torch.bfloat16
assert receipt["a"] is x_fp4
assert receipt["a_scale"] is x_scale
assert receipt["b"].shape == (64, 256)
assert receipt["b"].data_ptr() == weight_packed.data_ptr()
assert receipt["b_scale"].shape == (8, 256)
assert receipt["b_scale"].data_ptr() == weight_scale.data_ptr()
assert receipt["alpha"] is alpha
assert receipt["out_dtype"] == torch.bfloat16
assert receipt["out"] is None
assert receipt["kwargs"] == {"backend": "auto"}
def test_apply_flattens_quantizes_and_restores_the_leading_dims(distributed_setup, monkeypatch) -> None:
"""The finalized forward on a 3-D fp32 input: coerce to bf16, quantize the 2-D view with the
unit global scale, GEMM with the layer's alpha, add the bias, restore [batch, seq, out]."""
config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config())
layer = _language_linear(config, bias=True)
_fill_loaded(layer, global_scale=4.0)
layer.quant_method.process_weights_after_loading(layer)
layer.bias.data.fill_(3.0)
receipt: dict[str, object] = {}
def fake_quantize(x_2d, global_scale):
receipt.update(x_2d=x_2d, global_scale=global_scale)
return torch.zeros(x_2d.shape[0], 64, dtype=torch.uint8), torch.zeros(128, 8, dtype=torch.uint8)
def fake_linear(x_fp4, x_scale, weight_packed, weight_scale, alpha):
receipt.update(alpha=alpha, weight_packed=weight_packed, weight_scale=weight_scale)
return torch.zeros(x_fp4.shape[0], weight_packed.shape[0], dtype=torch.bfloat16)
monkeypatch.setattr(h3_nvfp4, "_quantize_activation_nvfp4", fake_quantize)
monkeypatch.setattr(h3_nvfp4, "_nvfp4_linear", fake_linear)
output = MiniMaxH3SerializedNVFP4LinearMethod._apply_finalized(layer, torch.ones(1, 5, 128), layer.bias)
assert receipt["x_2d"].shape == (5, 128)
assert receipt["x_2d"].dtype == torch.bfloat16
assert receipt["global_scale"] is layer._nvfp4_x_global_scale
assert receipt["alpha"] is layer._nvfp4_alpha
assert receipt["weight_packed"] is layer.weight_packed
assert output.shape == (1, 5, 256)
assert output.dtype == torch.bfloat16
assert torch.equal(output, torch.full((1, 5, 256), 3.0, dtype=torch.bfloat16))
def _tiny_conditioner_config(keep_bf16: tuple[str, ...] = ("mlp.down_proj", )) -> MiniMaxH3Qwen3VLConfig:
"""One language layer whose linears satisfy the FlashInfer tile geometry, and a small vision tower."""
config = MiniMaxH3Qwen3VLConfig()
config.arch_config = MiniMaxH3Qwen3VLArchConfig(
vocab_size=64,
hidden_size=128,
intermediate_size=256,
num_hidden_layers=2,
output_hidden_state_index=1,
num_hidden_layers_override=1,
num_attention_heads=1,
num_key_value_heads=1,
head_dim=128,
rope_scaling={
"mrope_interleaved": True,
"mrope_section": [32, 16, 16],
"rope_type": "default",
},
vision_depth=1,
vision_hidden_size=64,
vision_intermediate_size=128,
vision_num_heads=1,
vision_deepstack_visual_indexes=(),
vision_out_hidden_size=128,
)
config.quant_config = MiniMaxH3SerializedNVFP4Config.from_config(
_checkpoint_quantization_config(modules_to_not_convert=["model.visual", "lm_head", *keep_bf16]))
return config
def test_conditioner_loads_converter_named_tensors_end_to_end(distributed_setup) -> None:
"""The real chain: a conditioner built with the NVFP4 config, checkpoint keys spelled the way the
converter writes them, ``load_weights``, the strict missing-tensor check, then the post-load hook."""
model = MiniMaxH3Qwen3VLConditioner(_tiny_conditioner_config())
layer = model.language_model.layers[0]
assert isinstance(layer.self_attn.q_proj.quant_method, MiniMaxH3SerializedNVFP4LinearMethod)
assert isinstance(layer.self_attn.o_proj.quant_method, MiniMaxH3SerializedNVFP4LinearMethod)
assert isinstance(layer.mlp.down_proj.quant_method, UnquantizedLinearMethod)
assert layer.self_attn.q_proj.weight is None and layer.mlp.down_proj.weight is not None
packed_name, scale_name, global_scale_name = NVFP4_TENSOR_SUFFIXES
checkpoint: dict[str, torch.Tensor] = {}
for name, param in model.named_parameters():
if name.endswith("." + packed_name):
tensor = torch.full_like(param, 0x11)
elif name.endswith("." + scale_name):
tensor = torch.full_like(param, E4M3_ONE)
elif name.endswith("." + global_scale_name):
tensor = torch.full_like(param, 2.0)
else:
tensor = torch.zeros_like(param)
checkpoint["model." + name] = tensor
expected = {name for name, _ in model.named_parameters()}
assert sum(name.endswith("." + packed_name) for name in expected) == 6
loaded = model.load_weights(iter(checkpoint.items()))
assert loaded == expected
assert _process_quantized_text_encoder_weights(model, torch.device("cpu")) == 6
assert layer.self_attn.q_proj._nvfp4_alpha.item() == pytest.approx(0.5)
assert layer.mlp.up_proj._nvfp4_alpha.item() == pytest.approx(0.5)
# A tensor the checkpoint lacks is reported by name through the loaded set, which is what the
# loader's strict check subtracts, and the hook then refuses the unfilled layer.
missing = f"model.language_model.layers.0.mlp.up_proj.{scale_name}"
layer.mlp.up_proj.weight_scale.data.fill_(0xFF)
partial = {key: value for key, value in checkpoint.items() if key != missing}
loaded = model.load_weights(iter(partial.items()))
assert expected - loaded == {missing[len("model."):]}
with pytest.raises(ValueError, match="weight_scale was not loaded"):
_process_quantized_text_encoder_weights(model, torch.device("cpu"))
wrong_dtype = dict(checkpoint)
wrong_dtype[missing] = wrong_dtype[missing].to(torch.float8_e4m3fn)
with pytest.raises(ValueError, match="declared dtype"):
model.load_weights(iter(wrong_dtype.items()))
stray = dict(checkpoint)
stray["model.language_model.layers.0.self_attn.q_proj.weight"] = torch.zeros(128, 128)
with pytest.raises(ValueError, match="Unexpected"):
model.load_weights(iter(stray.items()))