DeepSpeed와 Accelerate를 활용한 놀라운 속도의 BLOOM 추론

요약

Hugging Face는 8×80 GB A100 GPU에서 DeepSpeed‑Inference 텐서 병렬화와 Accelerate 파이프라인 병렬화를 사용하여 176B 파라미터 BLOOM 모델이 토큰당 1밀리초 미만의 속도로 생성될 수 있음을 보여줍니다.

하드웨어 요구사항 및 설정

  • 최적 구성: 8 × 80 GB A100 GPU (352 GB bf16 가중치). 대안으로는 2 × 8 × 40 GB A100, 2 × 8 × 48 GB A6000, 또는 24 × 32 GB V100이 있습니다.
  • 단일 노드 장점: 노드 내부 GPU 인터커넥트가 노드 간 링크보다 빠르며, 일반적으로 더 높은 처리량을 제공합니다.
  • 저사양 옵션: CPU 또는 NVMe 오프로드를 사용하면 작은 GPU에서도 BLOOM을 실행할 수 있지만, 생성 지연 시간이 크게 증가합니다.
  • 양자화 추론: BitsAndBytes를 통한 8‑bit 모델은 GPU 메모리를 대략 절반만 사용하지만, 약간의 처리량 감소가 있습니다.

벤치마크 환경

  • 노드: Jean Zay HPC, 8 × 80 GB A100, 512 GB CPU RAM, GPFS 스토리지 (~3 GB/s 읽기 속도).
  • 작업: 짧은 프롬프트와 KV‑캐시를 활성화한 상태에서 100 토큰을 탐욕적으로 생성 (max_length=100, do_sample=False).
  • 측정 지표: 모델 로딩 시간(초) 및 토큰당 처리량(ms per token = 전체 시간 ÷ (배치 × 토큰)).

모델 로딩 시간

솔루션 로드 시간 (초)
Accelerate 121
DeepSpeed‑Inference (shard‑int8) 61
DeepSpeed‑Inference (shard‑fp16) 60
DeepSpeed‑Inference (unsharded) 662
DeepSpeed‑ZeRO 462

Pre‑sharded DeepSpeed 체크포인트는 ~1 분에 로드되며, 비분할 체크포인트는 10 분 이상 걸릴 수 있습니다.

토큰 생성 처리량 (8 × 80 GB A100)

솔루션 (dtype) Batch 1 Batch 8 Batch 16 Batch 32 Batch 64 Batch 128 Batch 256 Batch 512
Accelerate bf16 230.38 ms 31.78 ms 17.84 ms 10.89 ms OOM – – –
Accelerate int8 286.56 ms 40.92 ms 22.65 ms 13.27 ms OOM – – –
DeepSpeed‑Inference fp16 44.02 ms 5.70 ms 3.01 ms 1.68 ms 1.00 ms 0.69 ms OOM –
DeepSpeed‑Inference int8 89.09 ms 11.44 ms 5.88 ms 3.09 ms 1.71 ms 1.02 ms 0.71 ms OOM
DeepSpeed‑ZeRO bf16 283 ms 34.88 ms OOM – – – – –

핵심 관찰

  • DeepSpeed‑Inference는 텐서 병렬화(TP)와 맞춤형 융합 CUDA 커널 덕분에 배치 128에서 토큰당 1 ms 미만을 달성합니다.
  • Accelerate는 단순 파이프라인 병렬화(PP)를 사용하여 배치 32에서 토큰당 약 10 ms에 도달하지만, 배치 64를 넘어서는 GPU 메모리 한계에 부딪힙니다.
  • 양자화된 int8 실행은 메모리 사용량을 절반으로 줄이며, DeepSpeed‑Inference int8은 배치 128에서도 토큰당 약 1 ms를 유지하지만 Accelerate int8은 더 일찍 메모리 부족이 발생합니다.

양자화된 8‑bit 처리량 (4 × 80 GB A100)

솔루션 Batch 1 Batch 8 Batch 16 Batch 32 Batch 64 Batch 128
Accelerate int8 284.15 ms 40.14 ms 21.97 ms OOM – –
DeepSpeed‑Inference int8 156.51 ms 20.11 ms 10.38 ms 5.50 ms 2.96 ms OOM

솔루션 상세

HuggingFace Accelerate

  • 접근 방식: 레이어 크기와 사용 가능한 메모리에 따라 모델 가중치를 지연 로드하며, 각 레이어당 하나의 GPU만 활성화되는 단순 파이프라인 병렬화를 사용합니다.
  • 장점: 어떤 하드웨어 구성에서도 바로 작동하며, GPU 메모리가 부족할 경우 CPU나 디스크로 오프로드할 수 있습니다.
  • 단점: 대부분의 순전파 과정에서 GPU가 유휴 상태가 되어 최고 처리량이 제한되고, 배치 크기가 커지면 빠르게 OOM에 도달합니다.
  • 사용법:
    pip install transformers>=4.21.3 accelerate>=0.12.0
