torchao

torchao는 사용자 정의 고성능 데이터 타입, 양자화, 희소성을 지원하는 PyTorch 아키텍처 최적화 라이브러리입니다. torch.compile과 같은 PyTorch의 네이티브 기능과 조합하여 더 빠른 추론과 학습을 수행할 수 있습니다.

추가적인 torchao 기능은 아래 표를 참조하세요.

기능	설명
양자화 인식 학습(QAT)	정밀도 손실을 최소화하여 양자화 모델을 학습합니다(QAT README 참고)
Float8 학습	float8 형식을 활용한 고처리량 학습(자세한 내용은 torchtitan 및 Accelerate 문서를 참조)
희소성 지원	더 빠른 추론을 위한 반구조화(2:4) 희소성(자세한 내용은 반구조화(2:4) 희소성을 활용한 신경망 학습 가속화 블로그 게시물을 참고하세요)
옵티마이저 양자화	4비트 및 8비트 Adam 변형을 사용하여 옵티마이저 상태 메모리를 줄입니다
KV 캐시 양자화	메모리를 줄여 긴 컨텍스트 추론을 활성화합니다(KV Cache Quantization 참조)
커스텀 커널 지원	`torch.compile`과 호환되는 자체 연산자를 사용하세요
FSDP2	FSDP2와 조합하여 학습

라이브러리에 대한 자세한 내용은 torchao의 README.md를 참고하세요.

torchao는 아래 양자화 기법을 지원합니다.

A16W8 Float8 동적 양자화
A16W8 Float8 WeightOnly 양자화
A8W8 Int8 동적 양자화
A16W8 Int8 가중치 전용 양자화
A16W4 Int4 가중치 전용 양자화
A16W4 Int4 가중치 전용 양자화 + 2:4 희소성
자동 양자화

torchao는 모듈의 완전한 자격 명칭과 해당 양자화 설정을 매핑하는 사전(dict)을 지정하여 모듈 수준 구성을 지원합니다. 이를 통해 특정 레이어의 양자화를 건너뛰거나, 모듈별로 다른 양자화 설정을 사용할 수 있습니다.

하드웨어가 호환되는지 확인하려면 아래 표를 참조하세요.

구성 요소	호환성
CUDA 버전	✅ cu118, cu126, cu128
XPU 버전	✅ pytorch2.8
CPU	✅ `device_map="cpu"`를 변경합니다(아래 예시 참조)

다음 명령어로 PyPi 또는 PyTorch 인덱스에서 torchao를 설치합니다.

PyPi

PyTorch 인덱스

torchao >= 0.15.0이 필요합니다. 문자열 기반 API(예: TorchAoConfig("int4_weight_only"))는 제거되었습니다. 대신 AOBaseConfig 객체를 사용하세요(아래 예시 참조).

양자화 예제

TorchAO는 다양한 양자화 구성을 제공합니다. 각 구성은 group_size, scheme, layout 등의 매개변수를 통해 특정 하드웨어 및 모델 아키텍처에 맞게 최적화할 수 있습니다.

사용 가능한 구성의 전체 목록은 quantization API 문서를 참조하세요.

양자화 유형 및 설정을 수동으로 선택하거나, 양자화 유형을 자동으로 선택할 수 있습니다.

TorchAoConfig을 생성하고, 양자화할 가중치의 양자화 유형과 group_size를 지정합니다(int8 weight only 및 int4 weight only의 경우). cache_implementation을 "static"으로 설정하여 forward 메서드를 자동으로 torch.compile합니다.

하드웨어(예: A100 GPU, H100 GPU, CPU)에 따라 권장되는 양자화 방법의 예시를 보여 드리겠습니다.

cache_implementation="static"를 설정하면, generate()에서 disable_compile=True를 전달하지 않는 한 torchao는 첫 번째 추론 시 모델을 자동으로 컴파일합니다. 배치 크기나 max_new_tokens가 변경될 때마다 모델은 다시 컴파일됩니다.

H100 GPU

float8-dynamic-and-weight-only

int4-weight-only

A100 GPU

int8-dynamic-and-weight-only

int4-weight-only

Intel XPU

int8-dynamic-and-weight-only

int4-weight-only

CPU

int8-dynamic-and-weight-only

int4-weight-only

모듈별 양자화

1. 특정 레이어에 대한 양자화 건너뛰기

FqnToConfig를 사용하면 모든 레이어에 대한 기본 구성을 지정하면서 특정 레이어의 양자화를 건너뛸 수 있습니다.

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. 정규식 없이 다른 양자화 구성으로 서로 다른 레이어를 양자화

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. 정규 표현식을 사용하여 서로 다른 양자화 구성으로 서로 다른 레이어를 양자화

module_fqn이 정규식과 일치하는 모든 모듈의 config를 지정하기 위해 정규식을 사용할 수도 있습니다. 모든 정규식은 re:로 시작해야 하며, 예를 들어 re:layers\..*\.gate_proj는 layers.0.gate_proj와 같은 모든 layer와 일치합니다. 관련 문서는 여기에서 확인하세요.

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)

직렬화

save_pretrained를 사용하여 양자화된 모델을 safetensors 형식으로 저장하는 기능은 torchao >= v0.15에서만 지원됩니다. 그보다 낮은 버전에서는 torch.save를 사용하여 .bin 체크포인트를 unsafe 방식으로 수동으로 저장하는 것만 가능합니다.

로컬에 저장

push-to-huggingface-hub

양자화 모델 로드

quantized model을 로드하는 방법은 양자화 스킴에 따라 달라집니다. int8과 float8과 같은 양자화 스킴의 경우, 어떤 장치에서든 모델을 양자화하고 어떤 장치에서든 로드할 수 있습니다. 아래 예제는 CPU에서 모델을 양자화한 후 CUDA 또는 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))

int4의 경우, 레이아웃이 장치에 특화되어 있어 모델은 양자화된 것과 동일한 장치에서만 로드할 수 있습니다. 아래 예제는 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))

리소스

예상 성능을 더 잘 파악하려면, CUDA와 XPU 백엔드를 사용한 다양한 모델의 벤치마크를 확인해 보세요. 또한 아래 코드를 실행하여 직접 모델을 벤치마크할 수도 있습니다.

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"))

최적의 성능을 위해 torchao.quantization.utils.recommended_inductor_config_setter()를 호출하여 권장 설정을 사용할 수 있습니다.

추가 예제 및 문서에 대해서는 Other Available Quantization Techniques를 참조하세요.

이슈

Transformers 통합과 관련하여 문제가 발생하면 Transformers 저장소에서 이슈를 등록해 주세요. torchao와 직접적으로 관련된 문제는 torchao 저장소에서 이슈를 등록해 주세요.

Update on GitHub