BLOOM(BigScience Large Open-science Open-access Multilingual Language Model)

AI 스타트업 hugging face는 공개 협력 프로젝트인 BigScience는 통해 언어모델 BLOOM(176B)을 최근 공개하였다.

BLOOM(BigScience Large Open-science Open-access Multilingual Language Model)의 이름에서 알 수 있듯이 BLOOM은 Multi-lingual 언어모델인데 Multi-lingual이 왜 중요한 것일까?

AI 모델을 학습할 때 가장 쉬운 방법은 각 언어마다 모델을 따로따로 학습시키는 것이다. 하지만 N개의 언어가 추가될 때마다 N개의 학습 모델이 학습을 해야하므로 학습 비용 및 serving 비용이 선형적으로 증가할 수 밖에 없다. 이와 같은 현실적인 문제로 인해 multi-lingual 모델을 선호할 수 밖에 없다.

BLOOM overview

BLOOM의 하드웨어/소프트 웨어 설정 및 주요 기술 구성 요소는 다음과 같다.

Model 아키텍처

• 176 B 파라미터 (70 layers, 112 attention heads, 14336 hidden layers, 2048 seq. length)
• Megatron-LM GPT-2를 수정
• decoder-only 아키텍처
• 워드 임베딩 레이어에 적용된 layer normalization (StableEmbedding)
• ALiBI positional encodings w/ GeLU activation

학습 데이터셋

• 46개 언어
• 13개 프로그래밍 언어
• 전처리된 텍스트(1.6TB)를 350B unique tokens로 변환
• 모델의 어휘 크기 250,680 (tokens)
• 학습 데이터셋의 언어 분포

학습 하드웨어

• GPUs: 384 NVIDIA A100 80GB GPUs (48 nodes) + 32 spare gpus
• 8 GPUs per node Using NVLink 4 inter-gpu connects, 4 OmniPath links
• CPU: AMD EPYC 7543 32-Core Processor
• CPU memory: 512GB per node
• GPU memory: 640GB per node
• Inter-node connect: Omni-Path Architecture (OPA) w/ non-blocking fat tree
• NCCL-communications network: a fully dedicated subnet
• Disc IO network: GPFS shared with other nodes and users

학습 소프트웨어

• Megatron-DeepSpeed
• DeepSpeed
• PyTorch (pytorch-1.11 w/ CUDA-11.5)
• apex

Speeds, Sizes, Times

• 학습 시간: 2022/3/11 ~ 2022/7/5 (약 1백만 컴퓨팅 시간 소요)
• Checkpoint size: full checkpoint w/ optimizer states : 2.3TB, BF16 weights: 329GB
• 학습 throughput: 150 TFLOPS per GPU
• 1 epochs
• 학습 비용: $2M~5M (클라우드 비용으로 환산시, 예비실험 포함)

사용방법

• 직접 사용: 텍스트 생성, 언어모델에 의해 생성된 특성 연구
• 다운스트림 사용: 정보 추출, 질문 답변, 요약
• 이외 작업에 사용되면 범위 외 사용 및 오용 가능성이 있음

Megatron-DeepSpeed

BLOOM은 Megatron-DeepSpeed을 사용하여 학습되었다.

BigScience의 Megatron-DeepSpeed는 원본 Megatron-DeepSpeed repository를 fork하여 몇가지 addition을 추가하였다.

DP(Data Parallelism), TP(Tensor Parallelism), PP(Pipeline Parallelism)을 제외한 Megatron-DeepSpeed의 features는 다음과 같다.

BF16 Optimizer

• 거대한 LLM 모델을 FP16으로 학습하는 것은 절대해선 안된다.
• BF16은 FP32와 같은 지수를 가지므로 overflow가 발생하지 않는다.
• Optimizer는 FP32 weight를 업데이트한 뒤 연산 시 16bit 포맷으로 캐스트하여 사용한다.
• 학습을 정밀하게 유지하기 위해선 FP32으로 gradient accumulation을 구현하는 것이 중요하며, 이것이 BF16 Optimizer의 역할이다.

Fused CUDA Kernels

• Fused CUDA kernel을 사용하면 중간 결과를 메모리에 복사하지 않으므로 오버헤드를 절약하고 GPU idle 상태를 방지한다.

임베딩 LayerNorm

• 104B가 발산하는 것을 방지하기 위해 첫번째 단어 임베딩 후 추가 LayerNorm을 추가하면 학습에 훨씬 더 안정적인 것을 발견함
• StableEmbedding (normal Embedding w/ layernorm) & uniform xavier initialization) 사용

위치 인코딩

• 일반적인 위치 임베딩을 AliBi로 대체함 (논문: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation)
• 이를 통해 모델이 학습한 것보다 더 긴 입력 시퀀스를 처리 가능함 (즉, 길이가 2048인 시퀀스로 학습된 모델은 훨씬 더 긴 시퀀스도 처리 가능함)

레퍼런스

[1] bigscience/bloom

bigscience/bloom · Hugging Face

[2] The Technology Behind BLOOM Training

The Technology Behind BLOOM Training