torchao 소개 — 설치와 int4 양자화 시작

torchao 소개

모델이 커질수록 메모리와 추론 속도가 병목이 돼요. torchao는 이 문제를 가중치를 낮은 비트로 양자화해서 푸는 데 특화된 PyTorch 라이브러리예요. 공식 문서는 'Pre-train Llama-3.1-70B 1.5x faster with float8 training', 'Quantize Llama-3-8B to int4 for 1.89x faster inference with 58% less memory' 같은 수치로 그 효과를 보여줘요.

무엇을 하는 라이브러리인가요

torchao는 커스텀 데이터 타입과 최적화를 위한 라이브러리예요. 가중치, 그래디언트, 옵티마이저, 활성화를 양자화·희소화하면서 네이티브 PyTorch로 처리한다는 점이 특징이에요. 추론뿐 아니라 훈련 단계(float8 훈련, QAT)에도 쓸 수 있어요.

빠른 시작

설치부터 해요.

pip install torchao

모델 가중치를 int4로 양자화하는 기본 예시예요.

import torch
from torchao.quantization import Int4WeightOnlyConfig, quantize_

if torch.cuda.is_available():
    # quantize on CUDA
    quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="tile_packed_to_4d", int4_choose_qparams_algorithm="hqq"))
elif torch.xpu.is_available():
    # quantize on XPU
    quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="plain_int32"))

설치 옵션

환경에 맞춰 설치할 수 있어요.

pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu128  # Nightly
pip install torchao --index-url https://download.pytorch.org/whl/cpu                  # CPU only

개발 빌드는 --no-build-isolation 플래그가 필요해요 (예: USE_CUDA=1 pip install -e . --no-build-isolation). 일부 워크플로우를 가속하는 MSLK는 선택 의존성이고, 안정 버전엔 pip install mslk-cuda==1.0.0을 써요.

더 알아보기