diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..dee250ff47d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path. - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. @@ -40,6 +41,7 @@ Changelog **Bug Fixes** +- Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. diff --git a/examples/_shard_cast_utils.py b/examples/_shard_cast_utils.py new file mode 100644 index 00000000000..0f94707b625 --- /dev/null +++ b/examples/_shard_cast_utils.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared primitives for streaming MXFP4-to-NVFP4 checkpoint conversion.""" + +from __future__ import annotations + +import errno +import os +import shutil +from collections import defaultdict +from pathlib import Path +from typing import TYPE_CHECKING + +import torch + +from modelopt.torch.quantization.qtensor import MXFP4QTensor, NVFP4QTensor +from modelopt.torch.quantization.utils.numeric_utils import ( + E2M1_MAX, + E4M3_KMAX, + E4M3_KMIN, + E4M3_MAX, + E8M0_BIAS, + mxfp4_to_nvfp4_global_amax, + mxfp4_to_nvfp4_per_block_amax, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Collection + +__all__ = [ + "build_w13_amax_overrides", + "build_w13_kmax_overrides", + "dequantize_mxfp4_to_bf16", + "link_aux_files", + "link_or_copy", + "log", + "mxfp4_kmax", + "prepare_output_dir", + "quantize_mxfp4_to_nvfp4", + "quantize_mxfp4_to_nvfp4_lossless", + "validate_paths", +] + +_MXFP4_BLOCK = 32 +_MXFP4_BYTES_PER_BLOCK = 16 +_NVFP4_BLOCK = 16 + + +def dequantize_mxfp4_to_bf16( + mxfp4_weight: torch.Tensor, mxfp4_scale: torch.Tensor, device: str +) -> torch.Tensor: + """Dequantize packed MXFP4 weights and E8M0 scales to BF16.""" + packed = mxfp4_weight.to(device).contiguous().view(torch.uint8) + scale = mxfp4_scale.to(device).contiguous().view(torch.uint8) + original_shape = torch.Size((*packed.shape[:-1], packed.shape[-1] * 2)) + assert packed.shape[:-1] == scale.shape[:-1] and ( + 2 * packed.shape[-1] == scale.shape[-1] * _MXFP4_BLOCK + ), f"Incompatible MXFP4 shapes: weight {tuple(packed.shape)} vs scale {tuple(scale.shape)}" + return MXFP4QTensor(original_shape, torch.bfloat16, packed).dequantize( + dtype=torch.bfloat16, + scale=scale, + block_sizes=[_MXFP4_BLOCK], + ) + + +def _w13_pairs(expert_bases: list[str]) -> list[tuple[str, str]]: + groups: dict[str, dict[str, str]] = defaultdict(dict) + for base in expert_bases: + prefix, proj = base.rsplit(".", 1) + if proj in {"w1", "w3"}: + groups[prefix][proj] = base + + pairs: list[tuple[str, str]] = [] + for prefix, paths in groups.items(): + if "w1" not in paths or "w3" not in paths: + raise RuntimeError( + "w1/w3 of one expert are split across shards, so they cannot share " + f"scale_2 for the fused GEMM1: {prefix}" + ) + pairs.append((paths["w1"], paths["w3"])) + return pairs + + +def build_w13_kmax_overrides( + expert_bases: list[str], + get_scale: Callable[[str], torch.Tensor], + device: str, +) -> dict[str, int]: + """Return one shared E8M0 maximum exponent for each fused w1/w3 pair.""" + overrides: dict[str, int] = {} + for w1, w3 in _w13_pairs(expert_bases): + k1 = mxfp4_kmax(get_scale(w1), device) + k3 = mxfp4_kmax(get_scale(w3), device) + overrides[w1] = overrides[w3] = max(k1, k3) + return overrides + + +def build_w13_amax_overrides( + expert_bases: list[str], + get_amax: Callable[[str], torch.Tensor], +) -> dict[str, torch.Tensor]: + """Return one shared weight amax for each fused w1/w3 pair.""" + overrides: dict[str, torch.Tensor] = {} + for w1, w3 in _w13_pairs(expert_bases): + shared = torch.maximum(get_amax(w1).reshape(()), get_amax(w3).reshape(())) + overrides[w1] = overrides[w3] = shared + return overrides + + +def mxfp4_kmax(mxfp4_scale: torch.Tensor, device: str = "cpu") -> int: + """Return the largest non-zero unbiased exponent in an E8M0 scale tensor.""" + e8m0 = mxfp4_scale.to(device).contiguous().view(torch.uint8) + return mxfp4_to_nvfp4_global_amax(e8m0)[1]["k_max"] + + +def quantize_mxfp4_to_nvfp4( + mxfp4_weight: torch.Tensor, + mxfp4_scale: torch.Tensor, + weight_amax: torch.Tensor | None, + device: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]: + """Dequantize MXFP4 and requantize it to NVFP4 using an optional global amax.""" + bf16 = dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device) + synthesized = weight_amax is None + if weight_amax is None: + weight_amax = bf16.abs().max() + weight_scale_2 = (weight_amax.to(device).float() / (E2M1_MAX * E4M3_MAX)).reshape(()) + q_tensor, weight_scale, _ = NVFP4QTensor.quantize( + bf16, _NVFP4_BLOCK, None, weight_scale_2, try_tensorrt=False + ) + return q_tensor._quantized_data, weight_scale, weight_scale_2, synthesized + + +def quantize_mxfp4_to_nvfp4_lossless( + mxfp4_weight: torch.Tensor, + mxfp4_scale: torch.Tensor, + k_max: int, + device: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Closed-form MXFP4-to-NVFP4 cast with lossless-block accounting.""" + bf16 = dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device) + e8m0 = mxfp4_scale.to(bf16.device).contiguous().view(torch.uint8) + packed = mxfp4_weight.to(bf16.device).contiguous().view(torch.uint8) + blocks = packed.view(*packed.shape[:-1], e8m0.shape[-1], _MXFP4_BYTES_PER_BLOCK) + per_block_amax = mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) + + weight_scale_2 = torch.tensor( + 2.0 ** (k_max - E4M3_KMAX), dtype=torch.float32, device=bf16.device + ).reshape(()) + per_block_scale = ( + (per_block_amax / (E2M1_MAX * weight_scale_2)) + .clamp(min=2**E4M3_KMIN, max=E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + + k = e8m0.to(torch.int32) - E8M0_BIAS + lossless = (k >= (k_max - (E4M3_KMAX - E4M3_KMIN))) | (e8m0 == 0) + n_blocks = k.numel() + n_lossless = int(lossless.sum().item()) + + q_tensor, weight_scale, _ = NVFP4QTensor.quantize( + bf16, _NVFP4_BLOCK, per_block_scale, weight_scale_2, try_tensorrt=False + ) + return q_tensor._quantized_data, weight_scale, weight_scale_2, n_blocks, n_lossless + + +def link_or_copy(src: Path, dst: Path) -> None: + """Hard-link a file, copying when the filesystem cannot create the link.""" + try: + os.link(src, dst) + except OSError as exc: + copy_errnos = { + errno.EXDEV, + errno.EPERM, + errno.EACCES, + errno.EMLINK, + getattr(errno, "EOPNOTSUPP", errno.EXDEV), + getattr(errno, "ENOTSUP", errno.EXDEV), + } + if exc.errno not in copy_errnos: + raise + shutil.copy2(src, dst) + + +def link_aux_files( + src_dir: Path, + dst_dir: Path, + *, + skip_top_level: Collection[str] = (), + skip_dir_names: Collection[str] = (), + skip_file: Callable[[Path], bool] | None = None, +) -> None: + """Recursively link checkpoint sidecars while applying model-specific skips.""" + for root, dirs, files in os.walk(src_dir): + rel = Path(root).relative_to(src_dir) + at_top_level = rel == Path(".") + dirs[:] = [ + name + for name in dirs + if name not in skip_dir_names and not (at_top_level and name in skip_top_level) + ] + (dst_dir / rel).mkdir(parents=True, exist_ok=True) + for name in files: + relative_path = rel / name + if at_top_level and name in skip_top_level: + continue + if skip_file is not None and skip_file(relative_path): + continue + src = src_dir / relative_path + dst = dst_dir / relative_path + if dst.exists(): + dst.unlink() + link_or_copy(src, dst) + + +def log(message: str) -> None: + """Print a checkpoint-conversion progress message immediately.""" + print(message, flush=True) + + +def validate_paths(source_ckpt: Path, output_ckpt: Path) -> None: + """Reject overlapping source and output checkpoint directories.""" + source_resolved = source_ckpt.resolve() + output_resolved = output_ckpt.resolve() + if ( + output_resolved == source_resolved + or source_resolved in output_resolved.parents + or output_resolved in source_resolved.parents + ): + raise ValueError( + "--source_ckpt and --output_ckpt must be disjoint directories; " + f"got source={source_ckpt}, output={output_ckpt}" + ) + + +def prepare_output_dir(output_ckpt: Path, overwrite: bool) -> None: + """Create an empty output directory, replacing its contents when allowed.""" + if output_ckpt.exists(): + if not output_ckpt.is_dir(): + raise ValueError(f"--output_ckpt exists and is not a directory: {output_ckpt}") + if any(output_ckpt.iterdir()): + if not overwrite: + raise ValueError( + f"--output_ckpt is not empty: {output_ckpt}; pass --overwrite to replace it" + ) + for item in output_ckpt.iterdir(): + if item.is_dir() and not item.is_symlink(): + shutil.rmtree(item) + else: + item.unlink() + output_ckpt.mkdir(parents=True, exist_ok=True) diff --git a/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py b/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py index be1e41c842e..a5bf206b828 100644 --- a/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py +++ b/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py @@ -90,11 +90,9 @@ from __future__ import annotations import argparse -import errno import json -import os import re -import shutil +import sys from collections import defaultdict from pathlib import Path from typing import Any @@ -103,18 +101,22 @@ from safetensors import safe_open from safetensors.torch import save_file -from modelopt.torch.quantization.qtensor import MXFP4QTensor, NVFP4QTensor - -# Closed-form MXFP4 -> NVFP4 numerics shared with the GPT-OSS cast (PR #1372). -from modelopt.torch.quantization.utils.numeric_utils import ( - E2M1_MAX, - E4M3_KMAX, - E4M3_KMIN, - E4M3_MAX, - E8M0_BIAS, - mxfp4_to_nvfp4_global_amax, - mxfp4_to_nvfp4_per_block_amax, +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from examples._shard_cast_utils import ( # noqa: E402 + build_w13_amax_overrides, + build_w13_kmax_overrides, + dequantize_mxfp4_to_bf16, + link_aux_files, + mxfp4_kmax, + prepare_output_dir, + quantize_mxfp4_to_nvfp4, + quantize_mxfp4_to_nvfp4_lossless, + validate_paths, ) +from examples._shard_cast_utils import log as _log # noqa: E402 # Routed-expert weights in regular MoE layers. MTP experts remain in source format. _EXPERT_WEIGHT_RE = re.compile(r"^layers\.\d+\.ffn\.experts\.\d+\.w[123]\.weight$") @@ -129,10 +131,6 @@ _HF_SHARD_RE = re.compile(r"^model-(?P\d+)-of-(?P\d+)\.safetensors$") -def _log(msg: str) -> None: - print(msg, flush=True) - - def _amax_to_nvfp4_scale_2(amax: torch.Tensor) -> torch.Tensor: """``amax / (fp4_max * fp8_max) = amax / (6 * 448)``; returns a 0-d fp32 scalar.""" return (amax.float() / (6.0 * 448.0)).to(torch.float32).reshape(()) @@ -215,139 +213,10 @@ def _lookup_amax( return amax.get(f"{expert_path}_{which}_quantizer._amax") -def _dequantize_mxfp4_to_bf16( - mxfp4_weight: torch.Tensor, mxfp4_scale: torch.Tensor, device: str -) -> torch.Tensor: - block_size = 32 - packed = mxfp4_weight.to(device).contiguous().view(torch.uint8) - scale = mxfp4_scale.to(device).contiguous().view(torch.uint8) - original_shape = torch.Size((*packed.shape[:-1], packed.shape[-1] * 2)) - assert packed.shape[:-1] == scale.shape[:-1] and ( - 2 * packed.shape[-1] == scale.shape[-1] * block_size - ), f"Incompatible MXFP4 shapes: weight {tuple(packed.shape)} vs scale {tuple(scale.shape)}" - return MXFP4QTensor(original_shape, torch.bfloat16, packed).dequantize( - dtype=torch.bfloat16, - scale=scale, - block_sizes=[block_size], - ) - - def _synthesize_weight_amax( mxfp4_weight: torch.Tensor, mxfp4_scale: torch.Tensor, device: str ) -> torch.Tensor: - return _dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device).abs().max().cpu() - - -def _quantize_weight_nvfp4( - mxfp4_weight: torch.Tensor, - mxfp4_scale: torch.Tensor, - weight_amax: torch.Tensor | None, - device: str, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]: - """MXFP4 + UE8M0 → BF16 → NVFP4 packed. Synthesizes ``weight_amax`` from - the dequantized BF16 tensor when ``None`` is passed.""" - bf16 = _dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device) - synthesized = weight_amax is None - if synthesized: - weight_amax = bf16.abs().max() - assert weight_amax is not None - weight_scale_2 = _amax_to_nvfp4_scale_2(weight_amax.to(device)) - q_tensor, weight_scale, _ = NVFP4QTensor.quantize( - bf16, 16, None, weight_scale_2, try_tensorrt=False - ) - return q_tensor._quantized_data, weight_scale, weight_scale_2, synthesized - - -# --------------------------------------------------------------------------- -# Lossless MXFP4 -> NVFP4 weight cast (``--cast_mxfp4_to_nvfp4``). -# -# NVFP4 uses the same E2M1 nibble grid as MXFP4 with 16-element blocks and a -# two-level scale ``per_block_scale (E4M3) * scale_2 (fp32)``. Pinning -# ``scale_2 = 2^m`` (``m = k_max - 8``) and ``per_block_scale = 2^(k_j - m)`` -# makes ``per_block_scale * scale_2 = 2^k_j`` exactly, so each NVFP4 nibble -# equals the source MXFP4 nibble verbatim — bit-exact for every block whose -# ``k_j`` lands in E4M3's window (``k_max - k_j <= 17``). The closed-form -# per-block amax and the format constants are reused from the GPT-OSS cast -# (``cast_mxfp4_to_nvfp4``, PR #1372); the V4 twist is that w1/w3 share one -# ``scale_2`` (fused GEMM1), so ``k_max`` is taken over both projections. -# --------------------------------------------------------------------------- -_NVFP4_BLOCK = 16 # NVFP4 block size (elements) -_MXFP4_BYTES_PER_BLOCK = 16 # 32 E2M1 nibbles packed 2-per-byte - - -def _kmax_from_mxfp4_scale(mxfp4_scale: torch.Tensor, device: str = "cpu") -> int: - """Largest non-zero E8M0 exponent ``k_j = e8m0 - 127`` (0 if all-zero). - - Delegates to the GPT-OSS cast's ``k_max`` logic, which excludes the - all-zero sentinel (``e8m0 == 0`` => ``k == -127``). - """ - e8m0 = mxfp4_scale.to(device).contiguous().view(torch.uint8) - return mxfp4_to_nvfp4_global_amax(e8m0)[1]["k_max"] - - -def _build_w13_kmax_overrides(f, expert_weight_keys: list[str], device: str) -> dict[str, int]: - """Shared ``k_max`` per w1/w3 pair so the fused GEMM1 gets one ``scale_2``.""" - groups: dict[str, dict[str, str]] = defaultdict(dict) - for key in expert_weight_keys: - expert_path = key[: -len(".weight")] - base, proj = expert_path.rsplit(".", 1) - if proj in {"w1", "w3"}: - groups[base][proj] = expert_path - - overrides: dict[str, int] = {} - for paths in groups.values(): - if "w1" not in paths or "w3" not in paths: - continue - k1 = _kmax_from_mxfp4_scale(f.get_tensor(paths["w1"] + ".scale"), device) - k3 = _kmax_from_mxfp4_scale(f.get_tensor(paths["w3"] + ".scale"), device) - shared = max(k1, k3) - overrides[paths["w1"]] = shared - overrides[paths["w3"]] = shared - return overrides - - -def _quantize_weight_nvfp4_lossless( - mxfp4_weight: torch.Tensor, - mxfp4_scale: torch.Tensor, - k_max: int, - device: str, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: - """Closed-form bit-exact MXFP4 -> NVFP4 weight conversion. - - Pins ``scale_2 = 2^(k_max - 8)`` and the per-block E4M3 scale to - ``2^(k_j - m)`` so the NVFP4 nibbles equal the source MXFP4 nibbles for - every in-range block. ``k_max`` is shared across w1/w3 (fused GEMM1), so it - is passed in rather than derived per tensor. The closed-form per-block amax - (``6 * 2^k_j`` in range, data-derived out of range) is independent of - ``k_max``, so we reuse the GPT-OSS helper directly. Returns - ``(packed, weight_scale, weight_scale_2, n_blocks, n_lossless)``. - """ - bf16 = _dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device) - e8m0 = mxfp4_scale.to(bf16.device).contiguous().view(torch.uint8) # (out, nblk32) - packed = mxfp4_weight.to(bf16.device).contiguous().view(torch.uint8) # (out, nblk32*16) - blocks = packed.view(*packed.shape[:-1], e8m0.shape[-1], _MXFP4_BYTES_PER_BLOCK) - per_block_amax = mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) # (out, nblk16) fp32 - - m = k_max - E4M3_KMAX - weight_scale_2 = torch.tensor(2.0**m, dtype=torch.float32, device=bf16.device).reshape(()) - per_block_scale = ( - (per_block_amax / (E2M1_MAX * weight_scale_2)) - .clamp(min=2**-9, max=E4M3_MAX) - .to(torch.float8_e4m3fn) - ) - - # Lossless accounting against the (possibly shared) k_max. A block is lossy - # only if k_max - k_j > 17; all-zero blocks (e8m0 == 0) reconstruct to 0 - # regardless of scale and so are always lossless. - k = e8m0.to(torch.int32) - E8M0_BIAS - lossless = (k >= (k_max - (E4M3_KMAX - E4M3_KMIN))) | (e8m0 == 0) - n_blocks = k.numel() - n_lossless = int(lossless.sum().item()) - - q_tensor, weight_scale, _ = NVFP4QTensor.quantize( - bf16, _NVFP4_BLOCK, per_block_scale, weight_scale_2, try_tensorrt=False - ) - return q_tensor._quantized_data, weight_scale, weight_scale_2, n_blocks, n_lossless + return dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device).abs().max().cpu() def _build_w13_weight_amax_overrides( @@ -357,36 +226,21 @@ def _build_w13_weight_amax_overrides( device: str, ) -> tuple[dict[str, torch.Tensor], set[str]]: """Return shared w1/w3 amax overrides so fused GEMM1 has one scale.""" - groups: dict[str, dict[str, str]] = defaultdict(dict) - for key in expert_weight_keys: - expert_path = key[: -len(".weight")] - base, proj = expert_path.rsplit(".", 1) - if proj in {"w1", "w3"}: - groups[base][proj] = expert_path - - overrides: dict[str, torch.Tensor] = {} + expert_bases = [key[: -len(".weight")] for key in expert_weight_keys] synthesized_paths: set[str] = set() - for paths in groups.values(): - if "w1" not in paths or "w3" not in paths: - continue - - values: list[torch.Tensor] = [] - for proj in ("w1", "w3"): - expert_path = paths[proj] - weight_amax = _lookup_amax(amax, expert_path, "weight") - if weight_amax is None: - weight_amax = _synthesize_weight_amax( - f.get_tensor(expert_path + ".weight"), - f.get_tensor(expert_path + ".scale"), - device, - ) - synthesized_paths.add(expert_path) - values.append(weight_amax.reshape(())) - shared = torch.maximum(values[0], values[1]) - overrides[paths["w1"]] = shared - overrides[paths["w3"]] = shared - return overrides, synthesized_paths + def get_amax(expert_path: str) -> torch.Tensor: + weight_amax = _lookup_amax(amax, expert_path, "weight") + if weight_amax is None: + weight_amax = _synthesize_weight_amax( + f.get_tensor(expert_path + ".weight"), + f.get_tensor(expert_path + ".scale"), + device, + ) + synthesized_paths.add(expert_path) + return weight_amax + + return build_w13_amax_overrides(expert_bases, get_amax), synthesized_paths def convert_shard( @@ -407,10 +261,15 @@ def convert_shard( all_keys = list(f.keys()) expert_weight_keys = [k for k in all_keys if _EXPERT_WEIGHT_RE.match(k)] expert_weight_key_set = set(expert_weight_keys) + expert_bases = [key[: -len(".weight")] for key in expert_weight_keys] if cast: # Closed-form weight cast derives scales from the source E8M0 # exponents, not from calibrated weight amax. w1/w3 share k_max. - w13_kmax = _build_w13_kmax_overrides(f, expert_weight_keys, device) + w13_kmax = build_w13_kmax_overrides( + expert_bases, + lambda base: f.get_tensor(base + ".scale"), + device, + ) w13_weight_amax, w13_synth_paths = {}, set() else: w13_kmax = {} @@ -463,9 +322,9 @@ def convert_shard( if cast: k_max = w13_kmax.get(expert_path) if k_max is None: - k_max = _kmax_from_mxfp4_scale(s, device) + k_max = mxfp4_kmax(s, device) packed, weight_scale, weight_scale_2, n_blk, n_lossless = ( - _quantize_weight_nvfp4_lossless(w, s, k_max, device) + quantize_mxfp4_to_nvfp4_lossless(w, s, k_max, device) ) weight_synth = False stats["cast_blocks_total"] += n_blk @@ -473,7 +332,7 @@ def convert_shard( if n_lossless < n_blk: stats[f"cast_oor_tensors_{block_kind}"] += 1 else: - packed, weight_scale, weight_scale_2, weight_synth = _quantize_weight_nvfp4( + packed, weight_scale, weight_scale_2, weight_synth = quantize_mxfp4_to_nvfp4( w, s, weight_amax, device=device ) input_scale = _amax_to_nvfp4_scale_2(input_amax).to(weight_scale_2.device) @@ -514,57 +373,6 @@ def convert_shard( _SKIP_SUBDIR_NAMES = {"__pycache__"} -def _link_or_copy(src: Path, dst: Path) -> None: - try: - os.link(src, dst) - except OSError as e: - copy_errnos = { - errno.EXDEV, - errno.EPERM, - errno.EACCES, - getattr(errno, "EOPNOTSUPP", errno.EXDEV), - getattr(errno, "ENOTSUP", errno.EXDEV), - } - if e.errno not in copy_errnos: - raise - shutil.copy2(src, dst) - - -def _hard_link_aux(src: Path, dst: Path) -> None: - """Link everything that isn't a shard file, rewritten metadata, or - a cache/__pycache__ directory. Recurses into legit subdirectories - (``encoding/``, ``inference/`` etc.) preserving structure. - - Falls back to copying when source and destination are on different - filesystems, which is common with container mounts. - """ - for item in src.iterdir(): - if item.name in _SKIP_TOP_LEVEL: - continue - if _HF_SHARD_RE.match(item.name): - continue - target = dst / item.name - if item.is_file(): - if target.exists(): - target.unlink() - _link_or_copy(item, target) - elif item.is_dir(): - target.mkdir(exist_ok=True) - for root, dirs, files in os.walk(item): - dirs[:] = [d for d in dirs if d not in _SKIP_SUBDIR_NAMES] - rel = Path(root).relative_to(item) - (target / rel).mkdir(parents=True, exist_ok=True) - for fname in files: - # Never pull stale shards / indexes from inside subdirs. - if fname == "model.safetensors.index.json" or _HF_SHARD_RE.match(fname): - continue - src_f = Path(root) / fname - dst_f = target / rel / fname - if dst_f.exists(): - dst_f.unlink() - _link_or_copy(src_f, dst_f) - - def _build_moe_quantization(quantized_layer_names: list[str]) -> dict[str, Any]: return { "quant_algo": "MIXED_PRECISION", @@ -680,37 +488,6 @@ def _routed_experts_prefix(expert_proj: str) -> str: return match.group("experts") -def _validate_paths(source_ckpt: Path, output_ckpt: Path) -> None: - source_resolved = source_ckpt.resolve() - output_resolved = output_ckpt.resolve() - if ( - output_resolved == source_resolved - or source_resolved in output_resolved.parents - or output_resolved in source_resolved.parents - ): - raise ValueError( - "--source_ckpt and --output_ckpt must be disjoint directories; " - f"got source={source_ckpt}, output={output_ckpt}" - ) - - -def _prepare_output_dir(output_ckpt: Path, overwrite: bool) -> None: - if output_ckpt.exists(): - if not output_ckpt.is_dir(): - raise ValueError(f"--output_ckpt exists and is not a directory: {output_ckpt}") - if any(output_ckpt.iterdir()): - if not overwrite: - raise ValueError( - f"--output_ckpt is not empty: {output_ckpt}; pass --overwrite to replace it" - ) - for item in output_ckpt.iterdir(): - if item.is_dir() and not item.is_symlink(): - shutil.rmtree(item) - else: - item.unlink() - output_ckpt.mkdir(parents=True, exist_ok=True) - - def main(): p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter @@ -758,7 +535,7 @@ def main(): ) args = p.parse_args() - _validate_paths(args.source_ckpt, args.output_ckpt) + validate_paths(args.source_ckpt, args.output_ckpt) src_index_path = args.source_ckpt / "model.safetensors.index.json" assert src_index_path.exists(), ( @@ -771,7 +548,7 @@ def main(): shards = sorted(args.source_ckpt.glob("model-*-of-*.safetensors")) assert shards, f"no HF-style shards in {args.source_ckpt}" - _prepare_output_dir(args.output_ckpt, args.overwrite) + prepare_output_dir(args.output_ckpt, args.overwrite) _log(f"[config] {len(shards)} input shards device={args.device}") stats: dict[str, int] = defaultdict(int) @@ -822,7 +599,15 @@ def main(): sorted(quantized), ) _log(f"[aux] linking ancillary files from {args.source_ckpt}") - _hard_link_aux(args.source_ckpt, args.output_ckpt) + link_aux_files( + args.source_ckpt, + args.output_ckpt, + skip_top_level=_SKIP_TOP_LEVEL, + skip_dir_names=_SKIP_SUBDIR_NAMES, + skip_file=lambda path: ( + path.name == "model.safetensors.index.json" or _HF_SHARD_RE.match(path.name) is not None + ), + ) _log(f"[done] {args.output_ckpt} ({len(quantized)} quantized routed-expert modules)") diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index fea69221825..aaac7a5f04d 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -111,6 +111,7 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http | QWen3, 3.5 MOE, Next 6 | ✅ | - | - | - | ✅ | | QwQ | ✅ | - | - | - | ✅ | | DeepSeek V3, R1, V3.1, V3.27 | - | - | - | - | ✅ | +| Kimi K314 | - | - | - | - | ✅ | | GLM-4.78 | ✅ | - | - | - | ✅ | | Kimi K2 | - | - | - | - | ✅ | | MiniMax M2.1 | - | - | - | - | ✅ | @@ -137,7 +138,8 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http > *10.GPT-OSS ships with native MXFP4 weights; NVFP4 export is produced via the closed-form `--cast_mxfp4_to_nvfp4` cast (see [MXFP4 → NVFP4 cast](#mxfp4--nvfp4-cast-for-gpt-oss)).* \ > *11.Vision-language model (VLM): only the language model is quantized while the vision encoder is kept in high precision. Pass `--vlm` to the shell script (see [VLM quantization](#vlm-quantization)).* \ > *12.For VLMs, `int8_smoothquant` only supports TensorRT-LLM checkpoint export and is not compatible with the TensorRT-LLM torch backend.* \ -> *13.Nemotron VL automatically calibrates with image-text pairs; see [VLM calibration with image-text pairs](#vlm-calibration-with-image-text-pairs-eg-nemotron-vl).* +> *13.Nemotron VL automatically calibrates with image-text pairs; see [VLM calibration with image-text pairs](#vlm-calibration-with-image-text-pairs-eg-nemotron-vl).* \ +> *14.Kimi K3 uses the calibration-free [streaming converter](../kimi/README.md) because its routed experts are released as packed MXFP4 tensors; it does not use the in-memory `hf_ptq.py` flow.* > *The accuracy loss after PTQ may vary depending on the actual model and the quantization method. Different models may have different accuracy loss and usually the accuracy loss is more significant when the base model is small. If the accuracy after PTQ is not meeting the requirement, please try either modifying [hf_ptq.py](./hf_ptq.py) and disabling the KV cache quantization or using the [QAT](./../llm_qat/README.md) instead. For NVFP4 quantization specifically, we recommend `nvfp4_mlp_only`, `nvfp4_experts_only`, or `nvfp4_omlp_only` to achieve higher accuracy by restricting quantization to the MLP/expert layers (and optionally the `o_proj` layer) while keeping the attention QKV projections unquantized.* @@ -256,6 +258,12 @@ The cast pins each NVFP4 block's `scale_2 = 2^(k_max - 8)` and `_amax = 6 * 2^k_ [PTQ for DeepSeek](../deepseek/README.md) shows how to quantize the DeepSeek model with FP4 and export to TensorRT-LLM. +#### Kimi K3 + +[PTQ for Kimi K3](../kimi/README.md) shows how to cast the source MXFP4 routed +experts to NVFP4 and quantize KDA/MLA attention weights to 128x128 block FP8 +without loading the full model or running calibration. + #### VLM quantization Vision-language models are quantized through the same script. Add `--vlm` so the script runs the diff --git a/examples/kimi/README.md b/examples/kimi/README.md new file mode 100644 index 00000000000..70364db2c9a --- /dev/null +++ b/examples/kimi/README.md @@ -0,0 +1,36 @@ +# Kimi post-training quantization + +## Kimi-K3: NVFP4 routed experts and block-FP8 attention + +`moonshotai/Kimi-K3` is released with its routed experts already packed as +MXFP4. Loading the 2.8T-parameter model and running the normal in-memory +`examples/hf_ptq/hf_ptq.py` flow would both discard that source representation +and require impractical host memory. The Kimi-K3 converter therefore streams +one safetensors shard at a time: + +- routed-expert `w1`, `w2`, and `w3` weights are cast from MXFP4 to NVFP4; +- expert activation `input_scale` is fixed to `1.0`; +- KDA and MLA projection weights are quantized to FP8 in 128x128 blocks; +- attention activations are dynamically quantized by the inference runtime; +- shared experts, latent experts, routers, convolution weights, norms, the + vision tower, and `lm_head` remain BF16. + +The exact quantization map is recorded in +`modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention.yaml`. +Because the source checkpoint uses packed MXFP4 tensors rather than ordinary +Hugging Face `Linear.weight` tensors, run the streaming converter instead of +passing this recipe to `hf_ptq.py`: + +```bash +python examples/kimi/kimi_k3/quantize_to_nvfp4.py \ + --source_ckpt /models/moonshotai/Kimi-K3 \ + --output_ckpt /models/Kimi-K3-NVFP4 \ + --recipe huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention \ + --jobs 8 +``` + +The conversion is calibration-free: it does not run a forward pass, require a +dataset, or require a GPU. For multi-node conversion, launch one process per +node against shared storage and set `--rank`, `--world_size`, and a common +`--run_id`. Use `--help` for the full set of conversion and synchronization +options. diff --git a/examples/kimi/kimi_k3/quantize_to_nvfp4.py b/examples/kimi/kimi_k3/quantize_to_nvfp4.py new file mode 100644 index 00000000000..aab2172ea24 --- /dev/null +++ b/examples/kimi/kimi_k3/quantize_to_nvfp4.py @@ -0,0 +1,1046 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Kimi-K3: MXFP4 routed experts -> NVFP4, plus FP8-family attention. + +``moonshotai/Kimi-K3`` ships 2.8T parameters in a hybrid checkpoint: + +* routed experts (92 MoE layers x 896 experts x {w1,w2,w3}) are **MXFP4**, + stored compressed-tensors style as ``.weight_packed`` (E2M1 nibbles, + 2 per byte) + ``.weight_scale`` (uint8 E8M0, one per 32-element block); +* everything else -- MLA/KDA attention, LatentMoE, shared experts, the dense + layer-0 MLP, ``lm_head``, the MoonViT tower -- is **BF16**. + +This script rewrites the 96-shard release in place-for-place fashion: + +1. **Routed experts: closed-form MXFP4 -> NVFP4 cast** (``--cast_mxfp4_to_nvfp4``). + NVFP4 uses the same E2M1 nibble grid with 16-element blocks and a two-level + scale ``per_block_scale (E4M3) * scale_2 (fp32)``. Pinning + ``scale_2 = 2^m`` (``m = k_max - 8``) and ``per_block_scale = 2^(k_j - m)`` + makes ``per_block_scale * scale_2 == 2^k_j`` exactly, so every NVFP4 nibble + equals the source MXFP4 nibble verbatim -- bit-exact for every block whose + ``k_j`` lands in E4M3's window (``k_max - k_j <= 17``). Blocks outside the + window fall back to a data-derived per-block amax. The numerics are shared + with the GPT-OSS cast (``examples/hf_ptq/cast_mxfp4_to_nvfp4.py``, PR #1372) + via ``modelopt.torch.quantization.utils.numeric_utils``. + + As in DeepSeek-V4, ``w1``/``w3`` feed one fused GEMM1 and therefore must + share a single ``scale_2``, so ``k_max`` is taken over both projections. + +2. **Attention: BF16 -> block FP8** (``--attn_fp8_pb``), MXFP8 + (``--attn_mxfp8``), or static per-tensor FP8 (``--attn_fp8``). Block FP8 + uses one fp32 scale per 128x128 weight tile and dynamic per-token activation + scaling. MXFP8 uses one E8M0 scale per 32 weights + and dynamic activation scaling, so it remains calibration-free without a + coarse static activation scale. Static FP8 weight scales are derived as + ``amax(|W|) / 448`` and activation ``input_scale`` is pinned to + ``--input_scale``. Both cover the MLA projections (``q_a_proj``, ``q_b_proj``, + ``kv_a_proj_with_mqa``, ``kv_b_proj``), the KDA projections (``q_proj``, + ``k_proj``, ``v_proj``) and the per-layer ``o_proj`` / ``g_proj`` -- ~36B + parameters, 72GB of BF16 down to 36GB. The vLLM-fused KDA projections + include ``b_proj`` and ``f_a_proj`` and are quantized consistently with the + other members of ``in_proj_qkvgfab``. Short convolutions (``*_conv1d``), + ``A_log``, ``dt_bias``, every norm, and the MoE router ``gate`` stay BF16. + ``f_b_proj`` stays BF16 for per-tensor FP8/MXFP8, but is included in the + block-FP8 recipe because its TP-local shape is already 128-tile aligned. + + Not covered here, but the obvious next candidates if more memory is needed: + the LatentMoE ``routed_expert_{up,down}_proj`` (~9.5GB) and + ``shared_experts.*`` (~24GB), both of which are replicated rather than + sharded under expert parallelism. + +**No calibration anywhere.** Expert ``input_scale`` is pinned to +``--input_scale`` (default 1.0) rather than derived from an amax dump, so the +whole conversion is a closed-form tensor transform: no forward pass, no +dataset, no GPU required. ``--device cpu`` is the default and is what the +shipped configuration uses; the work is dominated by shard I/O. + +Usage (CPU partition, no GPU needed; ``--jobs`` shards convert in parallel): + + python quantize_to_nvfp4.py \\ + --source_ckpt /path/to/Kimi-K3 \\ + --output_ckpt /path/to/Kimi-K3-NVFP4 \\ + --recipe huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention \\ + --jobs 8 +""" + +from __future__ import annotations + +import argparse +import json +import math +import multiprocessing +import os +import re +import shutil +import sys +import time +from collections import defaultdict +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from examples._shard_cast_utils import ( # noqa: E402 + build_w13_amax_overrides, + build_w13_kmax_overrides, + dequantize_mxfp4_to_bf16, + link_aux_files, + mxfp4_kmax, + prepare_output_dir, + quantize_mxfp4_to_nvfp4, + quantize_mxfp4_to_nvfp4_lossless, + validate_paths, +) +from examples._shard_cast_utils import log as _log # noqa: E402 +from modelopt import __version__ as modelopt_version # noqa: E402 +from modelopt.recipe import load_recipe # noqa: E402 +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format # noqa: E402 +from modelopt.torch.quantization.qtensor import FP8QTensor, MXFP8QTensor # noqa: E402 +from modelopt.torch.quantization.utils.numeric_utils import E2M1_MAX, E4M3_MAX # noqa: E402 + +# -------------------------------------------------------------------------- +# Kimi-K3 tensor schema (from model.safetensors.index.json, 497220 tensors). +# -------------------------------------------------------------------------- +_LM = r"language_model\.model\.layers\.\d+" + +# Routed experts: MXFP4 source pair .weight_packed + .weight_scale. +_EXPERT_PACKED_RE = re.compile( + rf"^(?P{_LM}\.block_sparse_moe\.experts\.\d+\.w[123])\.weight_packed$" +) +# Prefix reported in the quantized-layer manifest (one entry per expert bank). +_EXPERT_BANK_RE = re.compile(rf"^(?P{_LM}\.block_sparse_moe\.experts)\.\d+\.w[123]$") + +# Attention projections eligible for FP8/MXFP8. Everything else under +# ``self_attn`` (conv1d kernels, A_log, dt_bias, o_norm, layernorms) stays BF16. +# vLLM fuses q/k/v/g/f_a/b into ``in_proj_qkvgfab``, so those six must use one +# quantization algorithm. ``f_b_proj`` remains BF16 except in the block-FP8 +# recipe. +_ATTN_QUANT_PROJ = frozenset( + { + # KDA (69 layers) + "q_proj", + "k_proj", + "v_proj", + # MLA (24 layers) + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + # present on all 93 layers + "o_proj", + "g_proj", + # vLLM fuses both into in_proj_qkvgfab with q/k/v/g. + "b_proj", + "f_a_proj", + } +) +# ``f_b_proj`` is a small KDA state-control projection. It remains BF16 for +# per-tensor FP8/MXFP8, but its TP-local [1536, 128] shape is naturally aligned +# for 128x128 block FP8 and the prior standalone-FP8 recipe found it safe. +_ATTN_FP8_PB_PROJ = _ATTN_QUANT_PROJ | {"f_b_proj"} +_ATTN_WEIGHT_RE = re.compile(rf"^(?P{_LM}\.self_attn\.(?P[a-z_0-9]+))\.weight$") + +_FP8_MAX = 448.0 +_FP8_PB_BLOCK = 128 +_NVFP4_BLOCK = 16 # NVFP4 block size (elements) + +_PUBLISHED_RECIPE = "huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention" + + +def _conversion_settings_from_recipe(recipe_path: str) -> dict[str, Any]: + """Translate the supported Kimi-K3 recipe into streaming-converter settings.""" + quant_cfg = load_recipe(recipe_path).quantize.model_dump()["quant_cfg"] + by_name = { + entry["quantizer_name"]: entry + for entry in quant_cfg + if isinstance(entry, dict) and "quantizer_name" in entry + } + + expert_weight_name = "*block_sparse_moe.experts.*weight_quantizer" + expert_input_name = "*block_sparse_moe.experts.*input_quantizer" + try: + expert_weight = by_name[expert_weight_name]["cfg"] + expert_input = by_name[expert_input_name]["cfg"] + except KeyError as exc: + raise ValueError("recipe must configure Kimi-K3 routed experts") from exc + + if ( + expert_weight.get("num_bits") != (2, 1) + or expert_weight.get("block_sizes", {}).get(-1) != _NVFP4_BLOCK + ): + raise ValueError("recipe routed-expert weights must use block-16 NVFP4") + + constant_amax = expert_input.get("constant_amax") + if constant_amax is None: + raise ValueError("recipe routed-expert inputs must set constant_amax") + input_scale = float(constant_amax) / (E2M1_MAX * E4M3_MAX) + + for projection in _ATTN_FP8_PB_PROJ: + name = f"*self_attn.{projection}*weight_quantizer" + try: + cfg = by_name[name]["cfg"] + except KeyError as exc: + raise ValueError(f"recipe must configure {projection} attention weights") from exc + if cfg.get("num_bits") != (4, 3) or cfg.get("block_sizes") != {-1: 128, -2: 128}: + raise ValueError(f"recipe {projection} weights must use 128x128 block FP8") + + return { + "cast_mxfp4_to_nvfp4": True, + "attn_fp8": False, + "attn_mxfp8": False, + "attn_fp8_pb": True, + "input_scale": input_scale, + } + + +# -------------------------------------------------------------------------- +# BF16 -> per-tensor FP8 (attention) +# -------------------------------------------------------------------------- +# Projections that serving stacks fuse into a single Linear, and therefore +# must share one per-tensor scale. SGLang builds ``fused_qkvg_proj`` over the +# wide KDA projections and ``fused_qkv_a_proj_with_mqa`` over the MLA pair. +# +# Sharing is not merely cosmetic. The KDA fusion is a MergedColumnParallelLinear, +# which keeps a per-shard scale array and requantizes to the max across shards, +# so independent scales would survive there. The MLA fusion is a *ReplicatedLinear* +# with a single scalar scale and no shard structure to requantize over, so two +# independent scales silently misinterpret half the fused weight. Deriving one +# amax per group up front makes the checkpoint correct under either layout. +_FUSED_ATTN_GROUPS: tuple[tuple[str, ...], ...] = ( + ("q_proj", "k_proj", "v_proj", "g_proj", "f_a_proj", "b_proj"), + ("q_a_proj", "kv_a_proj_with_mqa"), # MLA +) +_PROJ_TO_GROUP = {proj: i for i, g in enumerate(_FUSED_ATTN_GROUPS) for proj in g} + + +def _fused_attn_group_key(layer_prefix: str, proj: str) -> str | None: + """``.self_attn`` + group index, or None when the proj is standalone.""" + idx = _PROJ_TO_GROUP.get(proj) + return None if idx is None else f"{layer_prefix}#{idx}" + + +def compute_fused_attn_amax(shards: list[Path], device: str) -> dict[str, torch.Tensor]: + """One shared ``amax`` per fused attention group, over all shards. + + Members of a group are not guaranteed to land in the same shard, so this is + a separate pass over the attention weights (~36B parameters) before any + conversion happens. It reads only BF16 attention tensors, not the experts. + """ + group_amax: dict[str, torch.Tensor] = {} + for shard in shards: + with safe_open(str(shard), framework="pt", device="cpu") as f: + for key in f.keys(): # noqa: SIM118 + m = _ATTN_WEIGHT_RE.match(key) + if not m or m.group("proj") not in _ATTN_QUANT_PROJ: + continue + gk = _fused_attn_group_key(m.group("base").rsplit(".", 1)[0], m.group("proj")) + if gk is None: + continue + amax = f.get_tensor(key).to(device).float().abs().max().cpu() + prev = group_amax.get(gk) + group_amax[gk] = amax if prev is None else torch.maximum(prev, amax) + _log(f"[fused-attn] shared amax computed for {len(group_amax)} fused groups") + return group_amax + + +def _quantize_weight_fp8( + weight: torch.Tensor, device: str, amax: torch.Tensor | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + """Data-free per-tensor FP8: ``weight_scale = amax(|W|) / 448``. + + ``amax`` overrides the tensor's own maximum so every member of a fused + group is quantized against one shared scale. Returns + ``(fp8_weight, weight_scale)``. A degenerate all-zero weight would give a + zero scale, which would make dequantization ill-defined, so the scale is + floored at the smallest positive normal fp32 value. + """ + w = weight.to(device).to(torch.float32) + a = w.abs().max() if amax is None else amax.to(w.device).float() + weight_scale = (a / _FP8_MAX).clamp(min=torch.finfo(torch.float32).tiny).reshape(()) + q = (w / weight_scale).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + return q, weight_scale.to(torch.float32) + + +def _quantize_weight_mxfp8(weight: torch.Tensor, device: str) -> tuple[torch.Tensor, torch.Tensor]: + """Calibration-free MXFP8 with one E8M0 scale per 32 weights.""" + q_tensor, weight_scale = MXFP8QTensor.quantize(weight.to(device)) + return q_tensor._quantized_data, weight_scale + + +def _quantize_weight_fp8_pb(weight: torch.Tensor, device: str) -> tuple[torch.Tensor, torch.Tensor]: + """Calibration-free 128x128 FP8 weight blocks. + + ModelOpt crops the serialized FP8 weight back to its logical shape while + retaining scales for the zero-padded edge tiles. Unified HF checkpoints + store those scales as ``[out_block, 1, in_block, 1]``. + """ + q_tensor, weight_scale = FP8QTensor.quantize( + weight.to(device), + block_sizes={-2: _FP8_PB_BLOCK, -1: _FP8_PB_BLOCK}, + ) + return ( + q_tensor._quantized_data.contiguous(), + weight_scale[:, None, :, None].to(torch.float32).contiguous(), + ) + + +# -------------------------------------------------------------------------- +# Shard conversion +# -------------------------------------------------------------------------- +def convert_shard( + src_shard: Path, + dst_shard: Path, + device: str, + cast: bool, + attn_fp8: bool, + input_scale_value: float, + fused_attn_amax: dict[str, torch.Tensor] | None = None, + attn_mxfp8: bool = False, + attn_fp8_pb: bool = False, +) -> dict[str, Any]: + """Rewrite one HF-style shard and return its conversion report.""" + out: dict[str, torch.Tensor] = {} + stats: dict[str, int] = defaultdict(int) + banks: set[str] = set() + attn_modules: set[str] = set() + + input_scale = torch.tensor(input_scale_value, dtype=torch.float32).reshape(()) + + with safe_open(str(src_shard), framework="pt", device="cpu") as f: + all_keys = list(f.keys()) + key_set = set(all_keys) + + expert_bases = [ + m.group("base") for m in (_EXPERT_PACKED_RE.match(k) for k in all_keys) if m + ] + # Source E8M0 scale tensors are consumed and replaced by NVFP4 scales. + expert_scale_keys = {b + ".weight_scale" for b in expert_bases} + expert_packed_keys = {b + ".weight_packed" for b in expert_bases} + + if cast: + w13_kmax = build_w13_kmax_overrides( + expert_bases, + lambda base: f.get_tensor(base + ".weight_scale"), + device, + ) + w13_weight_amax = {} + else: + w13_kmax = {} + w13_weight_amax = build_w13_amax_overrides( + expert_bases, + lambda base: ( + dequantize_mxfp4_to_bf16( + f.get_tensor(base + ".weight_packed"), + f.get_tensor(base + ".weight_scale"), + device, + ) + .abs() + .max() + ), + ) + + for key in all_keys: + # Source MXFP4 E8M0 scales are rewritten below alongside the packed + # weight; skip them here so they are not emitted twice. + if key in expert_scale_keys: + continue + + if key in expert_packed_keys: + base = key[: -len(".weight_packed")] + scale_key = base + ".weight_scale" + assert scale_key in key_set, f"no paired weight_scale for {key}" + + w = f.get_tensor(key) + s = f.get_tensor(scale_key) + + if cast: + k_max = w13_kmax.get(base) + if k_max is None: + k_max = mxfp4_kmax(s, device) + packed, weight_scale, weight_scale_2, n_blk, n_lossless = ( + quantize_mxfp4_to_nvfp4_lossless(w, s, k_max, device) + ) + stats["cast_blocks_total"] += n_blk + stats["cast_blocks_lossless"] += n_lossless + if n_lossless < n_blk: + stats["cast_oor_tensors"] += 1 + else: + packed, weight_scale, weight_scale_2, _ = quantize_mxfp4_to_nvfp4( + w, + s, + w13_weight_amax.get(base), + device, + ) + + out[base + ".weight"] = packed.cpu() + out[base + ".weight_scale"] = weight_scale.cpu() + out[base + ".weight_scale_2"] = weight_scale_2.cpu() + out[base + ".input_scale"] = input_scale.clone() + + stats["experts_converted"] += 1 + bank = _EXPERT_BANK_RE.match(base) + assert bank is not None, f"unexpected expert path: {base}" + banks.add(bank.group("bank")) + continue + + attn = _ATTN_WEIGHT_RE.match(key) + if attn_fp8_pb and attn and attn.group("proj") in _ATTN_FP8_PB_PROJ: + base = attn.group("base") + q, weight_scale = _quantize_weight_fp8_pb(f.get_tensor(key), device) + out[key] = q.cpu() + out[base + ".weight_scale"] = weight_scale.cpu() + stats["attn_fp8_pb_converted"] += 1 + attn_modules.add(base) + continue + + if attn_fp8 and attn and attn.group("proj") in _ATTN_QUANT_PROJ: + base = attn.group("base") + gk = _fused_attn_group_key(base.rsplit(".", 1)[0], attn.group("proj")) + shared = (fused_attn_amax or {}).get(gk) if gk else None + q, weight_scale = _quantize_weight_fp8(f.get_tensor(key), device, shared) + if shared is not None: + stats["attn_fp8_fused_shared_scale"] += 1 + out[key] = q.cpu() + out[base + ".weight_scale"] = weight_scale.cpu() + out[base + ".input_scale"] = input_scale.clone() + stats["attn_fp8_converted"] += 1 + attn_modules.add(base) + continue + + if attn_mxfp8 and attn and attn.group("proj") in _ATTN_QUANT_PROJ: + base = attn.group("base") + q, weight_scale = _quantize_weight_mxfp8(f.get_tensor(key), device) + out[key] = q.cpu() + out[base + ".weight_scale"] = weight_scale.cpu() + stats["attn_mxfp8_converted"] += 1 + attn_modules.add(base) + continue + + out[key] = f.get_tensor(key) + stats["passthrough"] += 1 + + save_file(out, str(dst_shard)) + return { + "shard": src_shard.name, + "tensor_bytes": sum(t.numel() * t.element_size() for t in out.values()), + "stats": dict(stats), + "banks": sorted(banks), + "attn_modules": sorted(attn_modules), + } + + +def _convert_shard_task(job: dict[str, Any]) -> dict[str, Any]: + """ProcessPoolExecutor entry point (module-level so it is picklable).""" + torch.set_num_threads(job["threads"]) + result = convert_shard( + Path(job["src"]), + Path(job["dst"]), + job["device"], + job["cast"], + job["attn_fp8"], + job["input_scale"], + job.get("fused_attn_amax"), + job.get("attn_mxfp8", False), + job.get("attn_fp8_pb", False), + ) + _log(f"[shard] {result['shard']} done: {result['stats']}") + return result + + +# -------------------------------------------------------------------------- +# Ancillary files / config / index +# -------------------------------------------------------------------------- +_SKIP_TOP_LEVEL = { + "model.safetensors.index.json", # rewritten + "config.json", # rewritten (mark hybrid NVFP4 MoE + FP8 attention) + "hf_quant_config.json", # rewritten + "conversion_report.json", # rewritten + ".kimi_k3_conversion", # temporary distributed-conversion rendezvous + ".cache", # HF download sidecars referencing old shards +} +_SKIP_SUBDIR_NAMES = {"__pycache__"} + + +# Modules deliberately left in BF16. Kept explicit so a reader can see that the +# router gate, the KDA state parameters, the norms, the LatentMoE/shared-expert +# projections, the dense layer-0 MLP and the whole vision stack are untouched. +_EXCLUDE_MODULES = [ + "*block_sparse_moe.gate*", + "*block_sparse_moe.routed_expert_up_proj*", + "*block_sparse_moe.routed_expert_down_proj*", + "*block_sparse_moe.routed_expert_norm*", + "*block_sparse_moe.shared_experts*", + "*self_attn.q_conv1d*", + "*self_attn.k_conv1d*", + "*self_attn.v_conv1d*", + # f_b_proj is a small KDA state-control linear and stays BF16 except in the + # block-FP8 recipe. b_proj and f_a_proj are quantized because vLLM folds + # them into in_proj_qkvgfab. + "*self_attn.f_b_proj*", + "*self_attn.o_norm*", + "*self_attn.q_a_layernorm*", + "*self_attn.kv_a_layernorm*", + "*_res_proj*", + "*_res_norm*", + "*layernorm*", + "language_model.model.layers.0.mlp.*", + "language_model.lm_head", + "vision_tower*", + "mm_projector*", +] + + +def _module_name_aliases(name: str) -> list[str]: + """Every prefix spelling a loader may probe for one HF module path. + + Manifest keys are HF checkpoint paths, but loaders look modules up by their + runtime prefix, and for Kimi-K3 the two differ twice over. SGLang's + ``kimi_k3.py`` builds the language model with ``prefix=""`` even though the + attribute is ``language_model``, so lookups arrive as ``model.layers.N...``; + and its ``WeightsMapper`` rewrites ``block_sparse_moe`` -> ``mlp`` but does + not strip ``language_model.``. vLLM's K3 mapper strips + ``language_model.model.`` entirely and probes ``layers.N...``. A manifest + keyed only on HF names therefore matches nothing, every module falls back + to the unquantized method, and a 2.8T MoE tries to allocate BF16 experts. + Emitting all spellings costs a few KB and keeps the checkpoint loadable + under either convention. + """ + names = {name} + if "block_sparse_moe" in name: + names.add(name.replace("block_sparse_moe", "mlp")) + runtime_names: set[str] = set() + for n in names: + if n.startswith("language_model.model."): + suffix = n[len("language_model.model.") :] + runtime_names.add("model." + suffix) + runtime_names.add(suffix) + elif n.startswith("language_model."): + runtime_names.add(n[len("language_model.") :]) + names.update(runtime_names) + return sorted(names) + + +def _fused_attention_module_names(attn_modules: list[str]) -> set[str]: + """Return K3 runtime fused-linear prefixes implied by HF projection names. + + Different engines fuse the attention projections under different names, so + emit every spelling and let the loader match whichever it builds. + + SGLang (use_full_rank_gate): the wide [q, k, v, g] projections fuse into + ``fused_qkvg_proj`` while b / f_a / f_b stay standalone -- a clean split for + our FP8 policy, which quantizes exactly q/k/v/g and leaves the KDA + state-control projections BF16. The MLA pair fuses into + ``fused_qkv_a_proj_with_mqa`` (matches deepseek_v2). + + vLLM names the KDA fusion ``in_proj_qkvgfab`` and the MLA fusion + ``fused_qkv_a_proj``; the former also folds in b/f_a, so a mixed-precision + fused module there is vLLM's own concern. + """ + fused: set[str] = set() + for name in attn_modules: + parent, proj = name.rsplit(".", 1) + if proj in {"q_proj", "k_proj", "v_proj", "g_proj"}: + fused.add(parent + ".fused_qkvg_proj") # SGLang KDA + if proj in {"q_proj", "k_proj", "v_proj", "g_proj", "f_a_proj", "b_proj"}: + fused.add(parent + ".in_proj_qkvgfab") # vLLM KDA + if proj in {"q_a_proj", "kv_a_proj_with_mqa"}: + fused.add(parent + ".fused_qkv_a_proj_with_mqa") # SGLang MLA + fused.add(parent + ".fused_qkv_a_proj") # vLLM MLA + return fused + + +def _build_hf_quant_config( + expert_banks: list[str], + attn_modules: list[str], + attn_fp8: bool, + attn_mxfp8: bool = False, + attn_fp8_pb: bool = False, +) -> dict[str, Any]: + quantized_layers: dict[str, dict[str, Any]] = {} + for name in expert_banks: + for alias in _module_name_aliases(name): + quantized_layers[alias] = {"quant_algo": "NVFP4", "group_size": _NVFP4_BLOCK} + attention_algo = ( + "FP8" if attn_fp8 else "MXFP8" if attn_mxfp8 else "FP8_PB_WO" if attn_fp8_pb else None + ) + if attention_algo: + for name in attn_modules: + for alias in _module_name_aliases(name): + quantized_layers[alias] = {"quant_algo": attention_algo} + # vLLM's K3 model names these packed modules directly. Its mixed + # resolver does not currently infer either K3-specific fused name from + # the individual HF projection policies. + for name in _fused_attention_module_names(attn_modules): + for alias in _module_name_aliases(name): + quantized_layers[alias] = {"quant_algo": attention_algo} + + exclude_modules: list[str] = [] + for pattern in _EXCLUDE_MODULES: + if attn_fp8_pb and pattern == "*self_attn.f_b_proj*": + continue + exclude_modules.extend(_module_name_aliases(pattern)) + + return { + "producer": {"name": "modelopt", "version": modelopt_version}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "group_size": _NVFP4_BLOCK, + "quantized_layers": quantized_layers, + "exclude_modules": list(dict.fromkeys(exclude_modules)), + }, + } + + +def _rewrite_config_json(src_dir: Path, dst_dir: Path, hf_quant_config: dict[str, Any]) -> None: + """Copy ``config.json`` and replace the source MXFP4 manifest. + + The source ``quantization_config`` is compressed-tensors ``mxfp4-pack-quantized`` + and describes weights this script has just replaced, so leaving it in place + would make a loader dequantize the NVFP4 experts as MXFP4. It is replaced + wholesale by the ModelOpt mixed-precision manifest. + """ + cfg = json.loads((src_dir / "config.json").read_text()) + quant_cfg = convert_hf_quant_config_format(hf_quant_config) + # ``convert_hf_quant_config_format`` targets the llm-compressor layout and + # stamps ``quant_method="modelopt"``. Loaders gate their mixed-precision + # path on ``"modelopt_mixed"`` specifically (SGLang: + # ``ModelOptMixedPrecisionConfig.override_quantization_method``), so with + # the generic value the manifest is silently ignored and every quantized + # module falls back to the unquantized method -- which for a 2.8T MoE means + # allocating BF16 experts and OOMing at load. + if quant_cfg.get("quant_algo") == "MIXED_PRECISION": + quant_cfg["quant_method"] = "modelopt_mixed" + # ``text_config`` carries the source MXFP4 manifest for this model. + text_cfg = cfg.get("text_config") + if isinstance(text_cfg, dict): + text_cfg.pop("quantization_config", None) + cfg["quantization_config"] = quant_cfg + (dst_dir / "config.json").write_text(json.dumps(cfg, indent=2, sort_keys=True) + "\n") + + +def _write_index_and_manifest( + output_ckpt: Path, + src_index: dict, + results: list[dict[str, Any]], + hf_quant_config: dict[str, Any], + attn_fp8: bool, + attn_mxfp8: bool = False, + attn_fp8_pb: bool = False, +) -> None: + converted_shards = {r["shard"] for r in results} + weight_map: dict[str, str] = {} + for key, shard in src_index["weight_map"].items(): + if shard not in converted_shards: + continue + + expert = _EXPERT_PACKED_RE.match(key) + if expert: + base = expert.group("base") + weight_map[base + ".weight"] = shard + weight_map[base + ".weight_scale_2"] = shard + weight_map[base + ".input_scale"] = shard + continue + + weight_map[key] = shard + attn = _ATTN_WEIGHT_RE.match(key) + attention_projections = _ATTN_FP8_PB_PROJ if attn_fp8_pb else _ATTN_QUANT_PROJ + if ( + (attn_fp8 or attn_mxfp8 or attn_fp8_pb) + and attn + and attn.group("proj") in attention_projections + ): + base = attn.group("base") + weight_map[base + ".weight_scale"] = shard + if attn_fp8: + weight_map[base + ".input_scale"] = shard + + # Source expert ``weight_scale`` keys remain in ``weight_map`` and now + # describe NVFP4's per-16-element E4M3 scale rather than MXFP4's per-32 + # E8M0 scale. + metadata = dict(src_index.get("metadata", {})) + metadata["total_size"] = sum(r["tensor_bytes"] for r in results) + new_index = {"metadata": metadata, "weight_map": weight_map} + (output_ckpt / "model.safetensors.index.json").write_text(json.dumps(new_index, indent=2)) + _log(f"[index] wrote model.safetensors.index.json ({len(weight_map)} keys)") + + (output_ckpt / "hf_quant_config.json").write_text(json.dumps(hf_quant_config, indent=2)) + + +def _write_json_atomic(path: Path, value: Any) -> None: + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(json.dumps(value, indent=2, sort_keys=True)) + os.replace(tmp, path) + + +def _wait_for( + predicate, + description: str, + timeout_s: float, + poll_s: float = 5.0, +) -> None: + deadline = time.monotonic() + timeout_s + while not predicate(): + if time.monotonic() >= deadline: + raise TimeoutError(f"timed out waiting for {description} after {timeout_s:.0f}s") + time.sleep(poll_s) + + +def _rank0_ready(ready_path: Path, run_id: str, world_size: int, rank: int) -> bool: + """Check that rank 0 published matching rendezvous settings.""" + if not ready_path.exists(): + return False + ready = json.loads(ready_path.read_text()) + if ready.get("run_id") != run_id: + return False + published_world_size = ready.get("world_size") + if published_world_size != world_size: + raise ValueError( + f"rank {rank} has --world_size {world_size}, but rank 0 published " + f"{published_world_size} for run {run_id}" + ) + return True + + +def _merge_rank_reports( + rank_reports: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], dict[str, int], set[str], set[str]]: + results: list[dict[str, Any]] = [] + totals: dict[str, int] = defaultdict(int) + banks: set[str] = set() + attn_modules: set[str] = set() + for report in rank_reports: + results.extend(report["results"]) + for k, v in report["stats"].items(): + totals[k] += v + banks.update(report["banks"]) + attn_modules.update(report["attn_modules"]) + return results, dict(totals), banks, attn_modules + + +def main(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--source_ckpt", type=Path, required=True, help="original HF Kimi-K3 release") + p.add_argument("--output_ckpt", type=Path, required=True) + p.add_argument( + "--recipe", + help=( + "Kimi-K3 checkpoint-mirror recipe (recommended: " + f"{_PUBLISHED_RECIPE}); cannot be combined with the low-level format flags" + ), + ) + p.add_argument( + "--cast_mxfp4_to_nvfp4", + action="store_true", + help=( + "closed-form bit-exact cast of the source MXFP4 routed-expert weights to " + "NVFP4 (pin scale_2 = 2^(k_max-8), per-block scale = 2^(k_j-m) from the " + "source E8M0 scales) instead of dequantizing and re-quantizing from data" + ), + ) + attention = p.add_mutually_exclusive_group() + attention.add_argument( + "--attn_fp8", + action="store_true", + help="quantize MLA/KDA attention projections to data-free per-tensor FP8", + ) + attention.add_argument( + "--attn_mxfp8", + action="store_true", + help="quantize MLA/KDA attention projections to calibration-free MXFP8", + ) + attention.add_argument( + "--attn_fp8_pb", + action="store_true", + help=( + "quantize MLA/KDA attention projections to calibration-free 128x128 " + "block FP8 weights with dynamic per-token activations" + ), + ) + p.add_argument( + "--input_scale", + type=float, + default=None, + help="fixed activation input_scale for every quantized module (default: 1.0)", + ) + p.add_argument("--device", default="cpu", help="'cpu' (default) or 'cuda'") + p.add_argument("--jobs", type=int, default=8, help="shards converted in parallel per rank") + p.add_argument( + "--threads_per_job", type=int, default=8, help="torch intra-op threads per worker" + ) + p.add_argument( + "--rank", + type=int, + default=0, + help="distributed shard rank (use one rank per node; default 0)", + ) + p.add_argument( + "--world_size", + type=int, + default=1, + help="number of distributed shard ranks sharing --output_ckpt (default 1)", + ) + p.add_argument( + "--run_id", + default=os.environ.get("SLURM_JOB_ID", "local"), + help="rendezvous identifier shared by distributed ranks (defaults to SLURM_JOB_ID)", + ) + p.add_argument( + "--sync_timeout", + type=float, + default=24 * 60 * 60, + help="seconds rank 0 waits for distributed ranks (default 24 hours)", + ) + p.add_argument("--limit_shards", type=int, default=0, help="smoke test: convert only N shards") + p.add_argument("--overwrite", action="store_true") + args = p.parse_args() + + if args.recipe: + if ( + args.cast_mxfp4_to_nvfp4 + or args.attn_fp8 + or args.attn_mxfp8 + or args.attn_fp8_pb + or args.input_scale is not None + ): + p.error("--recipe cannot be combined with the low-level format flags or --input_scale") + try: + settings = _conversion_settings_from_recipe(args.recipe) + except (OSError, ValueError) as exc: + p.error(f"invalid Kimi-K3 recipe: {exc}") + for name, value in settings.items(): + setattr(args, name, value) + if args.input_scale is None: + args.input_scale = 1.0 + + input_scale = float(args.input_scale) + if not math.isfinite(input_scale) or input_scale <= 0: + p.error("--input_scale must be finite and > 0") + args.input_scale = input_scale + if args.jobs <= 0: + p.error("--jobs must be > 0") + if args.threads_per_job <= 0: + p.error("--threads_per_job must be > 0") + if args.world_size <= 0: + p.error("--world_size must be > 0") + if not 0 <= args.rank < args.world_size: + p.error("--rank must satisfy 0 <= rank < world_size") + if args.limit_shards < 0: + p.error("--limit_shards must be >= 0") + if args.sync_timeout <= 0: + p.error("--sync_timeout must be > 0") + if args.device.startswith("cuda") and args.jobs > 1: + p.error("--device cuda requires --jobs 1 so workers do not contend for one GPU") + + validate_paths(args.source_ckpt, args.output_ckpt) + src_index_path = args.source_ckpt / "model.safetensors.index.json" + assert src_index_path.exists(), f"{src_index_path} not found" + src_index = json.loads(src_index_path.read_text()) + + shards = sorted(args.source_ckpt.glob("model-*-of-*.safetensors")) + assert shards, f"no HF-style shards in {args.source_ckpt}" + if args.limit_shards: + shards = shards[: args.limit_shards] + + marker_dir = args.output_ckpt / ".kimi_k3_conversion" + ready_path = marker_dir / "ready.json" + if args.rank == 0: + prepare_output_dir(args.output_ckpt, args.overwrite) + marker_dir.mkdir() + _write_json_atomic(ready_path, {"run_id": args.run_id, "world_size": args.world_size}) + else: + _wait_for( + lambda: _rank0_ready( + ready_path, + args.run_id, + args.world_size, + args.rank, + ), + f"rank-0 rendezvous for run {args.run_id}", + args.sync_timeout, + ) + + assigned_shards = shards[args.rank :: args.world_size] + _log( + f"[config] rank={args.rank}/{args.world_size} " + f"shards={len(assigned_shards)}/{len(shards)} device={args.device} jobs={args.jobs} " + f"cast={args.cast_mxfp4_to_nvfp4} attn_fp8={args.attn_fp8} " + f"attn_mxfp8={args.attn_mxfp8} attn_fp8_pb={args.attn_fp8_pb} " + f"input_scale={args.input_scale}" + ) + + # Members of a fused attention group may live in different shards, so the + # shared amax is derived in a pass over all shards before any conversion. + fused_attn_amax = compute_fused_attn_amax(shards, args.device) if args.attn_fp8 else {} + + jobs = [ + { + "src": str(s), + "dst": str(args.output_ckpt / s.name), + "device": args.device, + "cast": args.cast_mxfp4_to_nvfp4, + "attn_fp8": args.attn_fp8, + "attn_mxfp8": args.attn_mxfp8, + "attn_fp8_pb": args.attn_fp8_pb, + "input_scale": args.input_scale, + "threads": args.threads_per_job, + "fused_attn_amax": fused_attn_amax, + } + for s in assigned_shards + ] + + if args.jobs > 1 and len(jobs) > 1: + # Use a 'spawn' pool, not the default 'fork'. The --attn_fp8 path runs + # compute_fused_attn_amax() in the parent first, which initializes + # torch's intra-op thread pool; forking after that leaves the workers + # with an inherited-but-dead thread state and they hang at 0% CPU + # before doing any work. 'spawn' starts each worker from a clean + # interpreter, at the cost of re-importing torch per worker. + ctx = multiprocessing.get_context("spawn") + with ProcessPoolExecutor(max_workers=args.jobs, mp_context=ctx) as ex: + results = list(ex.map(_convert_shard_task, jobs)) + else: + results = [_convert_shard_task(j) for j in jobs] + + local_totals: dict[str, int] = defaultdict(int) + local_banks: set[str] = set() + local_attn_modules: set[str] = set() + for r in results: + for k, v in r["stats"].items(): + local_totals[k] += v + local_banks.update(r["banks"]) + local_attn_modules.update(r["attn_modules"]) + + rank_report = { + "run_id": args.run_id, + "rank": args.rank, + "results": results, + "stats": dict(local_totals), + "banks": sorted(local_banks), + "attn_modules": sorted(local_attn_modules), + } + rank_path = marker_dir / f"rank-{args.rank:05d}.json" + _write_json_atomic(rank_path, rank_report) + + if args.rank != 0: + _log(f"[rank {args.rank}] complete; rank 0 will finalize checkpoint metadata") + return + + rank_paths = [marker_dir / f"rank-{rank:05d}.json" for rank in range(args.world_size)] + + def all_ranks_done() -> bool: + for path in rank_paths: + if not path.exists(): + return False + try: + if json.loads(path.read_text()).get("run_id") != args.run_id: + return False + except (json.JSONDecodeError, OSError): + return False + return True + + _wait_for(all_ranks_done, f"{args.world_size} rank reports", args.sync_timeout) + rank_reports = [json.loads(path.read_text()) for path in rank_paths] + results, totals, banks, attn_modules = _merge_rank_reports(rank_reports) + if len(results) != len(shards): + raise RuntimeError(f"expected {len(shards)} converted shards, got {len(results)}") + + _log("[stats]") + for k in sorted(totals): + _log(f" {k:32s} {totals[k]}") + + if args.cast_mxfp4_to_nvfp4: + tot = totals.get("cast_blocks_total", 0) + loss = totals.get("cast_blocks_lossless", 0) + pct = 100.0 * loss / tot if tot else 100.0 + _log(f"[cast] lossless MXFP4->NVFP4 blocks: {loss}/{tot} ({pct:.6f}%)") + _log(f"[cast] tensors with out-of-range blocks: {totals.get('cast_oor_tensors', 0)}") + + hf_quant_config = _build_hf_quant_config( + sorted(banks), + sorted(attn_modules), + args.attn_fp8, + args.attn_mxfp8, + args.attn_fp8_pb, + ) + _write_index_and_manifest( + args.output_ckpt, + src_index, + results, + hf_quant_config, + args.attn_fp8, + args.attn_mxfp8, + args.attn_fp8_pb, + ) + attention_algo = ( + "FP8" + if args.attn_fp8 + else "MXFP8" + if args.attn_mxfp8 + else "FP8_PB_WO" + if args.attn_fp8_pb + else "BF16" + ) + _log( + f"[config] rewriting config.json " + f"(MIXED_PRECISION: NVFP4 experts + {attention_algo} attention)" + ) + _rewrite_config_json(args.source_ckpt, args.output_ckpt, hf_quant_config) + _log(f"[aux] linking ancillary files from {args.source_ckpt}") + link_aux_files( + args.source_ckpt, + args.output_ckpt, + skip_top_level=_SKIP_TOP_LEVEL, + skip_dir_names=_SKIP_SUBDIR_NAMES | _SKIP_TOP_LEVEL, + skip_file=lambda path: path.suffix == ".safetensors", + ) + _write_json_atomic( + args.output_ckpt / "conversion_report.json", + { + "run_id": args.run_id, + "recipe": args.recipe, + "world_size": args.world_size, + "shards": len(results), + "stats": totals, + "expert_banks": len(banks), + "attention_modules": len(attn_modules), + "cast_mxfp4_to_nvfp4": args.cast_mxfp4_to_nvfp4, + "attn_fp8": args.attn_fp8, + "attn_mxfp8": args.attn_mxfp8, + "attn_fp8_pb": args.attn_fp8_pb, + "attention_quant_algo": attention_algo, + "input_scale": args.input_scale, + }, + ) + shutil.rmtree(marker_dir) + _log( + f"[done] {args.output_ckpt} " + f"({len(banks)} expert banks NVFP4, " + f"{len(attn_modules)} attention modules {attention_algo})" + ) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 0d18530b902..9c21ea682cd 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -283,12 +283,12 @@ def quantize( # try call trtllm fp4 quantization if possible if ( - fp4_compatible() - and weights_scaling_factor is None - and try_tensorrt + try_tensorrt and block_size == 16 and input.is_cuda and input.dtype in [torch.half, torch.bfloat16] + and weights_scaling_factor is None + and fp4_compatible() ): try: import tensorrt_llm # noqa: F401 diff --git a/modelopt_recipes/huggingface/README.md b/modelopt_recipes/huggingface/README.md index 33ea80ad348..c0361f50d7b 100644 --- a/modelopt_recipes/huggingface/README.md +++ b/modelopt_recipes/huggingface/README.md @@ -34,6 +34,9 @@ modelopt_recipes/huggingface/ .yaml [..yaml] # optional snippet helpers (see below) [README.md] # optional; describes what's model-specific + models/// + / + .yaml # exact published-checkpoint mirror ``` `` is the model-optimization workflow the recipe targets (e.g. @@ -43,6 +46,11 @@ Selecting a recipe at runtime uses the path relative to `modelopt_recipes/`, e.g. `--recipe huggingface///`. +Recipes that reproduce one exact published checkpoint may instead live under +`huggingface/models///`. This layout records the canonical +source checkpoint directly and avoids implying that the recipe applies to +every checkpoint sharing the same `model_type`. + ### Verifying a model's `model_type` The authoritative source for a model's `model_type` is the released diff --git a/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention.yaml b/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention.yaml new file mode 100644 index 00000000000..30d85bcccf6 --- /dev/null +++ b/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention.yaml @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Checkpoint-mirror recipe for nvidia/Kimi-K3-NVFP4. The source Kimi-K3 routed +# experts are already packed as MXFP4, so examples/kimi/kimi_k3/quantize_to_nvfp4.py +# applies this layout as a streaming checkpoint conversion rather than loading +# the complete model through examples/hf_ptq/hf_ptq.py. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + fp8: configs/numerics/fp8 + nvfp4: configs/numerics/nvfp4 + +metadata: + recipe_type: ptq + description: >- + Kimi-K3 checkpoint-mirror recipe with calibration-free MXFP4-to-NVFP4 + routed experts, expert input_scale fixed to 1.0, and 128x128 block-FP8 + weight-only KDA and MLA attention projections. Shared and latent experts, + routers, convolutions, norms, the vision tower, lm_head, and KV cache stay + unquantized. + +quantize: + algorithm: + method: max + layerwise: + enable: false + skip_forward_without_activation_calib: true + quant_cfg: + - $import: base_disable_all + # Routed experts are cast directly from the source MXFP4 representation. + - quantizer_name: "*block_sparse_moe.experts.*weight_quantizer" + cfg: + $import: nvfp4 + - quantizer_name: "*block_sparse_moe.experts.*input_quantizer" + cfg: + $import: nvfp4 + # 2688 = E2M1_MAX * E4M3_MAX = 6 * 448, so input_scale is exactly 1.0. + constant_amax: 2688.0 + # KDA projections. + - quantizer_name: "*self_attn.q_proj*weight_quantizer" + cfg: &fp8_pb_weight + $import: fp8 + block_sizes: + -1: 128 + -2: 128 + - quantizer_name: "*self_attn.k_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.v_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.b_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.f_a_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.f_b_proj*weight_quantizer" + cfg: *fp8_pb_weight + # MLA projections. + - quantizer_name: "*self_attn.q_a_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.q_b_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.kv_a_proj_with_mqa*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.kv_b_proj*weight_quantizer" + cfg: *fp8_pb_weight + # Present in both attention variants. + - quantizer_name: "*self_attn.o_proj*weight_quantizer" + cfg: *fp8_pb_weight + - quantizer_name: "*self_attn.g_proj*weight_quantizer" + cfg: *fp8_pb_weight + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 53c4ef99504..edc5d926810 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -347,6 +347,14 @@ everything else matches the general recipe. The `huggingface/models/` tier reproduces a **single published (or planned) checkpoint's** quant config verbatim: +- **`models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention`** mirrors + `nvidia/Kimi-K3-NVFP4`: the source MXFP4 routed experts are cast to NVFP4, + with activation `input_scale=1.0`, while KDA and MLA projection weights use + 128x128 block FP8. Attention activations are dynamic; shared and latent + experts, routers, convolutions, norms, the vision tower, `lm_head`, and KV + cache remain BF16. Because the 2.8T source uses packed MXFP4 expert tensors, + use the calibration-free streaming converter in `examples/kimi/` rather than + the in-memory `hf_ptq.py` flow. - **`models/mistralai/Mistral-Medium-3.5-128B/ptq/nvfp4-max-calib`** mirrors `nvidia/Mistral-Medium-3.5-128B-NVFP4`: decoder MLP layers 4–86 use NVFP4 W4A4, edge MLP layers 0–3 and 87 use FP8 W8A8, and all attention projections diff --git a/tests/examples/hf_ptq/test_kimi_k3_quantize_to_nvfp4.py b/tests/examples/hf_ptq/test_kimi_k3_quantize_to_nvfp4.py new file mode 100644 index 00000000000..093fbced696 --- /dev/null +++ b/tests/examples/hf_ptq/test_kimi_k3_quantize_to_nvfp4.py @@ -0,0 +1,428 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the calibration-free Kimi-K3 checkpoint converter.""" + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from modelopt.torch.quantization.qtensor import FP8QTensor, MXFP4QTensor, MXFP8QTensor + +_SCRIPT = ( + Path(__file__).resolve().parents[3] / "examples" / "kimi" / "kimi_k3" / "quantize_to_nvfp4.py" +) +_SPEC = importlib.util.spec_from_file_location("kimi_k3_quantize_to_nvfp4", _SCRIPT) +assert _SPEC is not None and _SPEC.loader is not None +k3_cast = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(k3_cast) + + +def test_published_recipe_resolves_to_streaming_conversion_settings(): + settings = k3_cast._conversion_settings_from_recipe(k3_cast._PUBLISHED_RECIPE) + + assert settings == { + "cast_mxfp4_to_nvfp4": True, + "attn_fp8": False, + "attn_mxfp8": False, + "attn_fp8_pb": True, + "input_scale": 1.0, + } + + +def test_recipe_rejects_explicit_input_scale(monkeypatch, tmp_path, capsys): + monkeypatch.setattr( + sys, + "argv", + [ + str(_SCRIPT), + "--source_ckpt", + str(tmp_path / "source"), + "--output_ckpt", + str(tmp_path / "output"), + "--recipe", + k3_cast._PUBLISHED_RECIPE, + "--input_scale", + "1.0", + ], + ) + + with pytest.raises(SystemExit): + k3_cast.main() + + assert "--recipe cannot be combined" in capsys.readouterr().err + + +def test_rank0_rendezvous_rejects_mismatched_world_size(tmp_path): + ready_path = tmp_path / "ready.json" + ready_path.write_text(json.dumps({"run_id": "run-1", "world_size": 4})) + + assert not k3_cast._rank0_ready(ready_path, "other-run", world_size=4, rank=1) + assert k3_cast._rank0_ready(ready_path, "run-1", world_size=4, rank=1) + with pytest.raises(ValueError, match="rank 1 has --world_size 2"): + k3_cast._rank0_ready(ready_path, "run-1", world_size=2, rank=1) + + +def test_module_name_aliases_strip_language_model_prefix(): + assert k3_cast._module_name_aliases("language_model.lm_head") == [ + "language_model.lm_head", + "lm_head", + ] + + +def _mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + qtensor, scale = MXFP4QTensor.quantize(weight, block_size=32) + expected_scale_shape = (*weight.shape[:-1], weight.shape[-1] // 32) + return qtensor._quantized_data, scale.reshape(expected_scale_shape) + + +def _write_source_checkpoint(tmp_path: Path) -> tuple[Path, str, dict[str, torch.Tensor]]: + source = tmp_path / "source" + source.mkdir() + shard_name = "model-00001-of-00001.safetensors" + expert_prefix = "language_model.model.layers.1.block_sparse_moe.experts.0" + + torch.manual_seed(7) + w1 = torch.randn(4, 64, dtype=torch.bfloat16) * 0.25 + w2 = torch.randn(4, 64, dtype=torch.bfloat16) * 2.0 + w3 = torch.randn(4, 64, dtype=torch.bfloat16) * 8.0 + q_proj = torch.randn(16, 64, dtype=torch.bfloat16) + b_proj = torch.randn(8, 64, dtype=torch.bfloat16) + f_a_proj = torch.randn(8, 64, dtype=torch.bfloat16) + f_b_proj = torch.randn(16, 8, dtype=torch.bfloat16) + passthrough = torch.randn(4, 4, dtype=torch.bfloat16) + + state: dict[str, torch.Tensor] = { + "language_model.model.layers.1.self_attn.q_proj.weight": q_proj, + "language_model.model.layers.1.self_attn.b_proj.weight": b_proj, + "language_model.model.layers.1.self_attn.f_a_proj.weight": f_a_proj, + "language_model.model.layers.1.self_attn.f_b_proj.weight": f_b_proj, + "language_model.model.layers.1.input_layernorm.weight": passthrough, + } + for proj, weight in (("w1", w1), ("w2", w2), ("w3", w3)): + packed, scale = _mxfp4(weight) + base = f"{expert_prefix}.{proj}" + state[base + ".weight_packed"] = packed + state[base + ".weight_scale"] = scale + + save_file(state, str(source / shard_name)) + index = { + "metadata": {"total_size": sum(t.numel() * t.element_size() for t in state.values())}, + "weight_map": dict.fromkeys(state, shard_name), + } + (source / "model.safetensors.index.json").write_text(json.dumps(index)) + (source / "config.json").write_text( + json.dumps( + { + "architectures": ["KimiK3ForConditionalGeneration"], + "text_config": { + "quantization_config": { + "quant_method": "compressed-tensors", + "format": "mxfp4-pack-quantized", + } + }, + } + ) + ) + (source / "tokenizer_config.json").write_text("{}") + return source, shard_name, state + + +def test_split_w1_w3_pair_fails_instead_of_using_independent_scales(): + base = "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + + with pytest.raises(RuntimeError, match="split across shards"): + k3_cast.build_w13_kmax_overrides( + [base], + lambda _: torch.tensor([127], dtype=torch.uint8), + "cpu", + ) + + +def test_link_aux_files_preserves_sidecars_and_skips_checkpoint_data(tmp_path): + source = tmp_path / "source" + output = tmp_path / "output" + (source / "assets").mkdir(parents=True) + (source / ".cache").mkdir() + (source / "tokenizer_config.json").write_text("{}") + (source / "model-00001-of-00001.safetensors").write_bytes(b"shard") + (source / "assets" / "config.txt").write_text("keep") + (source / "assets" / "nested.safetensors").write_bytes(b"skip") + (source / ".cache" / "stale.json").write_text("{}") + + k3_cast.link_aux_files( + source, + output, + skip_top_level=k3_cast._SKIP_TOP_LEVEL, + skip_dir_names=k3_cast._SKIP_SUBDIR_NAMES | k3_cast._SKIP_TOP_LEVEL, + skip_file=lambda path: path.suffix == ".safetensors", + ) + + assert (output / "tokenizer_config.json").read_text() == "{}" + assert (output / "assets" / "config.txt").read_text() == "keep" + assert not (output / "model-00001-of-00001.safetensors").exists() + assert not (output / "assets" / "nested.safetensors").exists() + assert not (output / ".cache").exists() + + +def test_convert_shard_casts_experts_and_quantizes_attention(tmp_path): + source, shard_name, source_state = _write_source_checkpoint(tmp_path) + output = tmp_path / "output" + output.mkdir() + + report = k3_cast.convert_shard( + source / shard_name, + output / shard_name, + device="cpu", + cast=True, + attn_fp8=True, + input_scale_value=1.0, + ) + + expert = "language_model.model.layers.1.block_sparse_moe.experts.0" + q_proj = "language_model.model.layers.1.self_attn.q_proj" + b_proj = "language_model.model.layers.1.self_attn.b_proj" + f_a_proj = "language_model.model.layers.1.self_attn.f_a_proj" + with safe_open(output / shard_name, framework="pt", device="cpu") as f: + keys = set(f.keys()) + for proj in ("w1", "w2", "w3"): + base = f"{expert}.{proj}" + assert base + ".weight_packed" not in keys + assert f.get_tensor(base + ".weight").dtype == torch.uint8 + assert f.get_tensor(base + ".weight_scale").dtype == torch.float8_e4m3fn + assert f.get_tensor(base + ".weight_scale_2").shape == torch.Size([]) + assert f.get_tensor(base + ".input_scale").item() == 1.0 + + # Fused GEMM1 requires gate/up (w1/w3) to share one global scale. + assert torch.equal( + f.get_tensor(expert + ".w1.weight_scale_2"), + f.get_tensor(expert + ".w3.weight_scale_2"), + ) + + assert f.get_tensor(q_proj + ".weight").dtype == torch.float8_e4m3fn + assert f.get_tensor(q_proj + ".weight_scale").shape == torch.Size([]) + assert f.get_tensor(q_proj + ".input_scale").item() == 1.0 + + # vLLM packs b/f_a with q/k/v/g, so the entire fused linear is FP8. + for base in (b_proj, f_a_proj): + assert f.get_tensor(base + ".weight").dtype == torch.float8_e4m3fn + assert f.get_tensor(base + ".weight_scale").shape == torch.Size([]) + assert f.get_tensor(base + ".input_scale").item() == 1.0 + assert torch.equal( + f.get_tensor("language_model.model.layers.1.input_layernorm.weight"), + source_state["language_model.model.layers.1.input_layernorm.weight"], + ) + + assert report["stats"]["experts_converted"] == 3 + assert report["stats"]["attn_fp8_converted"] == 3 + assert report["stats"]["cast_blocks_total"] == 24 + assert report["stats"]["cast_blocks_lossless"] == 24 + assert report["banks"] == ["language_model.model.layers.1.block_sparse_moe.experts"] + assert report["attn_modules"] == [b_proj, f_a_proj, q_proj] + + +def test_convert_shard_requantizes_w1_w3_with_shared_scale(tmp_path): + source, shard_name, _ = _write_source_checkpoint(tmp_path) + output = tmp_path / "output" + output.mkdir() + + report = k3_cast.convert_shard( + source / shard_name, + output / shard_name, + device="cpu", + cast=False, + attn_fp8=False, + input_scale_value=1.0, + ) + + expert = "language_model.model.layers.1.block_sparse_moe.experts.0" + with safe_open(output / shard_name, framework="pt", device="cpu") as f: + assert torch.equal( + f.get_tensor(expert + ".w1.weight_scale_2"), + f.get_tensor(expert + ".w3.weight_scale_2"), + ) + assert not torch.equal( + f.get_tensor(expert + ".w1.weight_scale_2"), + f.get_tensor(expert + ".w2.weight_scale_2"), + ) + + assert report["stats"]["experts_converted"] == 3 + assert "cast_blocks_total" not in report["stats"] + + +def test_convert_shard_quantizes_attention_to_mxfp8_without_input_scale(tmp_path): + source, shard_name, source_state = _write_source_checkpoint(tmp_path) + output = tmp_path / "output" + output.mkdir() + + report = k3_cast.convert_shard( + source / shard_name, + output / shard_name, + device="cpu", + cast=True, + attn_fp8=False, + input_scale_value=1.0, + attn_mxfp8=True, + ) + + base = "language_model.model.layers.1.self_attn.q_proj" + original = source_state[base + ".weight"] + with safe_open(output / shard_name, framework="pt", device="cpu") as f: + keys = set(f.keys()) + quantized = f.get_tensor(base + ".weight") + scale = f.get_tensor(base + ".weight_scale") + assert quantized.dtype == torch.float8_e4m3fn + assert scale.dtype == torch.uint8 + assert scale.shape == (original.shape[0], original.shape[1] // 32) + assert base + ".input_scale" not in keys + + restored = MXFP8QTensor(original.shape, original.dtype, quantized).dequantize(scale=scale) + assert torch.allclose(restored, original, rtol=0.05, atol=0.05) + + assert report["stats"]["attn_mxfp8_converted"] == 3 + config = k3_cast._build_hf_quant_config( + report["banks"], report["attn_modules"], attn_fp8=False, attn_mxfp8=True + ) + assert config["quantization"]["quantized_layers"][base] == {"quant_algo": "MXFP8"} + assert config["quantization"]["quantized_layers"]["layers.1.self_attn.in_proj_qkvgfab"] == { + "quant_algo": "MXFP8" + } + + +def test_convert_shard_quantizes_attention_to_block_fp8(tmp_path): + source, shard_name, source_state = _write_source_checkpoint(tmp_path) + output = tmp_path / "output" + output.mkdir() + + report = k3_cast.convert_shard( + source / shard_name, + output / shard_name, + device="cpu", + cast=True, + attn_fp8=False, + input_scale_value=1.0, + attn_fp8_pb=True, + ) + + bases = [ + "language_model.model.layers.1.self_attn.q_proj", + "language_model.model.layers.1.self_attn.b_proj", + "language_model.model.layers.1.self_attn.f_a_proj", + "language_model.model.layers.1.self_attn.f_b_proj", + ] + with safe_open(output / shard_name, framework="pt", device="cpu") as f: + keys = set(f.keys()) + for base in bases: + original = source_state[base + ".weight"] + quantized = f.get_tensor(base + ".weight") + scale = f.get_tensor(base + ".weight_scale") + assert quantized.dtype == torch.float8_e4m3fn + assert scale.dtype == torch.float32 + assert scale.shape == ( + (original.shape[0] + 127) // 128, + 1, + (original.shape[1] + 127) // 128, + 1, + ) + assert base + ".input_scale" not in keys + + restored = FP8QTensor(original.shape, original.dtype, quantized).dequantize( + scale=scale.squeeze(1).squeeze(-1), + block_sizes={-2: 128, -1: 128}, + ) + assert torch.allclose(restored, original, rtol=0.05, atol=0.05) + + assert report["stats"]["attn_fp8_pb_converted"] == 4 + config = k3_cast._build_hf_quant_config( + report["banks"], + report["attn_modules"], + attn_fp8=False, + attn_fp8_pb=True, + ) + quantization = config["quantization"] + assert quantization["quantized_layers"][bases[0]] == {"quant_algo": "FP8_PB_WO"} + assert quantization["quantized_layers"]["layers.1.self_attn.in_proj_qkvgfab"] == { + "quant_algo": "FP8_PB_WO" + } + assert "*self_attn.f_b_proj*" not in quantization["exclude_modules"] + + +def test_manifest_and_index_replace_source_mxfp4_schema(tmp_path): + source, shard_name, _ = _write_source_checkpoint(tmp_path) + output = tmp_path / "output" + output.mkdir() + report = k3_cast.convert_shard( + source / shard_name, + output / shard_name, + device="cpu", + cast=True, + attn_fp8=True, + input_scale_value=1.0, + ) + + hf_quant_config = k3_cast._build_hf_quant_config( + report["banks"], report["attn_modules"], attn_fp8=True + ) + source_index = json.loads((source / "model.safetensors.index.json").read_text()) + k3_cast._write_index_and_manifest( + output, + source_index, + [report], + hf_quant_config, + attn_fp8=True, + ) + k3_cast._rewrite_config_json(source, output, hf_quant_config) + + index = json.loads((output / "model.safetensors.index.json").read_text()) + weight_map = index["weight_map"] + expert = "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + assert expert + ".weight_packed" not in weight_map + assert weight_map[expert + ".weight"] == shard_name + assert weight_map[expert + ".weight_scale"] == shard_name + assert weight_map[expert + ".weight_scale_2"] == shard_name + assert weight_map[expert + ".input_scale"] == shard_name + assert index["metadata"]["total_size"] == report["tensor_bytes"] + + config = json.loads((output / "config.json").read_text()) + assert "quantization_config" not in config["text_config"] + quant = config["quantization_config"] + assert quant["quant_method"] == "modelopt_mixed" + assert quant["quant_algo"] == "MIXED_PRECISION" + assert len(quant["config_groups"]) == 2 + assert quant["quantized_layers"]["language_model.model.layers.1.block_sparse_moe.experts"] == { + "quant_algo": "NVFP4", + "group_size": 16, + } + assert quant["quantized_layers"]["language_model.model.layers.1.self_attn.q_proj"] == { + "quant_algo": "FP8" + } + assert quant["quantized_layers"]["layers.1.block_sparse_moe.experts"] == { + "quant_algo": "NVFP4", + "group_size": 16, + } + assert quant["quantized_layers"]["layers.1.mlp.experts"] == { + "quant_algo": "NVFP4", + "group_size": 16, + } + assert quant["quantized_layers"]["layers.1.self_attn.q_proj"] == {"quant_algo": "FP8"} + assert quant["quantized_layers"]["layers.1.self_attn.b_proj"] == {"quant_algo": "FP8"} + assert quant["quantized_layers"]["layers.1.self_attn.f_a_proj"] == {"quant_algo": "FP8"} + assert quant["quantized_layers"]["layers.1.self_attn.in_proj_qkvgfab"] == {"quant_algo": "FP8"} diff --git a/tests/unit/recipe/test_kimi_k3_recipe.py b/tests/unit/recipe/test_kimi_k3_recipe.py new file mode 100644 index 00000000000..1041021981c --- /dev/null +++ b/tests/unit/recipe/test_kimi_k3_recipe.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Kimi-K3 checkpoint-mirror PTQ recipe.""" + +from modelopt.recipe import load_recipe + +RECIPE = "huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention" + + +def test_kimi_k3_recipe_matches_published_quantization_map(): + config = load_recipe(RECIPE).quantize.model_dump() + quant_cfg = config["quant_cfg"] + + by_name = { + entry["quantizer_name"]: entry + for entry in quant_cfg + if isinstance(entry, dict) and "quantizer_name" in entry + } + + expert_input = by_name["*block_sparse_moe.experts.*input_quantizer"]["cfg"] + assert expert_input["constant_amax"] == 2688.0 + + attention_projections = { + "q_proj", + "k_proj", + "v_proj", + "b_proj", + "f_a_proj", + "f_b_proj", + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "g_proj", + } + for projection in attention_projections: + weight_cfg = by_name[f"*self_attn.{projection}*weight_quantizer"]["cfg"] + assert weight_cfg["num_bits"] == (4, 3) + assert weight_cfg["block_sizes"] == {-1: 128, -2: 128} + + assert not any("kv_cache" in name for name in by_name) diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py index 8523edc4052..a2a977f45b1 100644 --- a/tests/unit/torch/quantization/test_nvfp4_tensor.py +++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py @@ -17,8 +17,10 @@ from types import SimpleNamespace +import pytest import torch +import modelopt.torch.quantization.qtensor.nvfp4_tensor as nvfp4_tensor from modelopt.torch.quantization.qtensor.nvfp4_tensor import ( NVFP4QTensor, _cast_per_block_scale_to_fp8, @@ -28,6 +30,21 @@ _FP8_E4M3FN_MAX = 448.0 +@pytest.mark.parametrize("try_tensorrt", [False, True]) +def test_cpu_quantize_does_not_probe_cuda(monkeypatch, try_tensorrt): + """The optional TRT-LLM path must not query GPU capability for a CPU tensor.""" + + def fail_if_called(): + pytest.fail("fp4_compatible() was called for CPU quantization") + + monkeypatch.setattr(nvfp4_tensor, "fp4_compatible", fail_if_called) + weight = torch.randn(2, 16, dtype=torch.bfloat16) + + quantized, _, _ = NVFP4QTensor.quantize(weight, block_size=16, try_tensorrt=try_tensorrt) + + assert quantized._quantized_data.shape == (2, 8) + + class TestNVFP4ScaleClamping: """Per-block weight scales outside the FP8 E4M3FN range must be clamped, not turned into 0/NaN."""