torchao 양자화 API — quantize_와 워크플로우 컨피그

torchao 양자화 API

torchao의 양자화는 큰 함수 하나와 여러 컨피그로 나뉘어요. 흐름은 '원하는 양자화 방식을 컨피그로 고르고, quantize_에 넘긴다'는 것 하나로 기억하면 돼요.

메인 양자화 API

  • quantize_ — 모델의 linear 모듈 가중치를 컨피그에 따라 변환. 모델을 in-place로 수정해요.
  • FqnToConfig — 모듈/파라미터의 FQN(fully qualified name)에 따라 다른 컨피그를 적용하는 설정 클래스.

워크플로우 컨피그

비트 수와 구조에 따라 컨피그가 나뉘어요. 주요 예시예요.

  • float8: Float8DynamicActivationFloat8WeightConfig, Float8WeightOnlyConfig
  • int8: Int8DynamicActivationInt8WeightConfig, Int8WeightOnlyConfig
  • int4: Int4WeightOnlyConfig — 지금은 groupwise 양자화만 지원 (version 1/2)
  • intx: IntxWeightOnlyConfigtorch.intx(1≤x≤8) 가중치, PerGroup/PerAxis(0) 그레이놀러리티, ASYMMETRIC/SYMMETRIC 매핑 등 지정

prototype 단계인 mx(MXDynamicActivationMXWeightConfig), NVFP4(NVFP4DynamicActivationNVFP4WeightConfig), uintx(UIntxWeightOnlyConfig, gemlite 사용)도 있어요.

from torchao.quantization import Int8WeightOnlyConfig, quantize_

# 가중치만 int8로 양자화 (per-channel symmetric)
quantize_(model, Int8WeightOnlyConfig())

컨피그 이름만 봐도 어느 부분을 몇 비트로 줄이는지 알 수 있어요. 예를 들어 Float8DynamicActivationFloat8WeightConfig활성화와 가중치를 모두 float8 동적 대칭 양자화한다는 뜻이에요.

더 알아보기