HuggingFace가 Megatron-DeepSpeed를 통한 3D 병렬화로 384개 A100 GPU에서 176B 파라미터 모델 176일 만에 3.5개월 완성

Context

대규모 언어 모델 학습 방식의 기술적 세부사항이 공개되지 않으면서, 엔지니어들이 유사한 규모의 모델을 학습할 때 재현 불가능한 문제가 있었다. 176B 파라미터의 BLOOM 모델은 이러한 기술 격차를 해소하기 위해 하드웨어 구성, 소프트웨어 스택, 분산 학습 전략을 공개적으로 구축하고 공유하려고 기획되었다.

Technical Solution

• Megatron-DeepSpeed 채택: DeepSpeed의 ZeRO 셰딩과 파이프라인 병렬화를 Megatron-LM의 텐서 병렬화와 결합한 3D 병렬화 구현
• GPU 인프라 구성: 384개 NVIDIA A100 80GB GPU(48개 노드, 노드당 8 GPU)를 NVLink와 4개 OmniPath 링크로 연결
• BF16 옵티마이저 구현: 마지막 순간에 BF16 옵티마이저를 작성하고 fp32 옵티마이저 상태와 bf16+fp32 가중치로 2.3TB 크기의 체크포인트 생성
• 네트워크 아키텍처: NCCL 통신용 전용 서브넷 및 GPFS 공유 디스크 IO 구성으로 48개 노드 규모 학습에서만 나타나는 문제 발견 및 해결
• 학습 데이터 구성: 59개 언어 350B 토큰으로 2022년 3월~7월에 약 100만 컴퓨트 시간 소요

Impact

학습은 예상과 다르게 3.5개월간 비교적 안정적으로 진행되어 대부분의 시간에 1명의 모니터링 담당자만 필요했으며, 소수의 경우에만 여러 사람이 문제 해결에 참여했다.