TorchAO
TorchAO
TorchAO는 PyTorch용 아키텍처 최적화 라이브러리로, 추론과 훈련을 위한 고성능 dtype, 최적화 기법, 커널을 제공하며 torch.compile, FSDP 같은 네이티브 PyTorch 기능과의 컴포저빌리티를 특징으로 해요. 일부 벤치마크 수치는 여기 에서 확인할 수 있습니다.
출처: 문서
본문
최신 torchao nightly를 다음으로 설치하는 것을 권장합니다.
# 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 모델 양자화 (Quantizing HuggingFace Models)
transformers 및 diffusers 같은 torchao로 자신의 huggingface 모델을 양자화하고, 다음 예시 코드로 체크포인트를 허깅페이스 허브 에 저장할 수 있습니다.
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,
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)
또는 간단한 UI로 모델을 양자화하려면 TorchAO Quantization space 를 사용할 수 있습니다.