[llm-compressor] Sequential Pipeline: 레이어 단위 서브그래프 캘리브레이션
들어가며
GPTQ는 메모리를 많이 먹는다. 한 레이어의 헤시안 역행렬 하나만 해도 수 GB급이다. 70B 모델 전체에 GPTQ를 한 번에 적용하면 단일 GPU로는 불가능하다. 해법은 레이어 단위로 쪼개서 처리하는 것이다. 한 레이어의 캘리브레이션이 끝나면 헤시안을 해제하고 다음 레이어로 넘어간다. 중간 활성화는 CPU에 오프로드해서 GPU 메모리를 비운다. 이 전략을 구현한 것이 src/llmcompressor/pipelines/sequential/pipeline.py의 SequentialPipeline이다.
핵심 구조/코드 분석
SequentialPipeline.__call__: 세 단계 흐름
@CalibrationPipeline.register("sequential")
class SequentialPipeline(CalibrationPipeline):
@staticmethod
@handle_sequential_oom # OOM 시 디버깅 메시지 데코레이터
def __call__(
model: torch.nn.Module,
dataloader: DataLoader,
dataset_args: "DatasetArguments",
):
session = active_session()
# 1) 디바이스 준비
onload_device = get_main_device() # GPU 예: cuda:0
offload_device = torch.device(dataset_args.sequential_offload_device) # cpu 또는 cuda:1
dispatch_for_sequential(model, onload_device)
# 2) 서브그래프 분할 — 모델을 sequential_targets 단위로 쪼갬
sequential_targets = infer_sequential_targets(
model, dataset_args.sequential_targets
)
ignore = dataset_args.tracing_ignore
sample_input = next(iter(dataloader))
subgraphs = trace_subgraphs(
model,
sample_input,
sequential_targets,
ignore,
dataset_args.sequential_targets_per_subgraph,
)
num_subgraphs = len(subgraphs)
# 3) 캘리브레이션 에폭 시작
LifecycleCallbacks.calibration_epoch_start()
# QAT 비호환 Modifier 감지 — GPTQ/AWQ/AutoRound 는 QAC 강제 비활성
disable_qac = any(
type(mod).__name__ in DISABLE_QAC_MODIFIERS
for mod in session.lifecycle.recipe.modifiers
)
| 단계 | 하는 일 |
|---|---|
| 1. 디바이스 준비 | onload/offload 디바이스를 정하고 모델을 sequential 모드로 재배치 |
| 2. 서브그래프 분할 | torch.fx로 모델을 추적해 LlamaDecoderLayer 단위로 쪼갬 |
| 3. 캘리브레이션 | 각 서브그래프에 대해 두 번의 forward (아래 설명) |
sequential_targets는 보통 LlamaDecoderLayer, Qwen2DecoderLayer 같은 트랜스포머 디코더 블록 이름이다. 사용자가 명시하지 않으면 infer_sequential_targets가 HF 모델 정의의 no_split_params 속성을 참조해 자동 결정한다.
이중 forward: calibration + propagation
핵심 루프는 두 번의 forward로 구성된다.
with contextlib.ExitStack() as stack:
stack.enter_context(calibration_forward_context(model))
# 양자화 인식 캘리브레이션(QAC) 비활성 — GPTQ 등 호환 불가 Modifier 감지 시
if not dataset_args.quantization_aware_calibration or disable_qac:
stack.enter_context(DisableQuantization(model))
# 활성화 캐시 생성 — 데이터로더로부터 초기 입력 캡처
activations = IntermediatesCache.from_dataloader(
dataloader, onload_device, offload_device
)
for subgraph_index, subgraph in enumerate(subgraphs):
calib_desc = f"({subgraph_index + 1}/{num_subgraphs}): Calibrating"
prop_desc = f"({subgraph_index + 1}/{num_subgraphs}): Propagating"
num_batches = len(dataloader)
with disable_offloading(): # 이 서브그래프의 모듈을 GPU 에 고정
# === Pass 1: Calibration — Modifier 훅 활성, 통계 수집 ===
for batch_idx, inputs in _get_batches(
activations,
num_batches,
subgraph.input_names,
calib_desc,
sequential_prefetch,
):
session.state.current_batch_idx = batch_idx
subgraph.forward(model, **inputs)
LifecycleCallbacks.sequential_epoch_end(subgraph)
# === Pass 2: Propagation — 훅 비활성, 압축 후 출력 캡처 ===
with HooksMixin.disable_hooks():
for batch_idx, inputs in _get_batches(
activations,
num_batches,
subgraph.input_names,
prop_desc,
sequential_prefetch,
):
output = subgraph.forward(model, **inputs)
if subgraph_index < num_subgraphs - 1:
# 다음 서브그래프의 입력으로 쓸 수 있게 캐시 갱신
activations.update(batch_idx, output)
activations.delete(batch_idx, subgraph.consumed_names)
LifecycleCallbacks.calibration_epoch_end()
두 번의 forward가 필요한 이유가 이 파이프라인의 핵심이다.
- Pass 1: Calibration. 각 배치의 입력을 현재 서브그래프에 통과시키면서 Modifier 훅이 활성화 통계를 수집한다. 이 pass가 끝나면
sequential_epoch_end이벤트가 발생하고, Modifier가 이 레이어의 가중치를 최종화한다(GPTQ는 헤시안을 풀고 가중치를 업데이트). - Pass 2: Propagation. 훅을 끈 상태로 같은 배치를 다시 한 번 통과시킨다. 이제 가중치는 이미 양자화된 상태다. 이 pass의 출력이 "양자화된 레이어의 실제 출력"이므로, 이 값을 다음 레이어의 입력으로 사용해야 오차 전파를 정확히 시뮬레이션할 수 있다.
activations.update(batch_idx, output)이 다음 서브그래프의 입력을 갱신한다. activations.delete(batch_idx, subgraph.consumed_names)은 현재 서브그래프가 소비한 활성화를 캐시에서 제거해 메모리를 회수한다.
_get_batches와 프리페치
def _get_batches(
activations: IntermediatesCache,
num_batches: int,
input_names: list[str],
desc: str,
sequential_prefetch: bool = False,
) -> Iterator[tuple[int, dict]]:
batch_source = (
activations.iter_prefetch(input_names) # 백그라운드 스레드로 미리 오프로드 → 온로드
if sequential_prefetch
else activations.iter(input_names)
)
for batch_idx, inputs in tqdm(
enumerate(batch_source), total=num_batches, desc=desc
):
yield batch_idx, inputs
sequential_prefetch=True이면 IntermediatesCache가 다음 배치의 활성화를 백그라운드 스레드에서 GPU로 미리 옮긴다. 이는 "현재 배치 forward 중에 다음 배치 활성화를 오프로드 디바이스에서 가져오는" 오버랩을 만든다. GPU가 놀지 않고, 오프로드 I/O 시간이 숨겨진다.
왜 이 설계인가
1. 서브그래프 단위 분할. 전체 모델을 한 번에 GPU에 올리지 못해도 "현재 처리 중인 서브그래프 + 입력 활성화"만 GPU에 있으면 된다. 70B 모델도 80GB GPU 하나로 양자화 가능하다. 모델이 클수록 이 전략의 이점이 크다.
2. 이중 forward. GPTQ 같은 알고리즘은 **"이 레이어의 출력이 양자화 후에 어떻게 변하는가"**를 다음 레이어에 전달해야 한다. Pass 2가 없으면 다음 레이어는 "양자화되지 않은 이 레이어의 이상적 출력"을 입력으로 받게 되고, 오차 누적이 무시된다. 정확도 손실이 커진다.
3. 활성화 오프로드. 각 배치의 서브그래프 입력을 IntermediatesCache가 CPU(또는 다른 GPU)로 내린다. calib_samples * num_subgraphs 개의 활성화가 디스크나 RAM에 보관되는데, 실제로는 한 서브그래프 분량만 GPU에 로드된다.
4. disable_offloading 컨텍스트. 한 서브그래프 처리 중에 그 서브그래프의 모듈 가중치가 CPU로 오프로드되면, 이중 forward 시 매번 GPU로 다시 올려야 한다. 이 컨텍스트는 "현재 서브그래프는 메모리에 고정"을 보장한다.
5. QAC 자동 비활성화. 양자화 인식 캘리브레이션(QAC)은 forward 시 가중치가 이미 양자화된 상태로 계산한다. GPTQ·AWQ·AutoRound는 캘리브레이션 중 원본 가중치를 보아야 하므로 QAC와 충돌한다. llm-compressor는 DISABLE_QAC_MODIFIERS 목록에 이들을 등록해두고, 감지되면 자동으로 DisableQuantization 컨텍스트로 QAC를 끈다.
마무리
Sequential Pipeline은 llm-compressor의 기본값이자, 대형 LLM 양자화의 핵심이다. 이 파이프라인 덕분에 워크스테이션에서 70B 모델을 GPTQ로 양자화할 수 있다. 다음 글은 서브그래프 사이의 활성화를 오프로드하는 Intermediates Cache를 본다.
참고 자료
관련 포스트
- [llm-compressor] Basic Pipeline: 한 번의 forward로 끝내는 캘리브레이션
- [llm-compressor] Dataset Calibration: c4/wikitext/ultrachat 로더
- [llm-compressor] Quantization Calibration: update_weight_zp_scale와 observer 등록
- [llm-compressor] Intermediates Cache: 서브그래프 활성화 오프로드 캐시
- [llm-compressor] Data-Free & Independent Pipeline: 데이터 없는 파이프라인과 Modifier별 개별 실행
llm-compressor 의 다른글
- 이전글 [llm-compressor] Basic Pipeline: 한 번의 forward로 끝내는 캘리브레이션
- 현재글 : [llm-compressor] Sequential Pipeline: 레이어 단위 서브그래프 캘리브레이션
- 다음글 [llm-compressor] Data-Free & Independent Pipeline: 데이터 없는 파이프라인과 Modifier별 개별 실행
댓글