From 3a3ad3d20941a41826a41a72e2f3befaa003ca86 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:31:59 -0700 Subject: [PATCH] [feat]: NVFP4 text encoder for MiniMax H3 (#1837) --- fastvideo/layers/quantization/base_config.py | 9 + .../encoders/minimax_h3_checkpoint_nvfp4.py | 425 +++++++++++++++ .../models/encoders/minimax_h3_qwen3_vl.py | 21 +- fastvideo/models/loader/component_loader.py | 20 +- ...st_minimax_h3_qwen3_vl_checkpoint_nvfp4.py | 483 ++++++++++++++++++ 5 files changed, 946 insertions(+), 12 deletions(-) create mode 100644 fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py create mode 100644 fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py diff --git a/fastvideo/layers/quantization/base_config.py b/fastvideo/layers/quantization/base_config.py index c7dc2754a..c9ef3f6e4 100644 --- a/fastvideo/layers/quantization/base_config.py +++ b/fastvideo/layers/quantization/base_config.py @@ -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 diff --git a/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py b/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py new file mode 100644 index 000000000..0c96e636c --- /dev/null +++ b/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py @@ -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``:: + + .weight_packed uint8 [out, in // 2] two E2M1 values per byte + .weight_scale uint8 [out, in // 16] one E4M3 scale per 16 values, + FlashInfer 128x4 swizzled layout + .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", +] diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py index a67273195..9f4b5e8ea 100644 --- a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -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", ] diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 7f850744e..148a6d1bb 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -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() diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py new file mode 100644 index 000000000..74be1a9e8 --- /dev/null +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py @@ -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()))