[llm-compressor] CompressionSession: 전역 싱글톤 세션과 Lifecycle 래퍼

2026년 4월 13일수정: 2026년 4월 13일

들어가며

oneshot() 함수 내부를 보면 active_session() 호출이 제일 먼저 나온다. 이 함수가 반환하는 CompressionSession이 llm-compressor의 "실행 엔진"이다. 이 세션은 레시피, 모델, 데이터, 옵티마이저, 콜백을 한꺼번에 들고 있으며, Modifier들의 라이프사이클 전체(initialize → event* → finalize)를 관리한다.

흥미로운 점은 이 세션이 전역 싱글톤이라는 것이다. 사용자는 직접 CompressionSession()을 인스턴스화하지 않는다. 대신 active_session()이라는 전역 함수로 접근한다. 이 패턴은 Python 로깅(logging.getLogger())이나 PyTorch의 기본 디바이스 개념과 비슷하다. 이 글은 src/llmcompressor/core/session.py의 CompressionSession 클래스와 src/llmcompressor/core/session_functions.py의 전역 디스패처를 해부한다.

공식 문서

llm-compressor core README (있다면)

핵심 구조/코드 분석

`CompressionSession`: 얇은 래퍼

CompressionSession은 놀라울 정도로 얇다. 실제 상태와 로직은 전부 CompressionLifecycle에 위임된다.

class CompressionSession:
    """A session for compression that holds the lifecycle and state."""

    def __init__(self):
        self._lifecycle = CompressionLifecycle()  # 실제 상태는 Lifecycle 이 보유

    @property
    def lifecycle(self) -> CompressionLifecycle:
        return self._lifecycle

    @property
    def state(self) -> State:
        return self._lifecycle.state   # State 접근도 Lifecycle 경유

이 구조는 "CompressionSession = Lifecycle + 사용자 API"로 요약된다. 왜 래퍼를 따로 둘까? 두 가지 이유다.

사용자 API 안정성. CompressionLifecycle은 내부 구현이 자주 바뀔 수 있다. 사용자가 직접 의존하는 것은 CompressionSession의 공개 메서드로 제한되어 있어야 한다. 래퍼가 그 경계를 명시적으로 만든다.
콜백/이벤트 통합. CompressionSession은 향후 콜백 관리(_CallbackContainer)를 붙이기 위한 자리다. 현재 코드에는 이 구조가 완전히 활용되지 않지만, Lifecycle 상태 변화에 부작용을 붙이는 허브 역할을 한다.

`initialize`: 세션 준비

oneshot()이 호출하는 메서드다. 한 번 호출하면 세션이 실행 준비 상태가 된다.

def initialize(
    self,
    recipe: str | list[str] | Recipe | list[Recipe] | None = None,  # 레시피 경로/객체
    recipe_stage: str | list[str] | None = None,    # 실행할 스테이지 이름
    recipe_args: dict[str, Any] | None = None,      # 레시피 변수 오버라이드
    model: Any | None = None,                       # 압축 대상 모델
    teacher_model: Any | None = None,               # 증류용 교사 모델 (옵션)
    optimizer: Any | None = None,                   # 옵티마이저 (훈련 시)
    attach_optim_callbacks: bool = True,            # 옵티마이저에 lifecycle 콜백 부착
    train_data: Any | None = None,                  # 훈련 데이터 (옵션)
    val_data: Any | None = None,                    # 검증 데이터 (옵션)
    test_data: Any | None = None,                   # 테스트 데이터 (옵션)
    calib_data: Any | None = None,                  # 캘리브레이션 데이터 (oneshot 용)
    copy_data: bool = True,                         # 내부 사용 시 데이터 복사
    start: float | None = None,                     # 시작 에폭 (훈련 루프)
    steps_per_epoch: int | None = None,             # 에폭당 스텝 수
    batches_per_step: int | None = None,            # 스텝당 배치 수
    **kwargs,
) -> ModifiedState:
    mod_data = self._lifecycle.initialize(...)     # 내부는 그대로 Lifecycle 위임

    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,
        modifier_data=mod_data,
    )

