Views
No views yet
1import os
2import torch
3from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
4from torchao.quantization import Int8DynamicActivationInt8WeightConfig
5from torchao.quantization.quant_primitives import MappingType
6
7model_name = "OpenPipe/Qwen3-14B-Instruct"
8output_dir = "./Qwen3-14B-Instruct-da8w8-torchao-v0.16.0"
9os.makedirs(output_dir, exist_ok=True)
10modules_to_skip = ["lm_head"]
11
12# Step 1: Create quantization config
13quantization_config = TorchAoConfig(
14 Int8DynamicActivationInt8WeightConfig(
15 version=2,
16 act_mapping_type=MappingType.SYMMETRIC,
17 ),
18 modules_to_not_convert=modules_to_skip,
19)
20
21# Step 2: Load and quantize the model
22quantized_model = AutoModelForCausalLM.from_pretrained(
23 model_name,
24 dtype=torch.bfloat16,
25 device_map="cpu",
26 quantization_config=quantization_config,
27 trust_remote_code=True,
28)
29
30# Step 3: Save the quantized model (normal save; no in-place qdata change)
31quantized_model.save_pretrained(output_dir, safe_serialization=False)
32
33# Step 4: Save the tokenizer
34tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
35tokenizer.save_pretrained(output_dir)
36
37# Smoke test
38input_text = "What are we having for dinner?"
39input_ids = tokenizer(input_text, return_tensors="pt")
40output = quantized_model.generate(
41 **input_ids, max_new_tokens=30, cache_implementation="static"
42)
43print(tokenizer.decode(output[0], skip_special_tokens=True))Notes
safe_serialization=Falseis required because torchao's quantized tensor subclasses cannot currently be serialized in thesafetensorsformat.lm_headis excluded from quantization to preserve final-projection precision.
1pip install --extra-index-url https://download.pytorch.org/whl/cpu \
2 --extra-index-url https://wheels.vllm.ai/cpu/ \
3 torch==2.10.0+cpu \
4 vllm==0.18.0 \
5 torchao==0.16.0 \
6 transformers \
7 huggingface_hub1# vLLM CPU runtime tuning
2export VLLM_CPU_KVCACHE_SPACE=40 # GB of host memory for KV cache
3export VLLM_CPU_OMP_THREADS_BIND="0-63" # NUMA-local cores
4
5# TorchInductor
6export TORCHINDUCTOR_FREEZING=1
7export TORCHINDUCTOR_AUTOGRAD_CACHE=1
8export TORCHINDUCTOR_CACHE_DIR="./.torchinductor_cache/Qwen3-14B-Instruct-da8w8-torchao-v0.16.0"
9
10# Required CPU runtime libraries
11export LD_PRELOAD="<path to lib>/libtcmalloc_minimal.so.4:<path to lib>/libiomp5.so${LD_PRELOAD:+:$LD_PRELOAD}"find / -name 'libtcmalloc_minimal.so.4' and find / -name 'libiomp5.so', then substitute the resulting directory for <path to lib>.| Benchmark | BF16 Baseline | DA8W8 (this model) | Dynamic Quant Difference (baseline: BF16) |
|---|---|---|---|
| GSM8K (5-shot, exact-match strict) | 0.8795 | 0.8855 | +0.68% |
1mkdir -p "${TORCHINDUCTOR_CACHE_DIR}"
2
3lm_eval \
4 --model vllm \
5 --model_args pretrained=amd/Qwen3-14B-Instruct-da8w8-torchao-v0.16.0,tokenizer=OpenPipe/Qwen3-14B-Instruct,dtype=bfloat16 \
6 --tasks gsm8k \
7 --batch_size auto \
8 --trust_remote_code \
9 --num_fewshot 5 \
10 --log_samples \
11 --apply_chat_template \
12 --fewshot_as_multiturn \
13 --gen_kwargs "temperature=0.6,top_p=0.95,top_k=201" \
14 --output_path .