python bloom-inference-scripts/bloom-accelerate-inference.py \
  --name bigscience/bloom --batch_size 1 --benchmark
# 8‑bit 양자화 실행
pip install bitsandbytes
python bloom-inference-scripts/bloom-accelerate-inference.py \
  --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

DeepSpeed‑Inference

  • 접근 방식: 텐서 병렬화는 각 레이어를 GPU에 분할하고, 맞춤형 융합 커널은 메모리 복사와 커널 호출을 감소시킵니다.
  • 성능 요인:
    1. TP vs PP – 모든 GPU가 동시에 계산하여 활용도를 높입니다.
    2. 융합 커널 – 메모리 오버헤드가 낮고 커널 호출이 적습니다.
  • 사전 분할된 체크포인트 (microsoft/bloom-deepspeed-inference-fp16)는 약 1분에 로드되며, 비분할 체크포인트는 10‑20분에 로드되지만 로드 후 동일한 속도로 실행됩니다.
  • 양자화된 int8 (microsoft/bloom-deepspeed-inference-int8)은 메모리 요구량을 절반으로 줄이며 8 × 80 GB A100에서 배치 128 시 토큰당 1 ms 미만을 유지하거나 4 × 80 GB A100에서도 동작합니다.
  • 사용법:
    pip install deepspeed>=0.7.3
# 빠른 TP‑사전분할 fp16
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name microsoft/bloom-deepspeed-inference-fp16
# 원본 체크포인트 (로드 느림)
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name bigscience/bloom
# 8‑bit 버전 (메모리 절반)
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

DeepSpeed‑ZeRO Inference

  • 접근 방식: 모델 상태를 GPU에 분할(ZeRO‑3 스타일)하고, 여러 독립적인 생성 스트림을 병렬로 실행하여 GPU 수에 비례하는 실질적인 속도 향상을 제공합니다.
  • 제한 사항: 제공된 스크립트는 모든 GPU에서 동일한 입력을 실행합니다; 이론적인 8배 또는 16배 속도 향상을 얻으려면 GPU별 맞춤 스트림이 필요합니다.
  • 오프로드 옵션: CPU 오프로드 또는 NVMe 오프로드를 통해 단일 GPU에서 추론이 가능하지만 큰 지연이 발생합니다.
  • 사용법:
    pip install deepspeed
# 다중‑GPU ZeRO 추론
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 1 --benchmark
# CPU 오프로드 (단일 GPU)
deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark
# NVMe 오프로드 (단일 GPU)
deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 8 \
  --nvme_offload_path=/path/to/nvme_offload --benchmark

커뮤니티 서버 및 클라이언트 확장

  • 서버 구현:
    • Mayank Mishra는 데모 스크립트를 바로 실행 가능한 웹 서버로 패키징했습니다.
    • Nicolas Patry는 고성능 Rust 기반 서버를 만들었습니다.
  • 클라이언트‑측 프로젝트:
    • Thomas Wang은 맞춤형 CUDA 커널 가속 BLOOM 모델을 개발 중입니다.
    • HuggingFace JAX 팀은 BLOOM을 위한 JAX 추론 백엔드를 공개했습니다.
  • 업데이트 유지: transformers-bloom-inference 저장소는 최신 스크립트와 서버 구현을 모아놓았습니다.

실용적인 시사점

  • 단일 노드에서 최대 원시 처리량을 원한다면 사전 분할된 fp16 체크포인트와 함께 DeepSpeed‑Inference를 사용하세요; 배치 128에서 토큰당 1 ms 미만을 기대할 수 있습니다.
  • GPU 메모리가 제한적일 때는 BitsAndBytes(DeepSpeed 또는 Accelerate)를 통한 8‑bit 양자화로 전환하여 메모리 요구량을 절반으로 줄이고 약간의 속도 손실을 감수합니다.
  • Accelerate는 이기종 하드웨어에 가장 유연한 솔루션이지만, 파이프라인 병렬화는 DeepSpeed의 텐서 병렬화보다 높은 지연을 초래합니다.
  • ZeRO 추론은 다수의 GPU로 확장하거나 오프로드를 통해 단일 GPU에서 실행할 수 있는 경로를 제공하며, 대규모 GPU 클러스터가 없을 때 유용합니다.

“이 경우 OOM 없이 실행할 수 있었던 최대 배치 크기는 128이었습니다.” – Hugging Face 블로그, 8 × 80 GB A100에서 DeepSpeed‑Inference int8을 설명하면서.

이 게시물은 2022년 9월 기준 BLOOM 추론 성능 상태를 반영합니다. 최신 최적화에 대해서는 transformers-bloom-inference GitHub 저장소를 참고하세요.

Sources