인자	사용 용도
`recipe` + `recipe_stage` + `recipe_args`	어떤 레시피를 어떤 스테이지로 어떤 변수로 실행할지
`model`	압축할 `PreTrainedModel` 인스턴스
`teacher_model`	지식 증류 시 필요한 교사 모델 (oneshot PTQ 에서는 보통 None)
`calib_data`	캘리브레이션 DataLoader. basic/sequential 파이프라인이 순회
`train_data`/`val_data`/`test_data`	QAT 또는 Sparse Finetuning 에서 사용
`copy_data`	데이터를 복사해 부작용 방지 (예: DataLoader 상태 공유 방지)

눈여겨볼 점은 oneshot PTQ를 넘어 훈련 루프까지 염두에 둔 시그니처라는 것이다. start, steps_per_epoch, batches_per_step, attach_optim_callbacks는 모두 파인튜닝 시나리오에서 쓰인다. llm-compressor는 과거 SparseML의 계승 프로젝트라, "훈련 + 압축 통합 스케줄러"의 유산이 API에 남아 있다. 현재 사용되는 것은 대부분 recipe, model, calib_data 세 개다.

반환값 ModifiedState는 세션 실행 후 갱신된 모델/옵티마이저/손실의 스냅샷이다. 실제 모델 객체는 인플레이스로 수정되므로 이 반환값은 호출자가 굳이 받지 않아도 된다.

`finalize`: Modifier 종료 훅

def finalize(self, **kwargs) -> ModifiedState:
    mod_data = self._lifecycle.finalize(**kwargs)  # Modifier 각자의 on_finalize 호출
    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,
        modifier_data=mod_data,
    )

이 호출이 llm-compressor의 가장 중요한 순간이다. GPTQ 같은 Modifier는 캘리브레이션 데이터를 다 본 뒤에야 최종 가중치를 계산한다. on_finalize에서 헤시안 역행렬을 풀고, 양자화된 가중치를 모델에 주입하는 등 실제 가중치 변경이 이 시점에 일어난다. on_event는 통계 누적만 하고 on_finalize가 최종화를 수행하는 패턴이다.

`event`: 이벤트 브로드캐스트

파이프라인이 배치를 처리할 때마다 호출되는 훅이다.

def event(
    self,
    event_type: EventType,            # BATCH_START, BATCH_END, OPTIM_POST_STEP 등
    batch_data: Any | None = None,    # 이 이벤트가 관련된 배치 데이터
    loss: Any | None = None,          # 손실값 (훈련 시)
    **kwargs,
) -> ModifiedState:
    mod_data = self._lifecycle.event(
        event_type=event_type, batch_data=batch_data, loss=loss, **kwargs
    )
    return ModifiedState(...)

EventType enum은 Events 글에서 다룬다. 핵심은 "배치 단위 hook 지점"이 이 메서드를 통해 모든 활성 Modifier에 브로드캐스트된다는 점이다. Modifier 한 쪽에서는 배치 훅을 받고, 다른 쪽에서는 훈련 옵티마이저 훅을 받을 수 있다. Lifecycle이 등록된 Modifier들을 순회하면서 각자의 관심 이벤트만 처리한다.

`reset` / `reset_stage`: 상태 재설정

def reset(self):
    """세션을 초기 상태로 완전 초기화 — 모델, 레시피, 모든 Lifecycle 상태 제거"""
    self._lifecycle.reset()

def reset_stage(self):
    """새 스테이지를 시작하기 위한 재설정 — recipe와 model은 유지"""
    self.lifecycle.initialized_ = False
    self.lifecycle.finalized = False

두 메서드의 차이는 "모델과 레시피까지 버릴지"이다. 멀티스테이지 레시피에서 pruning_stage 다음에 quant_stage를 실행할 때는 reset_stage를 쓴다. 레시피와 모델은 그대로 유지하면서 Lifecycle만 initialized_=False로 되돌려, 다음 initialize 호출이 Modifier 리스트를 새로 해석하게 한다.

`get_serialized_recipe`: 레시피 직렬화

def get_serialized_recipe(self) -> str | None:
    recipe = self.lifecycle.recipe
    if recipe is not None and hasattr(recipe, "yaml"):
        return recipe.yaml()   # Recipe 객체의 yaml() 호출
    logger.warning("Recipe not found in session - it may have been reset")

체크포인트 저장 시 레시피를 함께 직렬화하기 위한 헬퍼다. Recipe Metadata 글에서 본 Recipe.yaml()을 호출한다. 이 문자열은 compressed-tensors 체크포인트의 config에 embeded되어 "이 모델은 이 레시피로 만들어졌다"를 기록한다.

