TorchAO

TorchAO 是一個用於 PyTorch 的架構最佳化庫，它提供高效能資料型別、最佳化技術和推理與訓練核心，並支援與 PyTorch 原生功能（如 torch.compile、FSDP 等）的組合使用。一些基準測試資料可以在此處找到。

我們建議使用以下命令安裝最新的 TorchAO 每夜版：

# Install the latest TorchAO nightly build
# Choose the CUDA version that matches your system (cu126, cu128, etc.)
pip install \
    --pre torchao>=10.0.0 \
    --index-url https://download.pytorch.org/whl/nightly/cu126

量化 HuggingFace 模型

您可以使用 TorchAO 量化自己的 Hugging Face 模型，例如 transformers 和 diffusers，並使用以下示例程式碼將檢查點儲存到 Hugging Face Hub，例如這個

程式碼

import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8WeightOnlyConfig

model_name = "meta-llama/Meta-Llama-3-8B"
quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
quantized_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

hub_repo = # YOUR HUB REPO ID
tokenizer.push_to_hub(hub_repo)
quantized_model.push_to_hub(hub_repo, safe_serialization=False)

或者，您可以使用 TorchAO 量化空間透過簡單的使用者介面來量化模型。