Views
No views yet
lm_head are kept in the original precision)QuarkW4A8Fp8MoEMethod), which dispatches through the ROCm AITER fused MoE
kernel. It is not intended to produce meaningful text — it exists so CI can load
a real W4A8 checkpoint and run a forward pass on GPU.2048, MoE intermediate 1024, 8 experts, top-2) are
multiples of 256 so the AITER W4A8 shuffle/GEMM tile constraints hold. The
vocab_size matches the tokenizer so token ids stay within the embedding table.Note: Quark quantizesnn.Linearmodules. MoE experts are stored as individualnn.Linearlayers intransformers~4.57; quantize with that version so the routed experts are captured.
1import argparse
2
3import torch
4from datasets import load_dataset
5from torch.utils.data import DataLoader
6from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
7
8from quark.torch import ModelQuantizer, export_safetensors
9from quark.torch.quantization.config.config import (
10 FP8E4M3PerTensorSpec,
11 Int4PerChannelSpec,
12 ProgressiveSpec,
13 QConfig,
14 QLayerConfig,
15)
16
17
18def get_config() -> QConfig:
19 # Quantize the routed experts only.
20 exclude_layers = ["*self_attn*", "*mlp.gate", "*lm_head"]
21 input_spec = FP8E4M3PerTensorSpec(
22 observer_method="min_max", scale_type="float", is_dynamic=True
23 ).to_quantization_spec()
24 # Progressive FP8 -> INT4 weight spec (Kimi-K2.5-W4A8 recipe).
25 weight_spec = ProgressiveSpec(
26 first_stage=FP8E4M3PerTensorSpec(
27 observer_method="min_max", scale_type="float", is_dynamic=False
28 ),
29 second_stage=Int4PerChannelSpec(
30 symmetric=True,
31 scale_type="float",
32 round_method="half_even",
33 is_dynamic=False,
34 ch_axis=0,
35 ),
36 ).to_quantization_spec()
37 return QConfig(
38 global_quant_config=QLayerConfig(input_tensors=input_spec, weight=weight_spec),
39 exclude=exclude_layers,
40 )
41
42
43def main() -> None:
44 parser = argparse.ArgumentParser()
45 parser.add_argument("--export-path", required=True)
46 parser.add_argument("--tokenizer", default="Qwen/Qwen1.5-MoE-A2.7B-Chat")
47 parser.add_argument("--hidden", type=int, default=2048)
48 parser.add_argument("--moe-intermediate", type=int, default=1024)
49 parser.add_argument("--experts", type=int, default=8)
50 parser.add_argument("--topk", type=int, default=2)
51 parser.add_argument("--layers", type=int, default=2)
52 parser.add_argument("--seed", type=int, default=0)
53 args = parser.parse_args()
54
55 torch.manual_seed(args.seed)
56 tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
57
58 # vocab_size MUST cover the tokenizer, else real prompts produce token ids
59 # beyond the embedding table -> out-of-bounds embedding lookup (GPU fault).
60 cfg = AutoConfig.for_model(
61 "qwen3_moe",
62 hidden_size=args.hidden,
63 intermediate_size=args.hidden,
64 moe_intermediate_size=args.moe_intermediate,
65 num_hidden_layers=args.layers,
66 num_attention_heads=16,
67 num_key_value_heads=2,
68 head_dim=128,
69 num_experts=args.experts,
70 num_experts_per_tok=args.topk,
71 vocab_size=len(tokenizer),
72 max_position_embeddings=2048,
73 )
74 model = AutoModelForCausalLM.from_config(cfg).to("cuda").eval().to(torch.bfloat16)
75
76 ds = load_dataset("mit-han-lab/pile-val-backup", split="validation")
77 samples = [
78 tokenizer(ds[i]["text"], return_tensors="pt", truncation=True,
79 max_length=64).input_ids.to("cuda")
80 for i in range(8)
81 ]
82 dataloader = DataLoader(samples, batch_size=1)
83
84 quantizer = ModelQuantizer(get_config())
85 with torch.no_grad():
86 model = quantizer.quantize_model(model, dataloader)
87
88 export_safetensors(
89 model, args.export_path, custom_mode="quark",
90 weight_format="real_quantized", pack_method="reorder",
91 )
92 tokenizer.save_pretrained(args.export_path)
93 # Symmetric INT4 export emits all-zero `*_zero_point_2` tensors that vLLM's
94 # W4A8 loader does not expect; drop them so the checkpoint loads directly.
95
96
97if __name__ == "__main__":
98 main()1VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 \
2 vllm serve amd/tiny-qwen3-moe-w4a8 --enforce-eager