From 224bbd03a2ef167a7f6c896b620085392a0ea1bd Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Mon, 17 Aug 2026 16:47:47 -0700 Subject: [PATCH] Add Kimi-K3 NVFP4 FP8-PB recipe Signed-off-by: Zhiyu Cheng --- CHANGELOG.rst | 2 + examples/hf_ptq/README.md | 10 +- examples/kimi/README.md | 36 + examples/kimi/kimi_k3/quantize_to_nvfp4.py | 1137 +++++++++++++++++ .../quantization/qtensor/nvfp4_tensor.py | 6 +- modelopt_recipes/huggingface/README.md | 8 + .../ptq/nvfp4_experts-fp8_pb_attention.yaml | 84 ++ modelopt_recipes/ptq.md | 8 + .../hf_ptq/test_kimi_k3_quantize_to_nvfp4.py | 320 +++++ tests/unit/recipe/test_kimi_k3_recipe.py | 55 + .../torch/quantization/test_nvfp4_tensor.py | 16 + 11 files changed, 1678 insertions(+), 4 deletions(-) create mode 100644 examples/kimi/README.md create mode 100644 examples/kimi/kimi_k3/quantize_to_nvfp4.py create mode 100644 modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention.yaml create mode 100644 tests/examples/hf_ptq/test_kimi_k3_quantize_to_nvfp4.py create mode 100644 tests/unit/recipe/test_kimi_k3_recipe.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ec9e7ffdf5e..fd8abf55e59 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 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. *Megatron Framework (M-LM / M-Bridge)* @@ -33,6 +34,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``. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 72b2e2348ab..17a56b1d5ed 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_sq` 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..e657b4d1b51 --- /dev/null +++ b/examples/kimi/kimi_k3/quantize_to_nvfp4.py @@ -0,0 +1,1137 @@ +# 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 errno +import json +import math +import multiprocessing +import os +import re +import shutil +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 + +from modelopt import __version__ as modelopt_version +from modelopt.recipe import load_recipe +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.quantization.qtensor import FP8QTensor, MXFP4QTensor, MXFP8QTensor, 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, +) + +# -------------------------------------------------------------------------- +# 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) +_MXFP4_BYTES_PER_BLOCK = 16 # 32 E2M1 nibbles packed 2-per-byte +_MXFP4_BLOCK = 32 + +_PUBLISHED_RECIPE = "huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention" + + +def _log(msg: str) -> None: + print(msg, flush=True) + + +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, + } + + +# -------------------------------------------------------------------------- +# MXFP4 -> NVFP4 (routed experts) +# -------------------------------------------------------------------------- +def _dequantize_mxfp4_to_bf16( + mxfp4_weight: torch.Tensor, mxfp4_scale: torch.Tensor, device: str +) -> torch.Tensor: + 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 _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_bases: 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 base in expert_bases: + prefix, proj = base.rsplit(".", 1) + if proj in {"w1", "w3"}: + groups[prefix][proj] = base + + 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"] + ".weight_scale"), device) + k3 = _kmax_from_mxfp4_scale(f.get_tensor(paths["w3"] + ".weight_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. 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) + 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**E4M3_KMIN, max=E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + + # Lossless accounting: a block is lossy only if k_max - k_j > 17. All-zero + # blocks (e8m0 == 0) reconstruct to 0 regardless of scale, 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 + + +# -------------------------------------------------------------------------- +# 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} + + w13_kmax = _build_w13_kmax_overrides(f, expert_bases, device) if cast else {} + + 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 = _kmax_from_mxfp4_scale(s, device) + packed, weight_scale, weight_scale_2, n_blk, n_lossless = ( + _quantize_weight_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: + bf16 = _dequantize_mxfp4_to_bf16(w, s, device) + weight_scale_2 = (bf16.abs().max().float() / (E2M1_MAX * E4M3_MAX)).reshape(()) + q_tensor, weight_scale, _ = NVFP4QTensor.quantize( + bf16, _NVFP4_BLOCK, None, weight_scale_2, try_tensorrt=False + ) + packed = q_tensor._quantized_data + + 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__"} + + +def _link_or_copy(src: Path, dst: Path) -> None: + try: + os.link(src, dst) + except OSError as e: + if e.errno in {errno.EXDEV, errno.EPERM, errno.EACCES, errno.EMLINK}: + shutil.copy2(src, dst) + else: + raise + + +def _hard_link_aux(src_dir: Path, target: Path) -> None: + """Hard-link (or copy) tokenizer/modeling/processor files into the output.""" + for root, dirs, files in os.walk(src_dir): + dirs[:] = [d for d in dirs if d not in _SKIP_SUBDIR_NAMES and d not in _SKIP_TOP_LEVEL] + rel = Path(root).relative_to(src_dir) + (target / rel).mkdir(parents=True, exist_ok=True) + for fname in files: + if rel == Path(".") and fname in _SKIP_TOP_LEVEL: + continue + if fname.endswith(".safetensors"): + continue # rewritten shards + src_f = Path(root) / fname + dst_f = target / rel / fname + if dst_f.exists(): + dst_f.unlink() + _link_or_copy(src_f, dst_f) + + +# 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) + 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 _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 _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 + ) + 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=1.0, + help="fixed activation input_scale for every quantized module (no calibration)", + ) + 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: + p.error("--recipe cannot be combined with the low-level format flags") + 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 not math.isfinite(args.input_scale) or args.input_scale <= 0: + p.error("--input_scale must be finite and > 0") + 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: ( + ready_path.exists() + and json.loads(ready_path.read_text()).get("run_id") == args.run_id + ), + 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}") + _hard_link_aux(args.source_ckpt, args.output_ckpt) + _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 a1fb2d94c57..3873ddc52fe 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -340,6 +340,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..7a897748a5c --- /dev/null +++ b/tests/examples/hf_ptq/test_kimi_k3_quantize_to_nvfp4.py @@ -0,0 +1,320 @@ +# 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 +from pathlib import Path + +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 _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_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_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..726e543b6fc 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,20 @@ _FP8_E4M3FN_MAX = 448.0 +def test_cpu_quantize_does_not_probe_cuda(monkeypatch): + """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=False) + + 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."""