torchao × Hugging Face — Transformers·Diffusers로 바로 양자화

torchao × Hugging Face 통합

Hugging Face의 transformers·diffusers는 torchao의 양자화 컨피그를 TorchAoConfig로 그대로 받아들여요. 그래서 모델을 불러올 때 quantization_config에 넘기면 로딩과 동시에 양자화가 됩니다.

필수 패키지 설치

pip install git+https://github.com/huggingface/transformers@main
pip install git+https://github.com/huggingface/diffusers@main
pip install torchao torch accelerate

Transformers로 LLM 양자화

Float8DynamicActivationInt4WeightConfig를 Llama-3.2-1B에 적용하는 예시예요.

from transformers import TorchAoConfig, AutoModelForCausalLM
from torchao.quantization import Float8DynamicActivationInt4WeightConfig

quantization_config = TorchAoConfig(
    quant_type=Float8DynamicActivationInt4WeightConfig(group_size=128, use_hqq=True)
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    torch_dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)

Diffusers로 이미지 생성 모델 양자화

from diffusers import FluxPipeline, FluxTransformer2DModel, TorchAoConfig

model_id = "black-forest-labs/Flux.1-Dev"
dtype = torch.bfloat16

quantization_config = TorchAoConfig("int8wo")
transformer = FluxTransformer2DModel.from_pretrained(
    model_id, subfolder="transformer",
    quantization_config=quantization_config, torch_dtype=dtype,
)
pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=dtype)
pipe.to("cuda")

image = pipe("A cat holding a sign that says hello world",
             num_inference_steps=4, guidance_scale=0.0).images[0]
image.save("output.png")

저장과 safetensors 지원

양자화한 모델은 save_pretrained로 저장하고 push_to_hub로 허브에 올릴 수 있어요. safetensors 직렬화는 일부 컨피그(Int4WeightOnlyConfig, Int8WeightOnlyConfig 등)만 지원하므로, 지원 목록에 없으면 safe_serialization=False를 써야 해요.

quantized_model.save_pretrained("llama3-8b-int4wo-128", safe_serialization=False)

참고

  • weight-only 양자화는 가중치만 저비트로 저장하고 연산은 bfloat16 같이 높은 정밀도로 해요 — 메모리는 줄고 활성화 연산 피크는 유지돼요.
  • 동적 활성화 양자화는 활성화를 그때그때 양자화해 추가 메모리를 아끼지만, 품질 트레이드오프가 있을 수 있어서 모델별 테스트가 권장돼요.

더 알아보기