[llm-compressor] Modifier Base: 모든 Modifier가 상속하는 기반 클래스

들어가며

llm-compressor의 모든 알고리즘(GPTQ, AWQ, SmoothQuant, SparseGPT, Wanda, QuIP, SpinQuant 등)은 Modifier를 상속한다. 이 베이스 클래스는 "Modifier는 이런 메서드를 구현하면 된다"는 계약을 정의하고, 라이프사이클 관리, 이벤트 디스패치, 상태 플래그 관리의 공통 기반을 제공한다. src/llmcompressor/modifiers/modifier.py를 해부한다.

핵심 구조/코드 분석

Pydantic 모델 + Mixin 상속

class Modifier(ModifierInterface, HooksMixin):
    """
    A base class for all modifiers to inherit from.

    Lifecycle:
    1. initialize
    2. on_event ->
        * on_start if self.start <= event.current_index
        * on_end if self.end >= event.current_index
    5. finalize
    """
    model_config = ConfigDict(extra="forbid")    # Pydantic: 알 수 없는 필드는 에러

    index: int | None = None        # 레시피 안에서의 순서 (디버깅용)
    group: str | None = None        # 그룹 이름 (예: "quantization", "pruning")
    start: float | None = None      # Modifier 작동 시작 인덱스 (에폭 또는 스텝)
    end: float | None = None        # Modifier 작동 종료 인덱스
    update: float | None = None     # Modifier 업데이트 간격

    initialized_: bool = False      # initialize 완료 플래그
    finalized_: bool = False        # finalize 완료 플래그
    started_: bool = False          # on_start 호출 완료 플래그
    ended_: bool = False            # on_end 호출 완료 플래그

Modifier는 두 가지를 상속한다.

• ModifierInterface — 추상 메서드만 선언
• HooksMixin — register_hook() 같은 PyTorch forward hook 관리 메서드 제공

또한 Pydantic BaseModel이기도 해서(부모 체인 어딘가에서) 필드는 Pydantic 필드로 취급된다. ConfigDict(extra="forbid")는 "알 수 없는 인자가 들어오면 조용히 무시하지 말고 에러"다. 이는 레시피 YAML 오타를 빨리 잡게 해준다. block_size를 block_sizes로 잘못 쓰면 즉시 실패한다.

initialize: 상태 전이 + 조건부 start

def initialize(self, state: State, **kwargs):
    if self.initialized_:
        raise RuntimeError("Cannot initialize a modifier that has already been initialized")
    if self.finalized_:
        raise RuntimeError("Cannot initialize a modifier that has already been finalized")

    # 실제 초기화 로직은 자식 클래스의 on_initialize 에 위임
    self.initialized_ = self.on_initialize(state=state, **kwargs)

    # start 조건이 이미 충족된다면 on_start 도 즉시 호출
    fake_start_event = Event(type_=EventType.BATCH_START, global_step=0)
    if self.should_start(fake_start_event):
        self.on_start(state, fake_start_event, **kwargs)
        self.started_ = True

initialize는 두 가지 일을 한다.

1. on_initialize(state)를 호출해 자식 클래스의 실제 초기화 로직을 수행. 반환값은 True면 초기화 성공.
2. "가짜 BATCH_START 이벤트"를 만들어 should_start를 확인한다. start=None이거나 start=0이면 즉시 on_start를 호출한다.

이 두 번째 동작이 중요하다. PTQ에서는 start와 end가 의미 없지만, Modifier 라이프사이클을 훈련용 코드와 통합하기 위해 남아 있다. initialize 직후 바로 "시작된 상태"로 들어가게 해주는 편의 로직이다.

finalize: 상태 전이 + on_finalize 위임

def finalize(self, state: State, **kwargs):
    if self.finalized_:
        raise RuntimeError("cannot finalize a modifier twice")
    if not self.initialized_:
        raise RuntimeError("cannot finalize an uninitialized modifier")

    self.finalized_ = self.on_finalize(state=state, **kwargs)

finalize는 단순하다. 초기화 여부와 중복 호출을 체크한 뒤 자식 클래스의 on_finalize를 호출한다. GPTQ/AWQ/SparseGPT처럼 최종화 시점에 실제 가중치 변경이 일어나는 Modifier는 이 훅에 핵심 로직을 둔다.

update_event: 이벤트 디스패치 + 상태 전이

가장 복잡한 메서드다. 이벤트를 받아 적절한 훅을 호출한다.

