Transformers documentation
torchao
You are viewing main version, which requires installation from source. If you'd like
regular pip install, checkout the latest stable version (v5.22.2).
torchao
torchao adalah pustaka optimasi arsitektur PyTorch dengan dukungan untuk tipe data berkinerja tinggi kustom, kuantisasi, dan sparsity. Pustaka ini dapat dikomposisikan dengan fitur native PyTorch seperti torch.compile untuk inferensi dan pelatihan yang lebih cepat.
Lihat tabel di bawah ini untuk fitur torchao tambahan.
| Fitur | Deskripsi |
|---|---|
| Pelatihan yang Menyadari Kuantisasi (QAT) | Latih model terkuantisasi dengan kehilangan akurasi minimal (lihat README QAT) |
| Pelatihan Float8 | Pelatihan berkapasitas tinggi dengan format float8 (lihat dokumentasi torchtitan dan Accelerate) |
| Dukungan Sparsity | Sparsity semi-terstruktur (2:4) untuk inferensi yang lebih cepat (lihat blog post Accelerating Neural Network Training with Semi-Structured (2:4) Sparsity) |
| Kuantifikasi Optimizer | Kurangi memori state optimizer dengan varian 4 dan 8-bit dari Adam |
| Kuantisasi KV Cache | Mengaktifkan inferensi konteks panjang dengan memori yang lebih rendah (lihat KV Cache Quantization) |
| Dukungan Kernel Kustom | gunakan ops yang kompatibel dengan torch.compile Anda sendiri |
| FSDP2 | Dapat dikomposisikan dengan FSDP2 untuk pelatihan |
Lihat README.md torchao untuk detail lebih lanjut tentang pustaka ini.
torchao mendukung teknik kuantisasi berikut.
- Kuantisasi Dinamis Float8 A16W8
- Kuantisasi WeightOnly Float8 A16W8
- Kuantisasi Dinamis Int8 A8W8
- Kuantisasi Hanya Bobot Int8 A16W8
- Kuantisasi Hanya Bobot Int4 A16W4
- Kuantisasi Hanya Bobot Int4 A16W4 + Sparsity 2:4
- Autoquantization
torchao juga mendukung konfigurasi pada level modul dengan menentukan kamus dari nama lengkap modul dan konfigurasi kuantisasi yang bersesuaian. Hal ini memungkinkan untuk melewati kuantisasi pada lapisan tertentu dan menggunakan konfigurasi kuantisasi yang berbeda untuk modul yang berbeda.
Periksa tabel di bawah untuk melihat apakah perangkat keras Anda kompatibel.
| Komponen | Kompatibilitas |
|---|---|
| Versi CUDA | ✅ cu118, cu126, cu128 |
| Versi XPU | ✅ pytorch2.8 |
| CPU | ✅ ubah device_map="cpu" (lihat contoh di bawah) |
Instal torchao dari PyPi atau indeks PyTorch dengan perintah berikut.
PyPi
Indeks PyTorch
# Updating 🤗 Transformers to the latest version, as the example script below uses the new auto compilation
# Stable release from Pypi which will default to CUDA 12.6
pip install --upgrade torchao transformersRilis Stabil dari indeks PyTorch
pip install torchao --index-url https://download.pytorch.org/whl/cu126 # options are cpu/cu118/cu126/cu128torchao >= 0.15.0 diperlukan. API berbasis string (misalnya, TorchAoConfig("int4_weight_only")) telah dihapus — gunakan objek AOBaseConfig sebagai gantinya (lihat contoh di bawah).
Contoh kuantisasi
TorchAO menyediakan berbagai konfigurasi kuantisasi. Setiap konfigurasi dapat dikustomisasi lebih lanjut dengan parameter seperti group_size, scheme, dan layout untuk mengoptimalkan arsitektur perangkat keras dan model tertentu.
Untuk daftar lengkap konfigurasi yang tersedia, lihat dokumentasi API kuantisasi.
Anda dapat memilih jenis dan pengaturan kuantisasi secara manual atau memilih jenis kuantisasi secara otomatis.
Buat TorchAoConfig dan tentukan jenis kuantisasi serta group_size dari bobot yang akan dikuantisasi (untuk int8 weight only dan int4 weight only). Atur cache_implementation menjadi "static" untuk secara otomatis torch.compile metode forward.
Kami akan menunjukkan contoh metode kuantisasi yang direkomendasikan berdasarkan perangkat keras, misalnya GPU A100, GPU H100, dan CPU.
torchao secara otomatis mengompilasi model selama inferensi pertama jika kita mengatur
cache_implementation="static". Model akan dikompilasi ulang setiap kali ukuran batch ataumax_new_tokensdiubah. Sampaikandisable_compile=Trueke generate() untuk melakukan kuantisasi tanpa kompilasi.
GPU H100
float8-dynamic-and-weight-only
int4-weight-only
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, Float8WeightOnlyConfig
quant_config = Float8DynamicActivationFloat8WeightConfig()
# or float8 weight only quantization
# quant_config = Float8WeightOnlyConfig()
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="auto",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import GemliteUIntXWeightOnlyConfig
# We integrated with gemlite, which optimizes for batch size N on A100 and H100
quant_config = GemliteUIntXWeightOnlyConfig(group_size=128)
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="auto",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))GPU A100
int8-dynamic-and-weight-only
int4-weight-only
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, Int8WeightOnlyConfig
quant_config = Int8DynamicActivationInt8WeightConfig()
# or int8 weight only quantization
# quant_config = Int8WeightOnlyConfig()
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="auto",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import GemliteUIntXWeightOnlyConfig, Int4WeightOnlyConfig
# For batch size N, we recommend gemlite, which may require autotuning
# default is 4 bit, 8 bit is also supported by passing `bit_width=8`
quant_config = GemliteUIntXWeightOnlyConfig(group_size=128)
# For batch size 1, we also have custom tinygemm kernel that's only optimized for this
# We can set `use_hqq` to `True` for better accuracy
# quant_config = Int4WeightOnlyConfig(group_size=128, use_hqq=True)
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="auto",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))Intel XPU
int8-dynamic-and-weight-only
int4-weight-only
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, Int8WeightOnlyConfig
quant_config = Int8DynamicActivationInt8WeightConfig()
# or int8 weight only quantization
# quant_config = Int8WeightOnlyConfig()
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="auto",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int4WeightOnlyConfig
quant_config = Int4WeightOnlyConfig(group_size=128, int4_packing_format="plain_int32")
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="auto",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device).to(quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))CPU
int8-dynamic-and-weight-only
int4-weight-only
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, Int8WeightOnlyConfig
quant_config = Int8DynamicActivationInt8WeightConfig()
# quant_config = Int8WeightOnlyConfig()
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="cpu",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))Kuantisasi bobot-only int4 pada CPU memerlukan versi torchao 0.15.0 atau lebih baru.
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.prototype.quantization.int4 import PrototypeInt4WeightOnlyConfig
quantization_config = TorchAoConfig(PrototypeInt4WeightOnlyConfig(group_size=128, int4_choose_qparams_algorithm="tinygemm"))
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="cpu",
quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device).to(quantized_model.dtype)
# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))Kuantisasi Per Modul
1. Lewati kuantisasi untuk lapisan tertentu
Dengan FqnToConfig, kita dapat menentukan konfigurasi default untuk semua lapisan sambil melewati kuantisasi untuk lapisan tertentu.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
model_id = "meta-llama/Llama-3.1-8B-Instruct"
from torchao.quantization import Int4WeightOnlyConfig, FqnToConfig
config = Int4WeightOnlyConfig(group_size=128)
# set default to int4 (for linears), and skip quantizing `model.layers.0.self_attn.q_proj`
quant_config = FqnToConfig({"_default": config, "model.layers.0.self_attn.q_proj": None})
quantization_config = TorchAoConfig(quant_type=quant_config)
quantized_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", dtype=torch.bfloat16, quantization_config=quantization_config)
# lm_head is not quantized and model.layers.0.self_attn.q_proj is not quantized
print("quantized model:", quantized_model)
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Manual Testing
prompt = "Hey, are you conscious? Can you talk to me?"
inputs = tokenizer(prompt, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128)
output_text = tokenizer.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)2. Mengkuantisasi lapisan yang berbeda dengan konfigurasi kuantisasi yang berbeda (tanpa regex)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
model_id = "facebook/opt-125m"
from torchao.quantization import Int4WeightOnlyConfig, FqnToConfig, Int8DynamicActivationInt4WeightConfig, IntxWeightOnlyConfig, PerAxis, MappingType
weight_dtype = torch.int8
granularity = PerAxis(0)
mapping_type = MappingType.ASYMMETRIC
embedding_config = IntxWeightOnlyConfig(
weight_dtype=weight_dtype,
granularity=granularity,
mapping_type=mapping_type,
)
linear_config = Int8DynamicActivationInt4WeightConfig(group_size=128)
quant_config = FqnToConfig({"_default": linear_config, "model.decoder.embed_tokens": embedding_config, "model.decoder.embed_positions": None})
# set `include_embedding` to True in order to include embedding in quantization
# when `include_embedding` is True, we'll remove input embedding from `modules_not_to_convert` as well
quantization_config = TorchAoConfig(quant_type=quant_config, include_embedding=True)
quantized_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", dtype=torch.bfloat16, quantization_config=quantization_config)
print("quantized model:", quantized_model)
# make sure embedding is quantized
print("embed_tokens weight:", quantized_model.model.decoder.embed_tokens.weight)
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Manual Testing
prompt = "Hey, are you conscious? Can you talk to me?"
inputs = tokenizer(prompt, return_tensors="pt").to("cpu", quantized_model.dtype)
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, cache_implementation="static")
output_text = tokenizer.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)3. Mengkuantisasi lapisan yang berbeda dengan konfigurasi kuantisasi yang berbeda (dengan regex)
Kita juga dapat menggunakan regex untuk menentukan konfigurasi untuk semua modul yang memiliki module_fqn yang cocok dengan regex tersebut. Semua regex harus diawali dengan re:, misalnya re:layers\..*\.gate_proj akan cocok dengan semua layer seperti layers.0.gate_proj. Lihat di sini untuk dokumentasinya.
import logging
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
# Configure logging to see warnings and debug information
logging.basicConfig(
level=logging.INFO, format="%(name)s - %(levelname)s - %(message)s"
)
# Enable specific loggers that might contain the serialization warnings
logging.getLogger("transformers").setLevel(logging.INFO)
logging.getLogger("torchao").setLevel(logging.INFO)
logging.getLogger("safetensors").setLevel(logging.INFO)
logging.getLogger("huggingface_hub").setLevel(logging.INFO)
model_id = "facebook/opt-125m"
from torchao.quantization import (
Float8DynamicActivationFloat8WeightConfig,
Int4WeightOnlyConfig,
IntxWeightOnlyConfig,
PerRow,
PerAxis,
FqnToConfig,
Float8Tensor,
Int4TilePackedTo4dTensor,
IntxUnpackedToInt8Tensor,
)
float8dyn = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
int4wo = Int4WeightOnlyConfig(int4_packing_format="tile_packed_to_4d")
intxwo = IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0))
qconfig_dict = {
# highest priority
"model.decoder.layers.3.self_attn.q_proj": int4wo,
"model.decoder.layers.3.self_attn.k_proj": int4wo,
"model.decoder.layers.3.self_attn.v_proj": int4wo,
# vllm
"model.decoder.layers.3.self_attn.qkv_proj": int4wo,
"re:model\.decoder\.layers\..+\.self_attn\.q_proj": float8dyn,
"re:model\.decoder\.layers\..+\.self_attn\.k_proj": float8dyn,
"re:model\.decoder\.layers\..+\.self_attn\.v_proj": float8dyn,
# this should not take effect and we'll fallback to _default
# since no full mach (missing `j` in the end)
"re:model\.decoder\.layers\..+\.self_attn\.out_pro": float8dyn,
# vllm
"re:model\.decoder\.layers\..+\.self_attn\.qkv_proj": float8dyn,
"_default": intxwo,
}
quant_config = FqnToConfig(qconfig_dict)
quantization_config = TorchAoConfig(quant_type=quant_config)
quantized_model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto",
torch_dtype=torch.bfloat16,
quantization_config=quantization_config,
)
print("quantized model:", quantized_model)
tokenizer = AutoTokenizer.from_pretrained(model_id)
for i in range(12):
if i == 3:
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Int4TilePackedTo4dTensor)
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Int4TilePackedTo4dTensor)
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Int4TilePackedTo4dTensor)
else:
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Float8Tensor)
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
assert isinstance(quantized_model.model.decoder.layers[i].self_attn.out_proj.weight, IntxUnpackedToInt8Tensor)
# Manual Testing
prompt = "What are we having for dinner?"
print("Prompt:", prompt)
inputs = tokenizer(
prompt,
return_tensors="pt",
).to(quantized_model.device, quantized_model.dtype)
# setting temperature to 0 to make sure result deterministic
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, temperature=0)
correct_output_text = tokenizer.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print("Response:", correct_output_text[0][len(prompt) :])
# Load model from saved checkpoint
reloaded_model = AutoModelForCausalLM.from_pretrained(
save_to,
device_map=torch.accelerator.current_accelerator(),
torch_dtype=torch.bfloat16,
# quantization_config=quantization_config,
)
generated_ids = reloaded_model.generate(**inputs, max_new_tokens=128, temperature=0)
output_text = tokenizer.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print("Response:", output_text[0][len(prompt) :])
assert(correct_output_text == output_text)Serialisasi
Menyimpan model terkuantisasi dengan save_pretrained (dalam format safetensors) hanya didukung untuk torchao >= v0.15. Untuk versi di bawahnya, yang mungkin dilakukan hanyalah menyimpan secara manual sebagai checkpoint .bin yang tidak aman menggunakan torch.save.
save-locally
push-to-huggingface-hub
# torchao >= 0.15
output_dir = "llama3-8b-int4wo-128"
quantized_model.save_pretrained("llama3-8b-int4wo-128")# torchao >= 0.15
USER_ID = "your_huggingface_user_id"
REPO_ID = "llama3-8b-int4wo-128"
quantized_model.push_to_hub(f"{USER_ID}/llama3-8b-int4wo-128")
tokenizer.push_to_hub(f"{USER_ID}/llama3-8b-int4wo-128")# torchao < 0.15 -> unsafe serialization
filename = "llama3-8b-int4wo-128/pytorch_model.bin"
torch.save(quantized_model.state_dict(), filename)Memuat model terkuantisasi
Memuat model terkuantisasi bergantung pada skema kuantisasi. Untuk skema kuantisasi seperti int8 dan float8, Anda dapat mengkuantisasi model di perangkat apa pun dan juga memuatnya di perangkat apa pun. Contoh di bawah ini menunjukkan cara mengkuantisasi model di CPU dan kemudian memuatnya di CUDA atau XPU.
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8WeightOnlyConfig
quant_config = Int8WeightOnlyConfig(group_size=128)
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="cpu",
quantization_config=quantization_config
)
# save the quantized model
output_dir = "llama-3.1-8b-torchao-int8"
quantized_model.save_pretrained(output_dir)
# reload the quantized model
reloaded_model = AutoModelForCausalLM.from_pretrained(
output_dir,
device_map="auto",
dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(reloaded_model.device.type)
output = reloaded_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))
Untuk int4, model hanya dapat dimuat pada perangkat yang sama dengan perangkat tempat model tersebut dikuantisasi karena layout-nya spesifik terhadap perangkat. Contoh di bawah ini mendemonstrasikan kuantisasi dan pemuatan model pada CPU.
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int4WeightOnlyConfig
from torchao.quantization.quantize_.workflows import Int4PackingFormat
quant_config = Int4WeightOnlyConfig(group_size=128, int4_packing_format=Int4PackingFormat.PLAIN_INT32)
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="cpu",
quantization_config=quantization_config
)
# save the quantized model
output_dir = "llama-3.1-8b-torchao-int4-cpu"
quantized_model.save_pretrained(output_dir)
# reload the quantized model
reloaded_model = AutoModelForCausalLM.from_pretrained(
output_dir,
device_map="cpu",
dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(reloaded_model.device.type)
output = reloaded_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))
Sumber daya
Untuk gambaran yang lebih baik tentang kinerja yang diharapkan, lihat benchmark untuk berbagai model dengan backend CUDA dan XPU. Anda juga dapat menjalankan kode di bawah ini untuk melakukan benchmark model secara mandiri.
from torch._inductor.utils import do_bench_using_profiling
from typing import Callable
def benchmark_fn(func: Callable, *args, **kwargs) -> float:
"""Thin wrapper around do_bench_using_profiling"""
no_args = lambda: func(*args, **kwargs)
time = do_bench_using_profiling(no_args)
return time * 1e3
MAX_NEW_TOKENS = 1000
print("int4wo-128 model:", benchmark_fn(quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static"))
bf16_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", dtype=torch.bfloat16)
output = bf16_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static") # auto-compile
print("bf16 model:", benchmark_fn(bf16_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static"))Untuk performa terbaik, Anda dapat menggunakan pengaturan yang direkomendasikan dengan memanggil
torchao.quantization.utils.recommended_inductor_config_setter()
Lihat Teknik Kuantisasi Lain yang Tersedia untuk contoh dan dokumentasi lebih lanjut.
Isu
Jika Anda mengalami masalah dengan integrasi Transformers, silakan buka issue di repositori Transformers. Untuk masalah yang langsung terkait dengan torchao, silakan buka issue di repositori torchao.
Update on GitHub