`active_session`: 전역 싱글톤 디스패처

세션이 전역으로 관리되는 이유는 편의성이다. src/llmcompressor/core/session_functions.py에 구현이 있다.

_global_session = CompressionSession()           # 프로세스 시작 시 한 번 생성
_local_storage = threading.local()               # 스레드별 세션 저장소
_local_storage.session = _global_session         # 메인 스레드는 전역 세션을 쓴다

@contextmanager
def create_session() -> Generator[CompressionSession, None, None]:
    """일시적으로 새 세션을 생성해 활성화. 컨텍스트 종료 시 이전 세션 복원."""
    global _local_storage
    orig_session = getattr(_local_storage, "session", None)
    new_session = CompressionSession()
    _local_storage.session = new_session
    try:
        yield new_session
    finally:
        _local_storage.session = orig_session    # 반드시 복원 — 예외 발생해도 OK


def active_session() -> CompressionSession:
    """현재 활성화된 세션 반환 — 스레드별로 다를 수 있음."""
    global _local_storage
    return getattr(_local_storage, "session", _global_session)

threading.local()을 쓴다는 점이 중요하다. 메인 스레드는 기본적으로 전역 세션을 공유하지만, 다른 스레드는 _local_storage.session이 초기화되지 않으므로 getattr의 fallback으로 전역 세션을 받게 된다. 이는 "대부분의 경우 전역 싱글톤처럼 동작하되, 필요하면 스레드별로 다른 세션을 쓸 수 있다"는 절충을 만든다.

create_session 컨텍스트 매니저는 테스트 격리에 특히 유용하다. 여러 테스트가 병렬로 돌아도 각 테스트는 자신만의 세션을 with 블록에서 생성해 서로 영향을 주지 않는다.

왜 이 설계인가

1. 세션 = 얇은 래퍼. CompressionSession은 자체 상태가 거의 없고 CompressionLifecycle에 모든 것을 위임한다. 이는 Lifecycle이 "진짜 상태 머신"이고, 세션은 "사용자 향 API + 싱글톤 경로"라는 역할 분리다. Lifecycle을 직접 쓸 수도 있지만, 대부분 사용자는 세션을 경유한다.

2. 전역 싱글톤 패턴. Python 로거와 동일한 접근. active_session() 한 줄로 어디서든 현재 세션을 가져올 수 있어, Modifier 내부에서 세션 인자를 받지 않아도 된다. 의존성 주입이 복잡해지는 것을 피한다. 단점은 테스트 격리가 어렵다는 것인데, create_session 컨텍스트 매니저가 이를 해결한다.

3. threading.local() 기반 스레드 격리. 단일 프로세스에서 여러 세션이 필요한 경우(예: 병렬 테스트, 여러 모델 동시 압축) 스레드별로 다른 세션을 쓸 수 있다. 일반 싱글톤보다 유연하다.

4. 훈련 루프 시그니처 유지. initialize가 optimizer, train_data, steps_per_epoch를 받는 것은 QAT(양자화 인식 훈련) 시나리오를 염두에 둔 것이다. oneshot PTQ만 쓰는 현재 주요 경로에서는 대부분 None이지만, 미래 확장성을 위한 자리를 비워두었다.

5. reset_stage vs reset. 멀티스테이지 레시피는 흔한 사용 패턴인데, 스테이지 사이에 "모델은 유지, 상태만 리셋" 기능이 없으면 불필요한 재로딩이 발생한다. 두 메서드의 분리로 세션 재사용 효율성을 높인다.

마무리

CompressionSession은 llm-compressor의 "세션 파사드"다. 얇은 레이어지만, 전역 싱글톤 패턴과 Lifecycle 위임 덕분에 사용자 코드가 간결해진다. 다음 글에서는 이 세션이 실제로 의존하는 Lifecycle 상태 머신을 분석한다.

참고 자료

llm-compressor 의 다른글

이전글 [llm-compressor] Recipe Metadata: 직렬화 헬퍼와 모델 메타데이터 구조
현재글 : [llm-compressor] CompressionSession: 전역 싱글톤 세션과 Lifecycle 래퍼
다음글 [llm-compressor] Lifecycle: Modifier 초기화-이벤트-종료 상태 머신