[llm-compressor] Quantization Calibration: update_weight_zp_scale와 observer 등록

들어가며

Quantization Base의 on_start는 update_weight_zp_scale, update_weight_global_scale 같은 함수들을 호출한다. 이 함수들은 src/llmcompressor/modifiers/quantization/calibration.py에 정의되어 있으며, "모듈 하나의 가중치 또는 활성화 관측자를 호출해 스케일을 결정하고 텐서에 저장"하는 작업을 한다. 이 글은 그 헬퍼 함수들의 역할과, observer/calibration의 연결 방식을 분석한다.

핵심 구조/코드 분석

update_weight_zp_scale: 가중치 스케일 결정

def update_weight_zp_scale(module: torch.nn.Module) -> None:
    """
    Update weight quantization parameters (scale, zero_point) for a module
    using the attached weight observer.
    """
    if not hasattr(module, "weight_observer"):
        return

    observer = module.weight_observer
    weight = module.weight.detach()

    # Observer 의 forward 가 (scale, zero_point) 반환
    scales, zero_points = observer(weight)

    # 계산된 스케일을 모듈에 저장 (compressed-tensors 형식)
    module.weight_scale = scales
    module.weight_zero_point = zero_points

핵심 흐름:

1. 모듈에 등록된 weight_observer 가져오기
2. 가중치 텐서를 observer의 forward에 통과
3. 반환된 (scale, zero_point)를 모듈의 속성으로 저장

weight_scale과 weight_zero_point는 compressed-tensors의 표준 속성 이름이다. 저장 시 이 이름으로 검색되어 체크포인트에 포함된다. vLLM·SGLang이 로딩 시 이 이름을 찾아 양자화 가중치를 복원한다.

update_weight_global_scale: 마이크로스케일 전역 스케일

def update_weight_global_scale(module: torch.nn.Module) -> None:
    """마이크로스케일 스킴에서만 사용되는 전역 스케일 계산"""
    if not hasattr(module, "weight_observer"):
        return

    observer = module.weight_observer
    if not observer.args.requires_global_scale():
        return   # 일반 스킴은 스킵

    weight = module.weight.detach()
    global_scale = observer.get_global_scale(weight)
    module.weight_global_scale = global_scale

NVFP4·MXFP4 같은 마이크로스케일 스킴은 "블록 스케일(E4M3 FP8)"과 "전역 스케일(FP32)"의 두 레벨을 쓴다. 이 함수가 전역 스케일을 계산한다. 일반 스킴(FP8, INT8, W4A16)은 requires_global_scale()이 False라 조기 반환되므로 오버헤드가 없다.

update_activation_zp_scale: 활성화 스케일

가중치와 달리, 활성화는 실행 시간에 결정된다. 캘리브레이션 루프가 observer를 업데이트하고, 최종 시점에 이 함수가 호출되어 값을 확정한다.

def update_activation_zp_scale(module: torch.nn.Module) -> None:
    """
    Finalize activation quantization parameters at end of calibration.
    """
    if not hasattr(module, "input_observer"):
        return

    observer = module.input_observer

    # observer 가 이미 누적한 통계로부터 최종 스케일 계산
    # (기본적으로 0 크기 더미 텐서로 forward 호출해도 누적된 state 로 계산)
    dummy = torch.zeros(1)
    scales, zero_points = observer(dummy)   # moving average 기반이면 누적 상태 사용

    module.input_scale = scales
    module.input_zero_point = zero_points

활성화 observer는 MovingAverageObserverBase 같은 누적 기반이다. 캘리브레이션 루프가 forward를 실행할 때마다 observer.update(...) 같은 내부 훅이 동작해 min/max를 갱신하고, 최종 시점에 forward(dummy)만 호출하면 현재 상태의 스케일을 반환한다.

sync_activation_observers: DDP 동기화

def sync_activation_observers(model: torch.nn.Module) -> None:
    """
    In DDP mode, all-reduce observer stats so all ranks have identical
    quantization parameters.
    """
    if not dist.is_initialized() or dist.get_world_size() == 1:
        return  # 단일 GPU 면 no-op

    for module in model.modules():
        if hasattr(module, "input_observer"):
            obs = module.input_observer
            if hasattr(obs, "min_val") and obs.min_val is not None:
                dist.all_reduce(obs.min_val, op=dist.ReduceOp.MIN)
                dist.all_reduce(obs.max_val, op=dist.ReduceOp.MAX)

