[llm-compressor] Intermediates Cache: 서브그래프 활성화 오프로드 캐시

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

들어가며

Sequential Pipeline은 모델을 서브그래프로 쪼개서 순차 처리한다. 한 서브그래프가 끝나면 그 출력이 다음 서브그래프의 입력이 되어야 한다. 하지만 수백 배치의 활성화를 모두 GPU에 보관하면 메모리가 부족하다. 해결책은 "오프로드 디바이스(보통 CPU)에 저장하고, 필요할 때 다시 GPU로 올리는" 캐시다. 이것이 src/llmcompressor/pipelines/cache.py의 IntermediatesCache다.

핵심 구조/코드 분석

`IntermediateValue` dataclass

@dataclass
class IntermediateValue:
    """오프로드된 값과 원래 디바이스를 함께 저장"""
    value: torch.Tensor | "IntermediateValue" | Any   # 오프로드된 텐서 또는 중첩 값
    device: torch.device | None                        # 원래 디바이스 (None 이면 오프로드 안함)

각 "저장된 값"은 텐서만이 아니다. dataclass나 튜플처럼 중첩 구조도 재귀적으로 저장할 수 있다. value가 텐서면 device는 "이 텐서를 어느 디바이스로 온로드할지"를 기억한다. 예를 들어 원본 입력이 cuda:0의 텐서였으면, 오프로드 시 CPU로 옮겨두고 device=cuda:0으로 기록해두었다가, 나중에 fetch할 때 다시 cuda:0으로 올린다.

`IntermediatesCache` 필드

class IntermediatesCache:
    batch_intermediates: list[IntermediateValues]    # 배치별 딕셔너리 리스트
    offload_device: torch.device | None              # 저장 시 옮길 디바이스 (cpu / cuda:1 / None)

    # 같은 텐서가 여러 곳에서 참조될 때 중복 저장 방지
    offload_values: WeakKeyDictionary[torch.Tensor, torch.Tensor] = WeakKeyDictionary()

    def __init__(
        self,
        batch_intermediates: list[IntermediateValues] | None = None,
        offload_device: torch.device | None = "cpu",
    ):
        self.batch_intermediates = batch_intermediates or []
        self.offload_device = offload_device

필드	의미
`batch_intermediates`	`batch_intermediates[batch_idx]`는 그 배치의 이름→값 딕셔너리
`offload_device`	`cpu` 또는 `cuda:1` 같은 대기 디바이스. `None`이면 오프로드 비활성
`offload_values`	같은 텐서가 두 번 오프로드되지 않도록 WeakKey 기반 중복 제거 맵

WeakKeyDictionary를 쓰는 이유는 중요하다. 입력 텐서가 여러 서브그래프에서 공유될 수 있는데, 같은 텐서를 두 번 CPU로 복사하면 메모리 낭비다. WeakKey 맵은 원본 텐서를 약참조로 보관해서 "이미 오프로드한 텐서는 기존 것을 재사용"하게 한다. 원본이 GC되면 맵에서 자동으로 제거된다.

생성자 팩토리: `empty`와 `from_dataloader`

@classmethod
def empty(cls, num_batches: int, offload_device: torch.device):
    """빈 캐시 생성 (나중에 update() 로 채움)"""
    batch_intermediates = [{} for _ in range(num_batches)]
    return cls(batch_intermediates, offload_device)

@classmethod
def from_dataloader(
    cls,
    dataloader: torch.utils.data.DataLoader,
    model_device: torch.device = torch.device("cpu"),
    offload_device: torch.device | None = torch.device("cpu"),
):
    """데이터로더의 배치들을 초기 활성화로 삼아 캐시를 채움"""
    # 각 배치를 iterate 하면서 model_device 에서 시작, offload_device 로 오프로드
    ...

Sequential Pipeline은 주로 from_dataloader를 쓴다. 데이터로더의 초기 배치가 모델의 "첫 서브그래프 입력"이 되고, 이후 서브그래프의 출력이 update()로 캐시에 쌓인다.

오프로드/온로드 메커니즘 (개념)

IntermediatesCache의 실제 fetch/store 구현은 다음과 같은 로직을 따른다.

# store (의사 코드)
def store(self, batch_idx: int, name: str, value):
    if isinstance(value, torch.Tensor):
        # 중복 제거 맵 조회
        if value in self.offload_values:
            offloaded = self.offload_values[value]
        else:
            offloaded = value.to(self.offload_device)
            self.offload_values[value] = offloaded
        self.batch_intermediates[batch_idx][name] = IntermediateValue(
            value=offloaded, device=value.device
        )
    elif is_dataclass(value):
        # 각 필드를 재귀적으로 오프로드
        ...
    elif isinstance(value, tuple):
        # 튜플 요소들을 재귀적으로 오프로드
        ...

# fetch (의사 코드)
def fetch(self, batch_idx: int, names: list[str]) -> dict:
    result = {}
    for name in names:
        iv = self.batch_intermediates[batch_idx][name]
        result[name] = iv.value.to(iv.device) if iv.device else iv.value
    return result

