Views
No views yet
[!NOTE] zentorch v5.2.1 for PyTorch v2.10.0 has to be built from source.
1import torch
2from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
3from torchao.quantization import Int8DynamicActivationInt8WeightConfig
4from torchao.quantization.quant_primitives import MappingType
5
6MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
7OUTPUT_DIR = "amd/Llama-3.1-8B-Instruct-da8w8-torchao-v0.16.0"
8modules_to_skip = ["lm_head"]
9
10quantization_config = TorchAoConfig(
11 Int8DynamicActivationInt8WeightConfig(
12 version=2,
13 act_mapping_type=MappingType.SYMMETRIC,
14 ),
15 modules_to_not_convert=modules_to_skip,
16)
17
18model = AutoModelForCausalLM.from_pretrained(
19 MODEL_ID,
20 dtype=torch.bfloat16,
21 device_map="cpu",
22 quantization_config=quantization_config,
23 trust_remote_code=True,
24)
25model.save_pretrained(OUTPUT_DIR, safe_serialization=False)
26
27tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
28tokenizer.save_pretrained(OUTPUT_DIR)
29
30# Smoke test
31inputs = tokenizer("What are we having for dinner?", return_tensors="pt")
32out = model.generate(**inputs, max_new_tokens=30, cache_implementation="static")
33print(tokenizer.decode(out[0], skip_special_tokens=True))[!NOTE]safe_serialization=Falseis required because torchao's quantized tensor subclasses cannot currently be serialized in thesafetensorsformat.
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_hubconda install -c conda-forge gperftools=2.17.2 llvm-openmp=18.1.8 --no-deps -y1# 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
8
9# Required CPU runtime libraries
10export 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.8453 | 0.8279 | -2.06% |
1lm_eval \
2 --model vllm \
3 --model_args pretrained=amd/Llama-3.1-8B-Instruct-da8w8-torchao-v0.16.0,tokenizer=meta-llama/Llama-3.1-8B-Instruct,dtype=bfloat16 \
4 --tasks gsm8k \
5 --batch_size auto \
6 --trust_remote_code \
7 --num_fewshot 5 \
8 --log_samples \
9 --gen_kwargs "max_gen_toks=2048" \
10 --apply_chat_template \
11 --output_path .