[llm-compressor] Lifecycle: Modifier 초기화-이벤트-종료 상태 머신
들어가며
CompressionSession 글에서 본 것처럼, CompressionSession은 얇은 래퍼이고 실제 상태 머신은 CompressionLifecycle이 들고 있다. Lifecycle은 세 가지 일을 한다. (1) Recipe를 파싱해 Modifier 리스트를 만든다. (2) 각 Modifier의 initialize → event* → finalize 호출 순서를 보장한다. (3) 이벤트 순서가 논리적으로 유효한지 검증한다 (예: BATCH_START 없이 BATCH_END가 오면 거부). 이 글은 src/llmcompressor/core/lifecycle.py를 해부한다.
공식 문서
핵심 구조/코드 분석
CompressionLifecycle dataclass
@dataclass
class CompressionLifecycle:
state: State = field(default_factory=State) # 모델/데이터/하드웨어 상태
recipe: Recipe = field(default_factory=Recipe) # 현재 실행 중인 레시피
initialized_: bool = False # initialize() 호출 완료 플래그
finalized: bool = False # finalize() 호출 완료 플래그
_last_event_type: EventType | None = EventType.BATCH_END # 마지막 이벤트 기억
_event_order: list[EventType] = field(
default_factory=lambda: [
EventType.BATCH_START, # 1) 배치 시작
EventType.LOSS_CALCULATED, # 2) 손실 계산 (훈련 시만)
EventType.OPTIM_PRE_STEP, # 3) 옵티마이저 스텝 전
EventType.OPTIM_POST_STEP, # 4) 옵티마이저 스텝 후
EventType.BATCH_END, # 5) 배치 종료
]
)
global_step: int = 0 # 현재 글로벌 스텝 (epoch * steps_per_epoch + batch)
| 필드 | 의미 |
|---|---|
state |
State 객체. 모델/옵티마이저/데이터/하드웨어 정보를 보관 |
recipe |
파싱된 Recipe 객체. recipe.modifiers가 실제 순회 대상 |
initialized_ |
Modifier들의 on_initialize가 모두 호출되었는지 |
finalized |
Modifier들의 on_finalize가 모두 호출되었는지 |
_last_event_type |
가장 최근에 호출된 이벤트. 순서 검증에 사용 |
_event_order |
허용되는 이벤트 순서. 배치 라이프사이클의 표준 훅 순서 |
global_step |
훈련 루프의 현재 스텝 번호 |
_event_order는 PyTorch 훈련 루프의 표준 순서를 그대로 반영한다. oneshot PTQ에서는 BATCH_START와 BATCH_END만 쓰고 중간 세 개(LOSS/OPTIM_PRE/OPTIM_POST)는 생략되지만, 코드는 훈련 시나리오까지 커버한다.
initialize: Modifier 일괄 준비
def initialize(
self,
recipe: RecipeInput | None = None, # 레시피 입력 (경로/문자열/Recipe/Modifier 리스트)
recipe_stage: RecipeStageInput | None = None, # 실행할 스테이지 이름
recipe_args: RecipeArgsInput | None = None, # 레시피 변수 오버라이드
**kwargs, # 나머지는 State.update 로 전달
) -> list[Any]:
self.state.update(**kwargs) # model/optimizer/data 등 State 에 주입
if self.initialized_: # 중복 initialize 방지
return
# 1) Recipe 객체 생성
if not recipe:
self.recipe = Recipe()
else:
self.recipe = Recipe.create_instance(
path_or_modifiers=recipe, target_stage=recipe_stage
)
if recipe_args:
self.recipe.args = {**recipe_args}
# 2) 각 Modifier 의 on_initialize 훅 호출
mod_data = []
for mod in self.recipe.modifiers:
data = mod.initialize(state=self.state, **kwargs)
if data is not None:
mod_data.append(data)
self.initialized_ = True
return mod_data
initialize의 핵심은 두 단계다. 먼저 입력을 Recipe 객체로 변환하고, 각 Modifier에 대해 mod.initialize(state=...)를 호출한다. 각 Modifier가 초기화 과정에서 반환하는 데이터는 mod_data 리스트에 모인다. 이는 로깅/디버깅 용도로 쓰인다.
중복 방지 체크(if self.initialized_: return)가 중요하다. reset_stage 없이 두 번째 initialize가 호출되면 조용히 무시된다. 이는 멀티스테이지 시나리오에서 실수로 레시피를 두 번 로딩하는 것을 방지한다.
finalize: Modifier 종료 훅
def finalize(self, **kwargs) -> list[Any]:
if not self.initialized_:
raise ValueError("Cannot finalize before initializing")
if self.finalized:
raise ValueError("Cannot finalize more than once")
mod_data = []
for mod in self.recipe.modifiers:
data = mod.finalize(state=self.state, **kwargs)
if data is not None:
mod_data.append(data)
self.finalized = True
return mod_data
이 호출 시점에 GPTQ가 헤시안을 풀고 실제 양자화 가중치를 모델에 주입한다. 순수 단방향 메서드다. initialize가 호출되지 않았거나 이미 finalize된 경우 예외를 던진다. 상태 머신을 엄격하게 지킨다.
event: 이벤트 브로드캐스트와 순서 검증
def event(
self, event_type: EventType, global_step: int | None = 0, **kwargs
) -> list[Any]:
if not self.initialized_:
raise ValueError("Cannot invoke event before initializing")
if self.finalized:
raise ValueError("Cannot invoke event after finalizing")
# INITIALIZE / FINALIZE 는 event() 로 호출 불가 — 전용 메서드를 쓰라는 가드
if event_type in [EventType.INITIALIZE, EventType.FINALIZE]:
raise ValueError(
f"Cannot invoke {event_type} event. Use the corresponding method instead."
)
# 이벤트 순서 검증
if not self._validate_event_order(event_type):
raise ValueError(
f"Lifecycle events must appear following order: {self._event_order}. "
f"Instead, {self._last_event_type} was called before {event_type}"
)
# LOSS_CALCULATED 이벤트는 loss 값이 필수
if event_type == EventType.LOSS_CALCULATED and (
"loss" not in kwargs or kwargs["loss"] is None
):
raise ValueError("Loss must be provided for loss calculated event")
if global_step is not None:
self.global_step = global_step
# 모든 Modifier 에 이벤트 전달
event = Event(type_=event_type)
mod_data = []
for mod in self.recipe.modifiers:
data = mod.update_event(state=self.state, event=event, **kwargs)
if data is not None:
mod_data.append(data)
return mod_data
가드가 다섯 개 있다. 초기화 전 이벤트 금지, 종료 후 이벤트 금지, INITIALIZE/FINALIZE 직접 호출 금지, 순서 위반 금지, 손실 이벤트의 필수 인자 체크. 각 위반은 ValueError로 즉시 실패해 "런타임에 조용히 잘못 동작하는" 시나리오를 원천 차단한다.
_validate_event_order: 순서 검증 알고리즘
def _validate_event_order(self, event_type: EventType) -> bool:
if event_type not in self._event_order:
return True # 정의된 순서에 없는 이벤트는 언제나 유효 (예: CALIBRATION_EPOCH_START)
if event_type == EventType.BATCH_START:
# 새 배치는 이전 배치가 BATCH_END 로 끝난 후에만 시작 가능
valid = self._last_event_type != EventType.BATCH_START
else:
# 나머지는 순서대로만 진행 가능 (역행 금지)
last_event_index = self._event_order.index(self._last_event_type)
curr_event_index = self._event_order.index(event_type)
valid = last_event_index <= curr_event_index
if valid:
self._last_event_type = event_type
return valid
두 가지 규칙이다.
BATCH_START는 직전 이벤트가BATCH_START가 아닐 때만 허용된다. 즉, 한 배치가 시작된 뒤 또 다른BATCH_START가 오면 거부한다.- 나머지 이벤트는 선형 순서를 지켜야 한다.
_event_order리스트에서 현재 이벤트의 인덱스가 마지막 이벤트의 인덱스 이상이어야 한다.
캘리브레이션 관련 이벤트(CALIBRATION_EPOCH_START, SEQUENTIAL_EPOCH_END 등)는 _event_order에 포함되지 않으므로 첫 줄에서 통과된다. 이는 캘리브레이션 훅이 배치 훅과 독립적으로 동작할 수 있음을 의미한다.
reset: 안전한 상태 초기화
def reset(self):
for mod in self.recipe.modifiers:
if not mod.initialized or mod.finalized:
continue
try:
mod.finalize(self.state) # 아직 살아있는 Modifier 는 finalize 시도
except Exception as e:
logger.warning(f"Exception during finalizing modifier: {e}")
self.__init__() # 모든 필드를 기본값으로 재설정
reset은 Lifecycle을 깨끗한 상태로 되돌린다. 단순한 __init__() 호출 전에, 아직 finalize되지 않은 Modifier들을 찾아서 먼저 finalize를 시도한다. 이는 GPU 메모리 누수 방지용이다. GPTQ 같은 Modifier는 캘리브레이션 중 헤시안 행렬을 큰 텐서로 들고 있는데, 이를 finalize 없이 그냥 버리면 GPU 메모리가 해제되지 않을 수 있다. 예외가 발생해도 경고만 내고 계속 진행한다.
왜 이 설계인가
1. dataclass로 직렬화 용이. Lifecycle이 @dataclass인 덕분에 Pydantic 모델로 쉽게 교체할 수 있고, 상태 스냅샷을 파일로 저장·복원하는 기능을 추가하기 쉽다.
2. 엄격한 상태 머신. initialized_, finalized 두 플래그와 _last_event_type으로 상태 전이를 선형 순서로 강제한다. 분산 실행이나 멀티스레드 환경에서 잘못된 순서의 호출이 오면 즉시 실패한다.
3. 순서 검증 예외 처리. _event_order에 없는 이벤트는 언제나 허용한다. 덕분에 새로운 이벤트 타입을 추가해도 기존 검증 로직을 건드릴 필요가 없다. 엄격한 핵심 이벤트와 유연한 확장 이벤트의 공존.
4. reset에서 finalize 시도. 메모리 누수와 GPU 리소스 해제를 위해 reset이 단순 초기화가 아닌 "graceful shutdown"을 수행한다. 예외는 경고로만 남겨 다음 실행을 방해하지 않는다.
5. global_step 필드 유지. 훈련 루프에서는 글로벌 스텝이 핵심 진행률 지표다. oneshot PTQ에서는 거의 쓰이지 않지만, 필드를 유지해 훈련 경로와 API를 통일한다.
마무리
Lifecycle은 Modifier 라이프사이클의 "감독관"이다. 각 Modifier는 자기 일만 신경 쓰면 되고, Lifecycle이 전체 흐름과 순서를 보장한다. 다음 글에서는 Lifecycle이 들고 있는 State 객체를 분석한다.
참고 자료
관련 포스트
llm-compressor 의 다른글
- 이전글 [llm-compressor] CompressionSession: 전역 싱글톤 세션과 Lifecycle 래퍼
- 현재글 : [llm-compressor] Lifecycle: Modifier 초기화-이벤트-종료 상태 머신
- 다음글 [llm-compressor] State & ModelLayer: 압축 상태 저장소
댓글