[llm-compressor] iMatrix Transform: 중요도 행렬 기반 가중치 리스케일

들어가며

iMatrix Observer 글에서 본 것처럼, llm-compressor는 llama.cpp의 importance matrix 개념을 도입했다. Observer는 "양자화 스케일 결정 시" iMatrix를 쓰는데, Transform 계층의 IMatrixGatherer Modifier는 그보다 앞선 단계다. 가중치를 iMatrix에 맞춰 리스케일해 양자화 친화적 분포를 만든다. src/llmcompressor/modifiers/transform/imatrix/base.py를 분석한다.

핵심 구조/코드 분석

IMatrixGatherer Modifier

class IMatrixGatherer(Modifier):
    """
    Collect per-input-channel importance (E[x²]) and optionally rescale weights.

    This modifier runs in a calibration pass to gather activation statistics,
    then either:
    1. Attaches them to modules as `_imatrix_importance` for later use by
       IMatrixMSEObserver
    2. Rescales weights directly based on importance
    """
    targets: list[str] = field(default_factory=lambda: ["Linear"])
    ignore: list[str] = field(default_factory=list)
    rescale_weights: bool = False       # True 면 가중치를 직접 리스케일
    alpha: float = 0.5                  # 리스케일 강도 (rescale_weights=True 시)

파라미터	기본값	의미
targets	["Linear"]	대상 레이어
rescale_weights	False	True면 직접 가중치 변환, False면 observer에 전달만
alpha	0.5	리스케일 강도 (SmoothQuant의 $\alpha$와 유사)

두 가지 동작 모드

IMatrixGatherer는 두 모드로 동작한다.

모드 1: Gatherer only (rescale_weights=False, 기본)

• 캘리브레이션 데이터를 순회하며 각 모듈에 _imatrix_importance 속성을 저장
• 실제 가중치는 건드리지 않음
• 이후 IMatrixMSEObserver가 이 importance를 사용해 양자화 스케일 결정

모드 2: Rescale (rescale_weights=True)

• importance를 수집한 뒤, SmoothQuant 공식으로 가중치 리스케일
• 이후 양자화는 일반 QuantizationModifier에 위임

on_start: Observer를 통한 통계 수집

def on_start(self, state: State, event: Event, **kwargs):
    # [iMatrix Observer](/opensource/llm-compressor/imatrix-observer/) 의 attach 메서드 재사용
    from llmcompressor.observers.imatrix import IMatrixMSEObserver

    self._observers = {}

    for name, module in match_named_modules(state.model, self.targets, self.ignore):
        if not isinstance(module, torch.nn.Linear):
            continue

        # Observer 인스턴스 생성 (dummy args)
        observer = IMatrixMSEObserver(
            base_name="weight",
            args=QuantizationArgs(num_bits=4, symmetric=True),
            module=module,
        )

        # attach → forward pre-hook 등록
        observer.attach(module)
        self._observers[module] = observer

attach 메서드는 이미 iMatrix Observer 글에서 설명했다. 모듈에 forward pre-hook을 걸어 입력의 제곱합을 누적한다.

on_finalize: 중요도 확정과 선택적 리스케일

def on_finalize(self, state: State, **kwargs) -> bool:
    for module, observer in self._observers.items():
        # detach → _imatrix_importance 계산 후 모듈에 저장
        observer.detach(module)

        if self.rescale_weights and hasattr(module, "_imatrix_importance"):
            self._rescale_module_weights(module)

    return True


def _rescale_module_weights(self, module):
    """
    Rescale weights using importance (SmoothQuant-like).
    W ← W * s_c  where s_c = imp^alpha / max(|W_col|)^(1-alpha)
    """
    W = module.weight.data                # (out, in)
    imp = module._imatrix_importance      # (in,)  —  E[x²] per input channel

    # 컬럼별 최대 절댓값
    w_max = W.abs().amax(dim=0)            # (in,)

    # 스케일 계산
    s = imp.pow(self.alpha) / w_max.pow(1 - self.alpha).clamp(min=1e-8)
    s = s / s.mean()                        # 정규화로 전역 스케일 유지

    # 적용: W ← W * diag(s)
    W_new = W * s.unsqueeze(0)
    module.weight.data = W_new

    # 이전 레이어에 1/s 를 흡수해야 하지만 여기서는 생략 (양자화 후처리에서 처리)
    # 또는 Sequential Pipeline 이 다음 레이어 입력을 1/s 배로 사용하도록 수정

    # 정리
    del module._imatrix_importance

_rescale_module_weights는 SmoothQuant와 유사한 공식을 쓰지만, 활성화 통계 대신 **$E[x^2]$ (MSE-friendly importance)**를 쓴다.

두 모드의 사용 사례

모드 1 (gather only): llm-compressor의 표준 흐름

imatrix_stage:
  transform_modifiers:
    IMatrixGatherer:
      targets: [Linear]

quant_stage:
  quantization_modifiers:
    QuantizationModifier:
      scheme: W4A16
      observers:
        weight: imatrix_mse      # IMatrixMSEObserver 사용

이 레시피는 (1) IMatrixGatherer가 데이터로 importance 수집, (2) QuantizationModifier가 IMatrixMSEObserver로 importance를 활용해 스케일 결정.

모드 2 (rescale): 구체적인 가중치 변환

transform_stage:
  transform_modifiers:
    IMatrixGatherer:
      targets: [Linear]
      rescale_weights: true
      alpha: 0.5

quant_stage:
  quantization_modifiers:
    QuantizationModifier:
      scheme: W4A16

가중치 자체를 사전 변환해서 일반 observer로도 좋은 결과를 얻는다.

왜 이 설계인가

1. Observer와 Modifier의 역할 분리 + 재사용. IMatrixGatherer는 IMatrixMSEObserver의 attach/detach를 재사용한다. 코드 중복이 없고, 두 지점에서 같은 수학을 공유한다.

2. Gather only 모드가 기본. 사용자가 실수로 가중치를 리스케일하지 않도록 안전한 기본값을 둔다. 리스케일은 명시적으로 rescale_weights=True를 설정해야 한다.

3. 2 패스 아키텍처. IMatrixGatherer가 1패스에서 통계 수집, QuantizationModifier + IMatrixMSEObserver가 2패스에서 활용. 두 단계 분리로 레시피 작성자가 세밀한 제어를 할 수 있다.

4. SmoothQuant 공식 재사용. 리스케일 공식은 SmoothQuant와 같은 $\alpha$ 파라미터 구조를 쓴다. 사용자는 두 알고리즘을 비슷한 방식으로 다룰 수 있다.

5. _imatrix_importance 정리. on_finalize 끝에서 del module._imatrix_importance로 메타데이터를 정리한다. 체크포인트 저장 시 이 속성이 포함되지 않도록 한다.

마무리

iMatrix Transform은 llama.cpp의 중요도 개념을 Python + PyTorch 세계에 완전히 통합한 결과다. QuIP/SpinQuant처럼 복잡한 회전은 없지만, 단순한 리스케일로도 극단 양자화에 도움이 된다. 이로써 Transform 섹션이 끝났다. 다음부터는 HF Integration 섹션이다.

참고 자료

• 소스 코드:
  • src/llmcompressor/modifiers/transform/imatrix/base.py
  • src/llmcompressor/observers/imatrix.py
• 관련 포스트:
  • Transform Overview
  • iMatrix Observer
  • SmoothQuant
  • QuIP