단일 GPU 체크가 먼저 온다. DDP가 아니면 즉시 반환해 오버헤드 없음. DDP 모드에서는 각 rank의 observer 내부 텐서를 all_reduce로 동기화한다. MIN 연산은 모든 rank의 min_val 중 가장 작은 값을, MAX는 가장 큰 값을 선택한다. 결과적으로 모든 rank가 "전체 데이터의 극단값"을 공유하게 된다.

initialize_quantization과 start_calibration

이 함수들은 QuantizationMixin에 정의되어 있고, calibration.py의 헬퍼를 사용한다. 흐름은 다음과 같다.

QuantizationMixin.initialize_quantization(modifier, model):
  └─ 각 대상 모듈에 대해:
       ├─ QuantizationScheme 부착 (compressed-tensors API)
       ├─ Observer 인스턴스 생성 (min_max / mse / imatrix)
       ├─ module.weight_observer, module.input_observer 로 저장
       └─ module.forward 를 fake_quantize 래퍼로 덮어씀

QuantizationMixin.start_calibration(modifier, model):
  └─ calibration hook 활성화
       └─ forward 시 activation observer 가 실시간 업데이트되도록

Observer 등록 패턴

각 모듈은 여러 observer를 가질 수 있다.

# 가중치 관측자
module.weight_observer = Observer.load_from_registry(
    "minmax",
    base_name="weight",
    args=scheme.weights,
    module=module,
)

# 입력 관측자 (activation 양자화 시)
module.input_observer = Observer.load_from_registry(
    "minmax",
    base_name="input",
    args=scheme.input_activations,
    module=module,
)

base_name이 "weight"인지 "input"인지로 observer의 역할이 구분된다. 각 observer는 같은 클래스라도 다른 이름을 가지므로 충돌 없이 공존한다.

calibration hook의 구조

활성화 observer는 수동으로 forward마다 호출되지 않는다. 대신 forward 자체를 패치해 자동으로 observer를 업데이트한다.

# 의사 코드
original_forward = module.forward

def wrapped_forward(self, input_):
    if calibration_enabled:
        self.input_observer.update(input_)   # observer 에 통계 누적
    output = original_forward(input_)
    return fake_quantize(output, self.input_scale, ...) if quantize_enabled else output

module.forward = wrapped_forward.__get__(module)

이 패치 덕에 캘리브레이션 루프가 단순히 model(batch)만 호출하면 자동으로 observer가 업데이트된다. 파이프라인 구현체는 이 내부를 몰라도 된다.

왜 이 설계인가

1. 모듈 속성 기반 상태. observer, scale, zero_point 모두 모듈의 속성으로 저장된다. 전역 딕셔너리 같은 외부 상태가 없으므로 모듈을 이동하거나 복제해도 연관 정보가 함께 따라간다.

2. 가중치/활성화 분리 처리. 가중치 스케일은 on_start에서 즉시, 활성화는 캘리브레이션 내내 점진적으로. 두 시간 스케일의 차이를 calibration 함수 집합으로 자연스럽게 표현한다.

3. 마이크로스케일 투명 지원. update_weight_global_scale은 일반 스킴에서 조기 반환해 오버헤드가 없다. 사용자는 스킴만 바꾸면 되고 호출 순서를 건드릴 필요가 없다.

4. DDP 대응 is_initialized 체크. sync_activation_observers가 DDP 여부를 먼저 확인해 단일 GPU 환경에서 no-op이 된다. 같은 코드가 모든 환경에서 동작.

5. forward 패치의 투명성. 파이프라인은 model(batch)만 부르면 되고, 내부적으로 observer 업데이트와 fake_quantize가 자동으로 일어난다. 사용자 코드가 calibration을 신경 쓸 필요가 없다.

마무리

Quantization Calibration은 "observer ↔ 모듈 ↔ 스케일 저장"의 세 점을 잇는 얇은 헬퍼 계층이다. 하지만 이 헬퍼 없이는 Quantization Base의 라이프사이클이 실제로 가중치를 업데이트할 수 없다. 다음 글은 그룹 크기 검증 로직을 본다.

참고 자료

• 소스 코드:
  • src/llmcompressor/modifiers/quantization/calibration.py
  • src/llmcompressor/modifiers/quantization/quantization/mixin.py
• 관련 포스트:
  • Quantization Base
  • Observers Base
  • Group Size Validation
  • GPTQ