BLOOM 176B 훈련 기술 개요

TL;DR

BLOOM의 176 B 파라미터 모델은 384 × 80 GB NVIDIA A100 GPU를 사용해 3.5 개월 동안 맞춤형 Megatron‑DeepSpeed 스택으로 훈련되었으며, 이 스택은 ZeRO 데이터 병렬 처리, 텐서 병렬 처리, 파이프라인 병렬 처리를 BF16 혼합 정밀도와 결합해 프랑스의 Jean Zay 슈퍼컴퓨터에서 1 M GPU‑시간 실행을 가능하게 했습니다.

프로젝트 요약

  • Hardware: 384 A100 80 GB GPU(48 노드, 노드당 8 GPU)와 AMD EPYC 7543 CPU, 노드당 512 GB RAM, Omni‑Path 인터커넥트, 전용 NCCL 서브넷, GPFS 스토리지.
  • Software: Megatron‑DeepSpeed (Microsoft DeepSpeed와 NVIDIA Megatron‑LM을 병합한 포크).
  • Architecture: GPT‑3 스타일 트랜스포머에 추가 개선 사항(ALiBi 위치 인코딩, 임베딩 LayerNorm)을 적용.
  • Dataset: 46개 언어에서 350 B 토큰(≈1.5 TB 정제 텍스트), 어휘 크기 250 680.
  • Training duration: 약 3.5 개월(≈1 M 컴퓨팅 시간).

주요 기여자

이 작업은 여섯 주요 그룹에 의존했습니다:

  1. Hugging Face BigScience 팀(내부 엔지니어 및 자금 지원).
  2. Microsoft DeepSpeed 팀(DeepSpeed 라이브러리 제공 및 통합 지원).
  3. NVIDIA Megatron‑LM 팀(프레임워크 및 조언).
  4. Jean Zay 슈퍼컴퓨터를 관리하는 IDRIS/GENCI 직원(컴퓨팅 자원 기부).
  5. PyTorch 핵심 팀(버그 수정 및 사용성 개선).
  6. BigScience 엔지니어링 작업 그룹의 자원봉사자들.

주목할 만한 인물로는 Olatunji Ruwase, Deepak Narayanan, Jeff Rasley, Jared Casper, Samyam Rajbhandari, Rémi Lacroix 등이 있습니다.

Megatron‑DeepSpeed 스택

Component DeepSpeed 제공 Megatron‑LM 제공
ZeRO 데이터 병렬 처리 ✅
텐서 병렬 처리 ✅
파이프라인 병렬 처리 ✅
BF16 옵티마이저 ✅
통합 CUDA 커널 ✅
DataLoader ✅

이 스택은 3‑D 병렬 처리를 구현합니다:

  • Data Parallelism (DP): 모델을 GPU 그룹에 복제하여 각 GPU가 데이터 조각을 처리합니다.
  • Tensor Parallelism (TP): 개별 텐서를 GPU에 분할하여 GPU당 메모리를 감소시킵니다.
  • Pipeline Parallelism (PP): 모델을 수직으로 GPU에 분할하고, 파이프라인에서 마이크로 배치를 처리하여 GPU 유휴 시간을 방지합니다.
  • ZeRO: 옵티마이저 상태, 그래디언트, 필요 시 가중치를 추가로 분할하여 대형 모델에 맞춥니다.

병렬 처리 세부 사항

ZeRO 데이터 병렬 처리

전체 모델을 복제하는 대신, 각 GPU는 파라미터, 그래디언트, 옵티마이저 상태의 일부만 보유하고 필요에 따라 전체 텐서를 실시간으로 재구성합니다. 이는 메모리 오버헤드를 크게 줄입니다.

텐서 병렬 처리

가중치 행렬은 GPU 간에 열 단위로 분할되며, 각 GPU는 행렬 곱의 해당 부분을 계산하고 로컬에서 활성화를 적용합니다. 이는 고속 인터커넥트를 필요로 하며, BLOOM에서는 TP 차수를 노드당 4(텐서 조각당 GPU 1개)로 제한했습니다.

파이프라인 병렬 처리

