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 저장소를 참고하세요.