[llm-compressor] Quantization Base: QuantizationModifier와 QuantizationMixin
들어가며
QuantizationModifier는 llm-compressor에서 가장 많이 쓰이는 Modifier다. GPTQ·AWQ 같은 특수 알고리즘을 쓰지 않는 "단순한" PTQ(FP8, W8A8, W4A16 등)의 핵심 엔진이며, 다른 알고리즘들도 내부적으로 이 Modifier의 QuantizationMixin을 활용한다. src/llmcompressor/modifiers/quantization/quantization/base.py의 코드를 해부한다.
공식 문서
핵심 구조/코드 분석
QuantizationModifier 시그니처와 파라미터
class QuantizationModifier(Modifier, QuantizationMixin):
"""
Enables PTQ and QAT for a given module or its submodules.
:param config_groups: dict of quantization schemes to apply to target modules
:param targets: list of layer names to quantize (default: Linear)
:param ignore: list of module class names or submodule names to skip
:param scheme: a single quantization scheme or preset name
:param kv_cache_scheme: optional kv-cache quantization args
"""
QuantizationModifier는 두 베이스를 상속한다. Modifier는 라이프사이클 훅(on_initialize, on_finalize, on_event 등)을 제공하고, QuantizationMixin은 실제 observer 등록·calibration 로직을 제공한다. 다중 상속은 역할 분리를 위해 사용된다. 라이프사이클과 양자화 메커니즘이 서로 독립적으로 진화할 수 있다.
주요 파라미터:
| 파라미터 | 의미 |
|---|---|
config_groups |
Modifier 하나가 여러 그룹을 가질 수 있음. 각 그룹은 (targets, scheme) 쌍 |
targets |
양자화 대상 레이어 타입. 정규식 re:.*proj 같은 형태 가능 |
ignore |
대상 중 제외할 모듈. ['lm_head']처럼 명시적 제외 |
scheme |
프리셋 이름(예: "FP8", "W4A16")이나 세부 dict |
kv_cache_scheme |
KV 캐시 양자화용 별도 설정. q_proj/k_proj 출력에 관측자 등록 |
on_initialize: 스킴 부착과 observer 등록
def on_initialize(self, state: State, **kwargs) -> bool:
if not QuantizationMixin.has_config(self):
raise ValueError(
"QuantizationModifier requires that quantization fields be specified"
)
QuantizationMixin.initialize_quantization(self, state.model)
return True
has_config는 scheme/config_groups 중 하나라도 있는지 확인한다. 둘 다 None이면 에러. initialize_quantization은 Mixin에서 제공하는 메서드로, 다음 일을 한다.
- 대상 모듈에
QuantizationScheme을 attach (compressed-tensors의 기본 API 사용) - 각
Linear의 forward 호출을 fake_quantize 래퍼로 덮어씀 - 가중치/활성화 observer 인스턴스 생성 (현재는 비활성 상태)
이 시점에는 observer가 등록되었지만 "캘리브레이션 모드"가 아니다. on_start가 호출되면 본격 작동한다.
on_start: 가중치 캘리브레이션
def on_start(self, state: State, event: Event, **kwargs):
self.started_ = True
QuantizationMixin.start_calibration(self, state.model)
named_modules = list(
match_named_modules(state.model, self.resolved_targets, self.ignore)
)
# 1) global scale 먼저 결정 (마이크로스케일 스킴용)
for _, module in named_modules:
update_weight_global_scale(module)
# 2) fused layer 처리 — q/k/v 같은 융합 가중치를 동기화
for module in state.model.modules():
update_fused_layer_weight_global_scales(module)
# 3) 실제 가중치 스케일/zero point 계산
for _, module in tqdm.tqdm(named_modules, desc="Calibrating weights"):
update_weight_zp_scale(module)
가중치 캘리브레이션은 한 번만 일어난다. 가중치는 시간에 따라 변하지 않으므로, on_start에서 모든 타겟 모듈을 돌며 update_weight_zp_scale을 호출해 그 자리에서 스케일을 확정한다. 활성화 observer는 아직 데이터를 못 봤으므로 on_event에서 점진적으로 수렴한다.
update_fused_layer_weight_global_scales는 미묘하다. q/k/v 같은 융합 레이어는 LLM 구현에 따라 한 모듈로 합쳐진 경우(single-matrix projection) vs. 분리된 경우(three projections)가 있다. 둘 다 같은 global scale을 써야 vLLM/SGLang 추론이 잘 작동하므로, 이 함수가 전체 모듈을 돌며 "융합 가능한 쌍"을 찾아 스케일을 동기화한다. 실행이 idempotent이므로 모든 모듈에 대해 안전하게 호출할 수 있다.
on_event: 활성화 observer 업데이트
def on_event(self, state: State, event: Event, **kwargs):
if event.type_ == EventType.CALIBRATION_EPOCH_START:
if not self.started_:
self.on_start(state, None)
if event.type_ == EventType.SEQUENTIAL_EPOCH_END:
QuantizationMixin.sync_activation_observers(self, state.model)
if event.type_ == EventType.CALIBRATION_EPOCH_END:
QuantizationMixin.sync_activation_observers(self, state.model)
if not self.ended_:
self.on_end(state, None)
세 이벤트를 처리한다.
- CALIBRATION_EPOCH_START: 아직
on_start가 불리지 않았다면 지금 호출. basic 파이프라인에서는 라이프사이클이on_start를 호출하지 않으므로 여기서 fallback. - SEQUENTIAL_EPOCH_END: Sequential Pipeline에서 한 서브그래프가 끝날 때. 이 시점에 DDP 환경이면 모든 rank의 activation 통계를 all-reduce해 동기화한다. 이는 rank마다 다른 배치를 처리하는 상황에서도 최종 스케일이 모든 rank에서 동일하게 되도록 보장한다.
- CALIBRATION_EPOCH_END: 전체 캘리브레이션이 끝나면 observer 동기화 후
on_end호출.
on_end / on_finalize: 정리
def on_end(self, state: State, event: Event, **kwargs):
"""Finish calibrating by removing observers and calibration hooks"""
self.ended_ = True
QuantizationMixin.end_calibration(self, state.model)
def on_finalize(self, state: State, **kwargs) -> bool:
if not self.ended_:
self.on_end(state, None)
end_calibration은 observer를 제거하고 calibration hook을 떼어낸다. 이후 모델은 "quantization 활성화 + observer 없음" 상태로 남아, 저장 시 compressed-tensors 포맷으로 직렬화된다.
on_finalize가 on_end를 호출하는 것은 방어적이다. 어떤 이유로 CALIBRATION_EPOCH_END가 발생하지 않아도 finalize가 걸리면 정리가 끝난다.
DDP 모드의 observer 동기화
QuantizationMixin.sync_activation_observers(self, state.model)
이 호출은 다중 GPU 환경에서만 의미가 있다. 각 rank는 자신의 데이터 샤드를 처리하므로 observer 내부 통계(예: min_val, max_val)가 rank마다 다르다. all-reduce로 "전 rank의 union of min/max"를 계산해 모든 rank에 동일하게 세팅한다. 이 덕분에 DDP 양자화 결과는 단일 GPU로 돌렸을 때와 동일하다.
왜 이 설계인가
1. Modifier + Mixin 분리. 라이프사이클과 양자화 메커니즘을 분리해 각자 독립적으로 진화 가능. 다른 Modifier(GPTQ, AWQ)도 QuantizationMixin만 재사용해 양자화 로직을 공유한다.
2. 가중치 vs 활성화 시점 차이. 가중치는 on_start에서 한 번, 활성화는 on_event에서 점진적으로. 두 종류의 양자화 대상이 시간 스케일이 다르다는 점을 인정하고 분리 처리한다.
3. Fused layer 전역 스케일 동기화. q/k/v 쌍이 다른 모듈로 정의된 경우에도 vLLM/SGLang이 올바르게 작동하려면 동일 스케일이 필요하다. update_fused_layer_weight_global_scales가 이 동기화를 idempotent하게 수행한다.
4. DDP 호환성. sync_activation_observers로 rank 간 통계를 일치시켜, 분산 환경에서도 단일 GPU와 동일한 결과를 보장한다.
5. 방어적 on_finalize. 이벤트가 누락되어도 정리가 보장된다. 파이프라인마다 이벤트 시퀀스가 조금씩 달라도 리소스가 누수되지 않는다.
마무리
Quantization Base는 llm-compressor의 양자화 중심 허브다. 단순 PTQ부터 DDP 시나리오까지 이 Modifier 하나로 커버된다. 다음 글은 여기서 호출되는 Quantization Calibration 헬퍼를 본다.
참고 자료
관련 포스트
llm-compressor 의 다른글
- 이전글 [llm-compressor] iMatrix Observer: 입력 채널 중요도 가중 MSE
- 현재글 : [llm-compressor] Quantization Base: QuantizationModifier와 QuantizationMixin
- 다음글 [llm-compressor] Quantization Calibration: update_weight_zp_scale와 observer 등록
댓글