[WIP] [Kimi-K3]

examples/model_free_ptq/kimi_k3_fp8_block.py

@@ -0,0 +1,33 @@
+from compressed_tensors.entrypoints.convert import CompressedTensorsDequantizer
+
+from llmcompressor import model_free_ptq
+
+MODEL_ID = "moonshotai/Kimi-K3"
+SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-BLOCK"
+
+# no attention because (q_proj|k_proj|v_proj|b_proj|f_a_proj) are all fused
+# and `b_proj` has weight shape [96, 7168] which is not divisible by 128
+ignore = [
+    "re:.*embed_tokens.*",
+    "re:.*self_attn.*",
+    "re:.*block_sparse_moe\.gate.*",
+    "re:.*self_attention_res_proj.*",
+    "re:.*mlp_res_proj.*",
+    "re:.*output_attn_res_proj.*",
+    "re:.*lm_head.*",
+    "re:.*vision_tower.*",
+    "re:.*mm_projector.*",
+]
+
+model_free_ptq(
+    model_stub=MODEL_ID,
+    save_directory=SAVE_DIR,
+    scheme="FP8_BLOCK",
+    ignore=ignore,
+    converter=CompressedTensorsDequantizer(
+        MODEL_ID,
+        ignore=ignore,
+    ),
+    max_workers=7,
+    device=[f"cuda:{i}" for i in range(7)],
+)

examples/quantization_w4a4_fp4/README.md

