본문으로 건너뛰기

[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_configscheme/config_groups 중 하나라도 있는지 확인한다. 둘 다 None이면 에러. initialize_quantization은 Mixin에서 제공하는 메서드로, 다음 일을 한다.

  1. 대상 모듈에 QuantizationScheme을 attach (compressed-tensors의 기본 API 사용)
  2. Linear의 forward 호출을 fake_quantize 래퍼로 덮어씀
  3. 가중치/활성화 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)

세 이벤트를 처리한다.

  1. CALIBRATION_EPOCH_START: 아직 on_start가 불리지 않았다면 지금 호출. basic 파이프라인에서는 라이프사이클이 on_start를 호출하지 않으므로 여기서 fallback.
  2. SEQUENTIAL_EPOCH_END: Sequential Pipeline에서 한 서브그래프가 끝날 때. 이 시점에 DDP 환경이면 모든 rank의 activation 통계를 all-reduce해 동기화한다. 이는 rank마다 다른 배치를 처리하는 상황에서도 최종 스케일이 모든 rank에서 동일하게 되도록 보장한다.
  3. 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_finalizeon_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 의 다른글