NVFP4 Example

Kimi-K3 NVFP4 Example

Overview

Kimi-K3 requires custom modeling files bundled with LLM Compressor, since it is not yet supported in Transformers. The example below quantizes the model to NVFP4 using calibration data.

This model has already been quantized to mxfp4, so we must first dequantize it before applying a new quantization

The full example script can be found here.

Code Walkthrough

# requires: einops, fla-core, tiktoken
import torch.distributed as dist
from compressed_tensors.compressors import ModelCompressor
from compressed_tensors.distributed import init_dist
from transformers import AutoConfig, AutoProcessor, CompressedTensorsConfig

from llmcompressor import oneshot
from llmcompressor.datasets.utils import get_rank_partition
from llmcompressor.modeling.kimi_k3 import KimiK3ForConditionalGeneration
from llmcompressor.modeling.patch.kimi_k3_patch import patch_kimi_k3_ignore
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.utils import load_context

# Small representative model with same MXFP4 quantization
MODEL_ID = "inference-optimization/Kimi-K3-0.40B-MXFP4"  # "moonshotai/Kimi-K3"

# Patch quantization config to disable quantization (for a later step)
config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
qconfig = CompressedTensorsConfig(dequantize=False, use_optimized_inference=False)

# Load model with the modified quantization config and disk offloading
# Patch an incomplete ignore list provided by the base checkpoint
init_dist()
with load_context(KimiK3ForConditionalGeneration), patch_kimi_k3_ignore():
    model = KimiK3ForConditionalGeneration.from_pretrained(
        MODEL_ID,
        quantization_config=qconfig,
        device_map="auto_offload",
        trust_remote_code=True,
        max_memory={},
        offload_folder="offload_folder",
    )
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)

# Decompress model upfront before calibrating
ModelCompressor.from_pretrained_model(model).decompress_model(model)

recipe = QuantizationModifier(
    targets="re:.*block_sparse_moe.*",
    scheme="NVFP4",
    ignore=[
        r"re:.*block_sparse_moe\.gate",
        "re:.*routed_expert.*",
        "re:.*mlp_res_proj$",
        "re:.*vision_tower.*",
        "lm_head",
    ],
    weight_observer="nvfp4_expanded_mse",
)

oneshot(
    model=model,
    tokenizer=processor.tokenizer,
    dataset="perfectblend",
    splits=get_rank_partition("train", 512),
    recipe=recipe,
    max_seq_length=2048,
    trust_remote_code_model=True,
    batch_size=16,
    shuffle_calibration_samples=False,
)

SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
model.save_pretrained(SAVE_DIR)
processor.save_pretrained(SAVE_DIR)

dist.destroy_process_group()