저장 시 offload_device(보통 cpu)로 옮기고, 가져올 때 원래 디바이스(iv.device)로 복원한다. 중요한 것은 "오프로드 비용은 저장 시 한 번, 온로드 비용은 가져올 때마다"라는 것이다. Sequential Pipeline이 이중 forward를 할 때, 두 pass 모두 같은 활성화를 가져오므로 두 번의 온로드가 발생한다.

`iter`와 `iter_prefetch`

Sequential Pipeline에서 가장 자주 호출되는 메서드는 iterator다.

def iter(self, input_names: list[str]) -> Iterator[dict]:
    """동기적으로 각 배치의 활성화를 yield"""
    for batch_idx in range(len(self.batch_intermediates)):
        yield self.fetch(batch_idx, input_names)

def iter_prefetch(self, input_names: list[str]) -> Iterator[dict]:
    """백그라운드 스레드에서 다음 배치를 미리 온로드"""
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(self.fetch, 0, input_names)
        for batch_idx in range(len(self.batch_intermediates)):
            current = future.result()   # 대기 중이던 결과 받기
            if batch_idx + 1 < len(self.batch_intermediates):
                future = executor.submit(self.fetch, batch_idx + 1, input_names)
            yield current

프리페치는 한 개의 백그라운드 스레드로 "다음 배치의 온로드"를 미리 시작한다. 현재 배치가 forward 중일 때 다음 배치가 CPU→GPU로 옮겨지고 있으므로, 메인 스레드가 다음 배치를 기다릴 필요가 없다. 이는 Sequential Pipeline의 sequential_prefetch=True 모드를 뒷받침한다.

단점은 GPU 메모리가 2배로 필요하다는 것이다. 현재 배치와 다음 배치가 동시에 GPU에 있기 때문이다. 따라서 GPU 메모리 여유가 있을 때만 활성화하는 것이 좋다.

`update`와 `delete`

def update(self, batch_idx: int, new_values: dict) -> None:
    """서브그래프 forward 후 새로운 활성화를 캐시에 저장"""
    for name, value in new_values.items():
        self.store(batch_idx, name, value)

def delete(self, batch_idx: int, names: list[str]) -> None:
    """현재 서브그래프가 소비한 활성화를 캐시에서 제거 (메모리 회수)"""
    for name in names:
        self.batch_intermediates[batch_idx].pop(name, None)

Sequential Pipeline의 propagation pass에서 activations.update(batch_idx, output)과 activations.delete(batch_idx, subgraph.consumed_names)이 차례로 호출된다. 새 활성화를 쌓으면서 오래된 것을 버리는 LRU 같은 효과를 낸다.

왜 이 설계인가

1. dataclass와 튜플의 재귀 지원. Transformer 모델의 중간 출력은 보통 BaseModelOutput 같은 dataclass다. 각 필드를 개별 텐서로 쪼개지 않고 구조를 유지한 채 캐싱할 수 있어야 Modifier가 원본 타입을 그대로 받을 수 있다. IntermediateValue의 재귀 구조가 이를 가능하게 한다.

2. WeakKeyDictionary 기반 중복 제거. 같은 텐서가 여러 곳에 참조될 때 한 번만 오프로드한다. 약참조 덕에 원본이 해제되면 오프로드 복사본도 자동 정리된다. 수동 GC가 필요 없다.

3. offload_device=None 옵션. 캐시를 GPU에만 두고 오프로드하지 않는 모드도 지원한다. 작은 모델에서는 CPU 오프로드 오버헤드가 오히려 속도를 떨어뜨리므로, 이 옵션이 유용하다.

4. ThreadPoolExecutor 기반 프리페치. 단일 백그라운드 스레드로 다음 배치 온로드를 겹친다. 멀티프로세스가 아닌 스레드를 쓰는 이유는 CUDA 컨텍스트 공유와 파이썬 GIL에서 벗어난 I/O(tensor copy) 때문이다. 온로드 작업은 GIL을 해제하는 CUDA 호출이라 스레드로 충분한 병렬성을 얻는다.

5. 단순한 key-value 인터페이스. store/fetch/update/delete의 네 메서드만 노출한다. Sequential Pipeline은 이 API만 쓰면 되고, 내부 저장 구조는 자유롭게 최적화할 수 있다.

마무리

Intermediates Cache는 Sequential Pipeline의 "조용한 영웅"이다. 이 캐시가 없었다면 70B 모델 양자화는 불가능했을 것이다. 프리페치와 오프로드의 단순한 조합이 메모리와 속도의 균형을 만든다. 다음 글은 Modifier 프레임워크의 기본 클래스인 Modifier Base를 본다.

참고 자료

소스 코드:
- src/llmcompressor/pipelines/cache.py
- src/llmcompressor/pipelines/sequential/pipeline.py
관련 포스트:

llm-compressor 의 다른글

이전글 [llm-compressor] Data-Free & Independent Pipeline: 데이터 없는 파이프라인과 Modifier별 개별 실행
현재글 : [llm-compressor] Intermediates Cache: 서브그래프 활성화 오프로드 캐시
다음글 [llm-compressor] Modifier Base: 모든 Modifier가 상속하는 기반 클래스

댓글

관련 포스트

llm-compressor 의 다른글