@@ -83,3 +83,28 @@ tokenizer.save_pretrained(SAVE_DIR)
 ```
 
 We have successfully created an `nvfp4` model!
+
+## Kimi-K3 (MoE + Multimodal) Example
+
+`kimi_k3_nvfp4.py` shows how to apply NVFP4 quantization to [Kimi-K3](https://huggingface.co/inference-optimization/Kimi-K3-0.18B), a multimodal Mixture-of-Experts model with a hybrid attention architecture (KDA linear attention + MLA full attention).
+
+### Model Architecture
+
+| Property | Value |
+|---|---|
+| Architecture | Kimi-K3 (MoE + custom attention) |
+| Total parameters | 0.18B |
+| Active parameters | ~0.10B |
+| Experts | 8 total, 2 active per token |
+| Layers | 4 (dense FFN, MoE FFN, KDA linear, MLA full attention) |
+
+### Key Considerations
+
+- **Custom model class**: Kimi-K3 uses `KimiK3ForConditionalGeneration` (not `AutoModelForCausalLM`), which requires `load_context` and `trust_remote_code=True`.
+- **MoE gate excluded**: The `block_sparse_moe.gate` layer (`KimiMoEGate`) controls expert routing and is excluded from quantization to prevent routing degradation.
+- **Vision tower excluded**: All `vision_tower.*` layers are excluded to preserve multimodal capabilities.
+- **Calibration samples**: 512 samples are used (MoE models benefit from more samples to ensure all experts are well-calibrated).
+
+```bash
+python3 kimi_k3_nvfp4.py
+```

examples/quantization_w4a4_fp4/kimi_k3_nvfp4.py

@@ -0,0 +1,96 @@
+from compressed_tensors.offload import dispatch_model
+from transformers import AutoTokenizer
+
+from datasets import load_dataset
+from llmcompressor import oneshot
+from llmcompressor.modeling.kimi_k3 import KimiK3ForConditionalGeneration
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.utils import load_context
+
+MODEL_ID = "inference-optimization/Kimi-K3-0.18B"
+
+# Load model.
+with load_context(KimiK3ForConditionalGeneration):
+    model = KimiK3ForConditionalGeneration.from_pretrained(
+        MODEL_ID,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+DATASET_ID = "HuggingFaceH4/ultrachat_200k"
+DATASET_SPLIT = "train_sft"
+
+# Select number of samples. 512 samples is a good place to start.
+# MoE models benefit from more samples for better expert calibration.
+NUM_CALIBRATION_SAMPLES = 512
+MAX_SEQUENCE_LENGTH = 2048
+
+# Load dataset and preprocess.
+ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
+ds = ds.shuffle(seed=42)
+
+
+def preprocess(example):
+    return {
+        "text": tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+        )
+    }
+
+
+ds = ds.map(preprocess)
+
+
+# Tokenize inputs.
+def tokenize(sample):
+    return tokenizer(
+        sample["text"],
+        padding=False,
+        max_length=MAX_SEQUENCE_LENGTH,
+        truncation=True,
+        add_special_tokens=False,
+    )
+
+
+ds = ds.map(tokenize, remove_columns=ds.column_names)
+
+# Configure the quantization algorithm and scheme.
+# In this case, we:
+#   * quantize the weights to fp4 with per-group-16 scaling
+#   * quantize the activations to fp4 with calibrated global scale
+# The MoE gate (KimiMoEGate) and vision tower are excluded from quantization.
+recipe = QuantizationModifier(
+    targets="Linear",
+    scheme="NVFP4",
+    ignore=[
+        "lm_head",
+        r"re:.*block_sparse_moe\.gate",
+        "re:.*vision_tower.*",
+    ],
+)
+
+# Apply quantization.
+oneshot(
+    model=model,
+    dataset=ds,
+    recipe=recipe,
+    max_seq_length=MAX_SEQUENCE_LENGTH,
+    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
+)
+
+print("\n\n")
+print("========== SAMPLE GENERATION ==============")
+dispatch_model(model)
+input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to(
+    model.device
+)
+output = model.generate(input_ids, max_new_tokens=100)
+print(tokenizer.decode(output[0]))
+print("==========================================\n\n")
+
+# Save to disk in compressed-tensors format.
+SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4"
+model.save_pretrained(SAVE_DIR)
+tokenizer.save_pretrained(SAVE_DIR)

examples/quantization_w4a4_fp4/qwen3_5_example.py

@@ -1,8 +1,8 @@
 import torch
 from compressed_tensors.utils import save_mtp_tensors_to_checkpoint
-from datasets import load_dataset
 from transformers import AutoProcessor, Qwen3_5MoeForConditionalGeneration
 
+from datasets import load_dataset
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import QuantizationModifier
 from llmcompressor.utils import load_context

examples/quantization_w4a4_fp4/qwen3_6_example.py

@@ -1,8 +1,8 @@
 import torch
 from compressed_tensors.utils import save_mtp_tensors_to_checkpoint
-from datasets import load_dataset
 from transformers import AutoProcessor, Qwen3_5MoeForConditionalGeneration
 
+from datasets import load_dataset
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import QuantizationModifier
 from llmcompressor.utils import load_context

examples/quantizing_moe/deepseek_v4_pro_example.py

@@ -5,12 +5,12 @@
     NVFP4,
     QuantizationScheme,
 )
-from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.models.deepseek_v4.modeling_deepseek_v4 import (
     DeepseekV4PreTrainedModel,
 )
 
+from datasets import load_dataset
 from llmcompressor import oneshot
 from llmcompressor.datasets.utils import get_rank_partition
 from llmcompressor.modifiers.quantization import QuantizationModifier

examples/quantizing_moe/glm5_example.py

@@ -5,9 +5,9 @@
     NVFP4,
     QuantizationScheme,
 )
-from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from datasets import load_dataset
 from llmcompressor import oneshot
 from llmcompressor.datasets.utils import get_rank_partition
 from llmcompressor.modifiers.quantization import QuantizationModifier

examples/quantizing_moe/glm5_gptq_example.py

@@ -1,6 +1,6 @@
-from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from datasets import load_dataset
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.utils import load_context

examples/quantizing_moe/inkling_example.py

@@ -0,0 +1,121 @@
+import torch
+from compressed_tensors.offload import init_dist
+from compressed_tensors.quantization.quant_scheme import (
+    FP8_BLOCK,
+    NVFP4,
+    QuantizationScheme,
+)
+from transformers import AutoTokenizer, InklingForConditionalGeneration
+
+from datasets import load_dataset
+from llmcompressor import oneshot
+from llmcompressor.datasets.utils import get_rank_partition
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.utils import load_context
+
+# Load the model
+init_dist()
+# if torch.distributed.get_rank() == 0:
+#     torch.cuda.memory._record_memory_history(max_entries=10000000)
+# model_id = "thinkingmachines/Inkling"
+model_id = "inference-optimization/Inkling-0.6B-A0.6B"
+with load_context(InklingForConditionalGeneration):
+    model = InklingForConditionalGeneration.from_pretrained(
+        model_id,
+        device_map="auto_offload",
+        max_memory={},
+        offload_folder="/mnt/nvme-data/engine/kylesayrs/offload_folder",
+    )
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# Select calibration dataset.
+DATASET_ID = "HuggingFaceH4/ultrachat_200k"
+DATASET_SPLIT = "train_sft"
+
+# Select number of samples. 512 samples is a good place to start.
+# Increasing the number of samples can improve accuracy.
+NUM_CALIBRATION_SAMPLES = 4  # 512
+MAX_SEQUENCE_LENGTH = 1024  # 2048
+
+# Load dataset and preprocess.
+ds = load_dataset(
+    DATASET_ID, split=get_rank_partition(DATASET_SPLIT, NUM_CALIBRATION_SAMPLES)
+)
+ds = ds.shuffle(seed=42)
+
+
+def preprocess(example):
+    return {
+        "text": tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+        )
+    }
+
+
+ds = ds.map(preprocess)
+
+
+# Tokenize inputs.
+def tokenize(sample):
+    return tokenizer(
+        sample["text"],
+        padding=False,
+        max_length=MAX_SEQUENCE_LENGTH,
+        truncation=True,
+        add_special_tokens=False,
+    )
+
+
+ds = ds.map(tokenize, remove_columns=ds.column_names)
+
+# Configure the quantization algorithm to run.
+recipe = QuantizationModifier(
+    config_groups={
+        "attention": QuantizationScheme(
+            targets=[r"re:.*attn\..*"],
+            **FP8_BLOCK,
+        ),
+        "mlp": QuantizationScheme(
+            targets=[r"re:.*mlp\..*"],
+            **NVFP4,
+        ),
+    },
+    ignore=[
+        r"re:.*sconv$",
+        r"re:.*mlp\.gate$",  # technically not necessary `InklingTopkRouter`
+        r"re:.*shared_experts.*",
+        r"re:audio_tower.*",
+        r"re:vision_tower.*",
+    ],
+)
+
+try:
+    # Apply algorithms.
+    num_experts = getattr(model.config, "n_routed_experts", 256)
+    oneshot(
+        model=model,
+        dataset=ds,
+        batch_size=1,
+        recipe=recipe,
+        shuffle_calibration_samples=False,
+        # sequential_targets=["InklingAttention", "ExpertMLP"],
+        # sequential_targets_per_subgraph=(num_experts // 4 + 10),
+    )
+finally:
+    # if torch.distributed.get_rank() == 0:
+    #     torch.cuda.memory._dump_snapshot("inkling_memory.pickle")
+    pass
+
+# Save to disk compressed.
+# Note: base checkpoint generation_config needs fixing for newer transformers versions
+model.generation_config.top_p = None
+SAVE_DIR = (
+    "/mnt/nvme-data/engine/kylesayrs/"
+    + model_id.rstrip("/").split("/")[-1]
+    + "-NVFP4-FP8"
+)
+model.save_pretrained(SAVE_DIR, save_compressed=True, save_original_format=False)
+tokenizer.save_pretrained(SAVE_DIR)
+
+torch.distributed.destroy_process_group()