[transformers] Hugging Face Transformers: logits_to_keep을 활용한 메모리 최적화

PR 링크: huggingface/transformers#46660 상태: Merged | 변경: +18 / -12

들어가며

LLM 추론 환경에서 Continuous Batching을 사용할 때, 모델의 lm_head는 전체 시퀀스 길이에 대해 로짓(logits)을 계산합니다. 하지만 실제 다음 토큰 예측에 필요한 로짓은 특정 인덱스에만 국한되는 경우가 많습니다. 기존 방식은 전체 로짓을 계산한 뒤 외부에서 슬라이싱을 수행했는데, 이는 불필요한 메모리 점유와 연산 낭비를 초래했습니다. 이번 PR은 logits_to_keep 파라미터를 활용해 모델 내부에서 미리 로짓을 슬라이싱함으로써 메모리 풋프린트를 줄이고 처리량을 개선하는 최적화를 다룹니다.

코드 분석

src/transformers/generation/continuous_batching/model_runner.py

핵심 변경 사항은 ModelRunner가 모델의 logits_to_keep 지원 여부를 확인하고, 이를 batch_data에 전달하는 로직입니다.

Before:

# 기존에는 모델이 전체 로짓을 반환한 후 외부에서 슬라이싱 수행
logits = model(**batch_data).logits
logits_indices = batch_data["logits_indices"]
logits = logits[:, logits_indices, :]

After:

# 모델이 logits_to_keep을 지원하면 내부에서 슬라이싱 수행
if self.supports_logits_to_keep(model):
    batch_data["logits_to_keep"] = batch_data["logits_indices"]

# ... (forward pass 이후)
if "logits_to_keep" not in batch_data:
    logits = logits[:, logits_indices, :]  # 모델이 처리하지 않은 경우에만 외부 슬라이싱

supports_logits_to_keep 메서드를 추가하여 모델이 해당 기능을 지원하는지 동적으로 확인합니다. 이를 통해 lm_head 연산 직전에 불필요한 토큰의 로짓 생성을 방지하여 GPU 메모리 사용량을 최적화합니다.

왜 이게 좋은가

이번 최적화의 핵심은 '불필요한 연산의 조기 차단'입니다. lm_head는 보통 모델의 마지막 레이어에서 큰 가중치를 차지하며, 전체 시퀀스에 대해 연산할 경우 메모리 대역폭을 많이 소모합니다. logits_to_keep을 사용하면 필요한 토큰에 대해서만 lm_head 연산을 수행하므로 다음과 같은 이점이 있습니다.

1. 메모리 풋프린트 감소: 대규모 프리필(prefill) 배치에서 메모리 사용량이 줄어듭니다.
2. 처리량(Throughput) 향상: 불필요한 연산이 제거되어 토큰 생성 속도가 소폭 상승합니다. 실제 벤치마크 결과 gsm8k 데이터셋 등에서 약 0.1% ~ 1.1%의 처리량 향상을 보였습니다.
3. 범용성: 기존 로직을 파괴하지 않고, 지원하는 모델에 대해서만 선택적으로 적용되도록 설계되어 안정적입니다.

일반적인 교훈으로, 추론 엔진을 설계할 때는 '데이터가 어디서 생성되고 어디서 소비되는가'를 파악하여, 소비되지 않을 데이터를 최대한 앞단에서 제거하는 것이 성능 최적화의 핵심임을 알 수 있습니다.

참고 자료

• https://huggingface.co/docs/transformers/main/en/model_doc/llama#transformers.LlamaForCausalLM.forward.logits_to_keep

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.