모델 레이어를 단계로 나누고 마이크로 배치를 단계별로 흐르게 하여 모든 GPU가 바쁘게 동작합니다. "chunks"(또는 GAS) 하이퍼파라미터는 마이크로 배치 수를 제어하여 파이프라인 버블과 마이크로 배치 크기의 균형을 맞춥니다. BLOOM은 GPU 간 메모리 균형을 위해 72개의 파이프라인 단계(임베딩 단계 2개 포함)를 사용했습니다.

결합된 DP + PP + TP (3‑D 병렬 처리)

최종 훈련 구성은 데이터 분배에 DP, 텐서 분할에 TP, 레이어 분배에 PP를 사용하여 384‑GPU 클러스터를 효율적으로 활용했습니다.

BF16 옵티마이저

FP16으로 훈련하면 이전 실험에서 발산이 발생했습니다(예: 104 B 모델). BLOOM은 맞춤형 BF16Optimizer를 사용해 BF16 혼합 정밀도로 전환했으며, 이 옵티마이저는:

  • FP32의 지수 범위를 유지하여 오버플로를 방지합니다.
  • 모든 누적을 FP32로 수행합니다.
  • 파이프라인 마이크로 배치 전반에 걸쳐 그래디언트를 FP32로 누적합니다.

이 옵티마이저는 176 B 모델에 대해 안정적인 손실 곡선을 가능하게 했습니다.

통합 CUDA 커널

GPU 유휴 시간을 최소화하기 위해 Megatron‑LM의 맞춤형 통합 커널을 다음에 사용했습니다:

  • LayerNorm
  • 스케일링, 마스킹 및 소프트맥스 결합
  • 편향이 추가된 GeLU (PyTorch JIT 사용)

이 커널들은 중간 결과를 레지스터에 유지함으로써 메모리 트래픽을 감소시킵니다.

데이터셋 처리

훈련 데이터 파이프라인:

  • 정제된 다국어 텍스트 1.5 TB를 토큰화하여 350 B 토큰으로 변환했습니다.
  • 고정 시퀀스 길이 2048에 대한 샘플별 인덱스를 생성했습니다.
  • 에포크별 순서를 섞어 균일한 노출을 보장했습니다.
  • 재시작 시 재계산을 방지하기 위해 인덱스를 디스크에 저장했습니다.
  • 여러 데이터셋을 가중치를 조정하여 혼합했습니다.

아키텍처 조정

  • Embedding LayerNorm: 첫 번째 임베딩 뒤에 LayerNorm을 추가하여 훈련을 안정화했으며, 이는 bitsandbytes의 StableEmbedding 구현에서 영감을 받았습니다.
  • ALiBi 위치 인코딩: 절대 위치 임베딩을 선형 바이어스 어텐션(ALiBi)으로 교체하여 훈련 길이(2048)보다 긴 시퀀스로의 외삽을 가능하게 했습니다.

엔지니어링 과제

  • 하드웨어 오류: 주당 1–2개의 GPU 오류가 발생했으며, 3 시간마다 체크포인트를 저장해 오류당 손실 작업을 약 1.5 시간으로 제한했습니다.
  • 소프트웨어 버그: CUDA_LAUNCH_BLOCKING=1 설정, 옵티마이저 그룹 분할, 다중 사용자 작업 제어를 위한 맞춤형 SLURM 킬 스위치가 필요했습니다.
  • 다운타임: 교착 상태, 디스크 공간 부족 등으로 5–10 시간의 중단이 있었지만 전체 훈련은 계획된 3.5개월 내에 완료되었습니다.
  • 온콜 조정: 유럽과 캐나다 서부에 분산되어 전용 호출 장치 없이 24/7 모니터링이 가능했습니다.

결론

가장 집중적인 단계는 두 달에 걸친 준비 작업이었으며, 여기에는 BF16 옵티마이저의 후기 개발과 대규모 병렬 처리 디버깅이 포함되었습니다. 스택이 안정화된 후 176 B 모델은 원활하게 훈련되었으며, 충분한 컴퓨팅 자원과 협업 도구가 제공될 경우 오픈소스 팀도 최첨단 다국어 모델을 훈련할 수 있음을 보여주었습니다.

리소스

핵심 논문

  • Megatron‑LM: Efficient Large‑Scale Language Model Training on GPU Clusters (arXiv:2104.04473)
  • DeepSpeed ZeRO: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv:1910.02054)
  • ALiBi: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (arXiv:2108.12409)

Sources