def update_event(self, state: State, event: Event, **kwargs):
    if not self.initialized_:
        raise RuntimeError("Cannot update an uninitialized modifier")
    if self.finalized_:
        raise RuntimeError("Cannot update a finalized modifier")

    # 1) 일반 on_event 훅 — 모든 이벤트에 대해 호출
    self.on_event(state, event, **kwargs)

    # 2) BATCH_START + 시작 조건 충족 → on_start 호출 후 on_update
    if (
        event.type_ == EventType.BATCH_START
        and not self.started_
        and self.should_start(event)
    ):
        self.on_start(state, event, **kwargs)
        self.started_ = True
        self.on_update(state, event, **kwargs)
        return

    # 3) BATCH_END + 종료 조건 충족 → on_end 호출 후 on_update
    if (
        event.type_ == EventType.BATCH_END
        and not self.ended_
        and self.should_end(event)
    ):
        self.on_end(state, event, **kwargs)
        self.ended_ = True
        self.on_update(state, event, **kwargs)
        return

    # 4) 그 외: started && !ended 인 "활성" 상태면 on_update 호출
    if self.started_ and not self.ended_:
        self.on_update(state, event, **kwargs)

이벤트 상황	호출되는 훅
모든 이벤트	on_event (필터 없는 일반 훅)
BATCH_START + should_start=True + 아직 시작 안 함	on_start → on_update
BATCH_END + should_end=True + 아직 종료 안 함	on_end → on_update
활성 상태 (started_ && !ended_)	on_update

이 복잡한 디스패치는 훈련 시나리오를 염두에 둔 것이다. PTQ에서는 start=end=0이거나 None이므로 on_event와 on_update만 실질적으로 호출된다.

should_start / should_end 조건

def should_start(self, event: Event) -> bool:
    if self.start is None:
        return False
    current = event.current_index
    return self.start <= current and (self.end is None or current < self.end)

def should_end(self, event: Event):
    current = event.current_index
    return self.end is not None and current >= self.end

start가 None이면 "영원히 시작 안 함"(자동 시작은 initialize에서 처리). end가 None이면 "종료 없음". 이 로직은 간단하지만, Modifier 가 특정 에폭 구간에서만 작동하는 훈련 시나리오를 지원한다.

추상 메서드 계약

@abstractmethod
def on_initialize(self, state: State, **kwargs) -> bool:
    """자식 클래스는 반드시 구현. True 반환 시 초기화 성공"""
    raise NotImplementedError()

def on_finalize(self, state: State, **kwargs) -> bool:
    """기본 구현은 True 반환. 자식이 오버라이드 가능"""
    return True

def on_start(self, state: State, event: Event, **kwargs):
    """기본 구현은 빈 훅. 자식이 오버라이드 가능"""
    pass

def on_update(self, state: State, event: Event, **kwargs):
    """기본 구현은 빈 훅"""
    pass

def on_end(self, state: State, event: Event, **kwargs):
    """기본 구현은 빈 훅"""
    pass

def on_event(self, state: State, event: Event, **kwargs):
    """기본 구현은 빈 훅"""
    pass

on_initialize만 abstract이고 나머지는 빈 기본 구현을 가진다. 즉 "최소한의 Modifier"는 on_initialize만 구현하면 된다. 예를 들어 data_free 파이프라인의 FP8 weight-only Modifier는 on_initialize에서 가중치를 그 자리에서 양자화하고 끝낼 수도 있다.

왜 이 설계인가

1. Template Method 패턴. 베이스 initialize/finalize/update_event는 상태 검사와 디스패치만 하고, 실제 로직은 on_* 훅에 위임한다. 자식 클래스는 필요한 훅만 오버라이드하면 되고, 상태 플래그 관리는 부모가 자동으로 처리한다.

2. 엄격한 상태 검사. 중복 initialize, 초기화 전 update, finalize 후 update 등 잘못된 순서의 호출은 모두 RuntimeError로 즉시 실패한다. 조용한 버그를 방지한다.

3. 훈련과 PTQ 통합. start/end/update 필드와 on_start/on_end/on_update 훅은 훈련 시나리오용이지만, PTQ에서도 그대로 작동한다. 공통 API 덕분에 미래의 QAT 기능 추가가 쉽다.

4. extra="forbid" Pydantic 설정. 레시피 YAML에 오타가 있으면 Modifier 생성 시 에러가 난다. 사용자 실수를 빨리 잡는다.

5. ModifierInterface 분리. 추상 메서드만 모은 별도 인터페이스를 상속하는 이유는 "덕 타이핑 방지"다. 다른 라이브러리가 자체 Modifier 유사 객체를 만들어 플러그인으로 등록하려면 이 인터페이스를 구현해야 한다.

마무리

Modifier Base는 "모든 알고리즘의 공통 골격"이다. 자식 클래스는 on_initialize 하나만 구현해도 되고, 필요하면 on_finalize를 추가해 최종 가중치 변경을 수행할 수 있다. 다음 글은 이 베이스에서 Modifier 인스턴스를 문자열 이름으로 생성하는 Modifier Factory를 본다.

참고 자료

• 소스 코드:
  • src/llmcompressor/modifiers/modifier.py
  • src/llmcompressor/modifiers/interface.py
  • src/llmcompressor/modifiers/utils/hooks.py
• 관련 포스트:
  • Modifier Factory
  • Modifier Interface
  • Events
  • Lifecycle
  • Recipe DSL