[llm-compressor] Basic Pipeline: 한 번의 forward로 끝내는 캘리브레이션

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

들어가며

BasicPipeline은 네 개의 파이프라인 중 가장 단순하다. 캘리브레이션 데이터로더를 순회하며 모델 전체에 대해 하나의 forward를 수행하고, Modifier의 훅이 필요한 통계를 누적한다. 그게 전부다. src/llmcompressor/pipelines/basic/pipeline.py의 코드를 해부한다.

핵심 구조/코드 분석

`BasicPipeline.call`

@CalibrationPipeline.register("basic")
class BasicPipeline(CalibrationPipeline):
    @staticmethod
    def __call__(
        model: torch.nn.Module,                 # 캘리브레이션 대상 모델
        dataloader: DataLoader,                 # 캘리브레이션 데이터
        dataset_args: Union["DatasetArguments", None],
    ):
        session = active_session()
        dispatch_model(model)                   # 모델을 적절한 디바이스에 자동 배치
        model_device = get_execution_device(model)
        use_loss_mask = (
            getattr(dataset_args, "use_loss_mask", False) if dataset_args else False
        )

        # AWQ 마스킹 지원을 위한 loss_masks 초기화
        if use_loss_mask:
            session.state.loss_masks = []

        LifecycleCallbacks.calibration_epoch_start()  # Modifier 들에게 epoch 시작 알림

        with contextlib.ExitStack() as stack:
            stack.enter_context(calibration_forward_context(model))  # 캘리브레이션 전용 컨텍스트

            for batch_idx, batch in enumerate(
                tqdm.tqdm(dataloader, desc="Calibrating")
            ):
                # AWQ 마스킹용 loss mask 수집
                if use_loss_mask:
                    session.state.loss_masks.append(batch.get("loss_mask"))

                session.state.current_batch_idx = batch_idx
                batch = tensors_to_device(batch, model_device)   # 배치를 GPU 로 이동
                model(**batch)                                   # forward 실행

        LifecycleCallbacks.calibration_epoch_end()  # Modifier 들에게 epoch 종료 알림

핵심 호출	의미
`dispatch_model(model)`	`compressed-tensors`의 자동 디바이스 배치. 단일 GPU면 전체를 GPU에, 여러 GPU면 파트별로 나눠 배치
`get_execution_device(model)`	입력 텐서를 올려야 할 디바이스 조회
`calibration_forward_context(model)`	Modifier 훅이 forward에 끼어들 수 있도록 하는 컨텍스트 매니저
`LifecycleCallbacks.calibration_epoch_start/end()`	내부적으로 `CompressionLifecycle.event(EventType.CALIBRATION_EPOCH_*)` 호출
`session.state.current_batch_idx = batch_idx`	어떤 배치가 현재 처리 중인지 기록 (Modifier가 배치별 가중치 적용 시 참조)
`model(**batch)`	실제 forward. Modifier 훅이 이 호출 안에서 활성화 통계를 수집

`run_calibration`: 편의 함수

def run_calibration(model: torch.nn.Module, dataloader: DataLoader):
    pipeline = BasicPipeline()
    pipeline(model, dataloader, None)

dataset_args가 필요 없을 때 쓰는 간편 래퍼다. 테스트나 간단한 사용자 코드에서 from llmcompressor.pipelines.basic import run_calibration 한 줄로 시작할 수 있다.

왜 이 설계인가

1. 단순함의 미덕. 이 파이프라인은 forward 루프 그 이상도 이하도 아니다. Modifier 가 forward 훅에 끼어들어 필요한 통계를 수집하므로, 파이프라인 자체는 배치 순회만 담당하면 된다. 역할이 명확하다.

2. dispatch_model로 멀티 GPU 지원. compressed-tensors의 dispatch_model을 호출하면 모델이 자동으로 여러 GPU에 분산되거나 CPU offload된다. basic 파이프라인은 "생성(generation)처럼 일반 추론과 동일한 디바이스 배치"를 사용한다. 즉 추론에 쓰는 디바이스 구성이 캘리브레이션에도 그대로 적용된다.

3. calibration_forward_context로 훅 활성화. PyTorch forward 호출 밖에서는 Modifier 훅이 비활성화되어 실제 추론에 방해가 되지 않는다. 이 컨텍스트 매니저가 훅을 "이 블록 안에서만" 활성화한다.

4. loss_masks 리스트. AWQ가 캘리브레이션 시 패딩 토큰을 제외하려면 loss mask가 필요하다. 일반 배치 데이터에 loss_mask 필드가 있으면 수집해 session.state.loss_masks에 쌓는다. 없으면 use_loss_mask=False로 스킵.

5. current_batch_idx 전파. GPTQ·AWQ 같은 Modifier는 "어떤 배치의 기여인지"를 알고 있어야 하는 경우가 있다. 세션 상태에 현재 배치 인덱스를 넣으면 Modifier가 session.state.current_batch_idx로 접근할 수 있어, 인자 전달 없이도 문맥을 공유한다.