[transformers] Hugging Face Transformers: Continuous Batching에 Tensor Parallelism 도입하기

PR 링크: huggingface/transformers#45821 상태: Merged | 변경: +807 / -183

들어가며

대규모 언어 모델(LLM)을 서빙할 때, 여러 요청을 동시에 처리하는 Continuous Batching은 필수적인 기술입니다. 하지만 모델의 크기가 커질수록 단일 GPU 메모리로는 모델을 로드하기 어렵고, 추론 속도 또한 제한됩니다. 이번 Hugging Face transformers의 PR은 Continuous Batching 환경에서 Tensor Parallelism(TP)을 지원함으로써, 여러 GPU에 모델 가중치를 분산시켜 추론 성능을 비약적으로 향상시켰습니다.

주요 코드 분석

1. 벤치마크 도구의 TP 지원 확장

benchmark_v2/benchmark_scripts/continuous_batching_overall.py에서는 TP 환경을 고려한 모델 로딩 로직이 추가되었습니다.

Before:

model = AutoModelForCausalLM.from_pretrained(
    self.model_id, attn_implementation=self.attn_impl, device_map="auto"
)

After:

placement = {"tp_plan": "auto"} if self.tp_size > 1 else {"device_map": 0}
model = AutoModelForCausalLM.from_pretrained(self.model_id, attn_implementation=self.attn_impl, **placement)

기존의 device_map 방식 대신 tp_plan을 도입하여, TP 사용 시 모델이 자동으로 여러 GPU에 분산 배치되도록 개선되었습니다. 이는 TP가 활성화되었을 때 NCCL 통신 그룹이 올바르게 설정되도록 보장합니다.

2. KV Cache의 TP 대응

TP 환경에서는 KV Cache 또한 분산되어야 합니다. 리뷰 과정에서 kv_is_tp 플래그를 도입하여 TP 계획에 포함된 레이어에 대해서만 캐시 처리를 최적화하도록 했습니다.

kv_is_tp = "layers.*.self_attn.k_proj" in config.tp_plan and "layers.*.self_attn.v_proj" in config.tp_plan

이러한 조건부 로직은 불필요한 연산을 방지하고, TP 그룹 간의 데이터 불일치를 막는 핵심적인 역할을 합니다.

왜 이게 좋은가

이번 최적화의 핵심은 NCCL Graph Mixing 방지와 프로세스 간 일관된 시딩(Seeding)입니다.

1. 성능 향상: 벤치마크 결과에 따르면, ifeval_default 태스크에서 TP2를 적용했을 때 기존 대비 약 1.92배의 처리량(tok/s) 향상을 보였습니다. 이는 GPU 간 연산 병렬화가 효율적으로 이루어지고 있음을 의미합니다.
2. 일관성: 파이썬의 기본 hash() 함수는 프로세스마다 salt 값이 달라 분산 환경에서 문제를 일으킬 수 있습니다. 이를 수정하여 프로세스 간 일관성을 확보한 점은 분산 시스템 설계의 모범 사례입니다.
3. 교훈: 분산 추론 환경에서는 단순히 모델을 나누는 것을 넘어, 캐시 관리, 난수 생성, 통신 그룹 설정 등 세밀한 동기화가 필수적임을 보여줍니다.

결론

이번 PR은 Continuous Batching의 실용성을 한 단계 끌어올렸습니다. 특히 transformers 라이브러리가 대규모 모델 서빙을 위한 프로덕션 레벨의 도구로 진화하고 있음을 잘 보여줍니다.

참고 자료

• https://pytorch.org/docs/stable/distributed.html
• https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#tensor-parallelism

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