[리뷰] Scaling Language Models: Methods, Analysis & Insights from Training Gopher

Google의 Deepmind는 최근 Gopher라는 280B LLM(Large Language Model)을 발표하였다. Gopher는 GPT-3와 마찬가지로 transformer 기반 LLM(Large Language Model)로 MassiveWeb, C4, Wikipedia, Github, books, news 기사 등을 포함한 Massive Text(10.5TB)으로 학습되었다. Gopher는 44M~7.1B 모델과 같이 작은 사이즈를 갖는 Gopher family와 Gopher(280B)로 구성되어 있으며, 모든 모델은 300B 토큰(12.8% of MassiveText)으로 학습되었다.

Gopher Family & Model Architecture

DeepMind는 최소 44M부터 최대 280B 파라미터의 6개의 Transformer LMs (Language Models)를 실험하였으며 280B 파라미터 모델을 Gopher로 명명하고 나머지 모델들의 집합은 Gopher Family로 명명하였다.

Gopher는 기존 transformer 아키텍처에 두가지 수정사항을 반영하였다.

• LayerNorm 대신 RSMNorm을 사용함
• Absolute positional encoding이 아닌 Relative positional encoding로 교체함. 학습했던 것 보다 더 긴 시퀀스를 평가할 수 있도록 허용함.

또한 32K의 단어를 가진 SentencePiece를 사용하여 텍스트를 토큰화시키고 open-vocabulary 모델링을 지원하기 위해 바이트 수준 backoff를 사용한다.

Training

다음은 DeepMind가 Gopher Family를 학습하기 위해 사용한 방법이다.

• 300B tokens (with 2048 token context window using Adam optimizer)
• 첫 1500 steps 동안 1*e-7에서 최대 lr(e.g., Gopher 280B: 4*e-5)로 warm-up
• cosine 스케줄을 사용하여 10x decay
• Model size가 증가할수록 max. lr를 줄이고 각 배치의 토큰 개수를 증가함
• Gopher의 batch size를 학습 동안 batch 당 3M → 6M 토큰으로 증가시킴
• 학습 throughput을 증가시키고 메모리를 줄이기 위해 BF16를 사용함
• 7.1B보다 더 작은 모델은 FP32 parameters와 BF16 activations의 Mixed Precision을 이용하여 학습함
• 7.1B과 280B 모델은 BF16 parameters와 activations를 사용함
• BF16 parameters는 안정성을 유지하기 위해 stochastic rounding을 사용하여 업데이트됨. 나중에 stochastic rounding은 mixed precision 학습 성능을 완전히 복원하지 못함을 발견함

Infrastructure

DeepMind는 JAX와 Haiku를 이용하여 학습 및 평가 코드 베이스를 구축하였다. 특히 data & model parallelism 둘다를 효율적으로 표현하기 위해 JAX’s pmap transformation를 사용하였다. 그리고 모든 Gopher family 모델은 TPUv3 칩에서 학습과 평가되었다.

Gopher의 16bits parameters와 32bits Adam state는 2.5 TiB를 차지한다. (각 TPUv3 코어의 가용 메모리 용량은 16GB) 이 메모리 문제를 해결하기 위해 TPU 메모리에 맞도록 아래 3가지 방법을 적용하였다.

• Optimizer state partitioning (from DeepSpeed ZeRO)
• Model parallelism (from Megatron-LM)
• Rematerialisation

Gopher 학습하는 동안, data & model parallelism 둘다 빠른 cross-chip 통신덕분에 TPUv3s 상에 낮은 오버헤드를 가지며(10%의 오버헤드)를 갖는다. 그러므로 pipelining은 training scale이 1024-chip pod를 초과할 때까지 불필요하다. 요약하면 Gopher 학습 시, TPU pods 이내 data & model parallelism을 사용하고 TPU pods간에는 pipelining을 사용한다.

Training Dataset

DeepMind는 웹페이지, 책, 뉴스 기사와 코드와 같은 다중 소스로부터 Large English-Language Text 데이터셋의 컬렉션인 MassiveText로 Gopher와 Gopher family 모델들을 학습하였다. MassiveText은 2.35B 다큐먼트 (또는 10.5TB 텍스트)를 포함한다. DeepMind는 MassiveText 데이터셋에서 12.8%인 300B 토큰을 Gopher 학습에 사용하였는데 서브셋(books, news, etc)당 특정 샘플링 비율로 MassiveText로부터 sub-sample하였다. 또한 다운스트림 성능을 최대화하기 위해 샘플링 비율을 튜닝하였다.

결과

DeepMind는 152 tasks에서 Gopher family의 성능을 측정한 후, 이 결과를 이전 언어 모델들의 SOTA 성능과 비교하였다.

124 tasks를 대상으로 SOTA 언어모델 성능과 Gopher의 성능을 비교하였을 때 100 tasks(81%)에서 현 SOTA 언어모델을 능가하는 결과를 보여준다. 비교 대상 SOTA 언어모델은 GPT-3(175B), Jurassic-1(178B), MT-NLG(530B)이다.

Gopher는 Reading Comprehension, Humanities & Ethics, STEM & Medicine에서 특히 가장 균일한 개선이 있음을 확인할 수 있다. 일반적인 경향은 추론이 많은 작업(e.g., Abstract Algebra)에서 개선이 적고, General Knowledge와 같이 지식 집약적인 테스트에서 더 크고 더 일관적인 개선을 보여준다.

Lessons Learned

Adafactor

Adam Optimizer 대신 Adafactor Optimizer를 사용하였다. Adafactor를 사용한 이유는 더 적은 memory footprint를 사용하기 때문이다. 작은 규모의 Gopher 모델의 경우, Adafactor를 사용하여 pretraining시 안정적이고 성능이 좋다는 사실을 발견한 반면, 큰 규모의 Gopher 모델(7.1B)에서 Adafactor를 사용하면 불안정성의 증가와 함께 Adam에 비해 성능이 줄어드는 효과를 발견하였다. 7.1B 파라미터의 Gopher 모델을 Adafactor로 학습할 때 Adam 베이스라인과 비교하여 minor loss 발산을 볼 수 있다. Adam의 max. lr은 1.2*e-4인 반면, Adafactor는 max. lr가 6*e-5에도 불구하고 불안정함을 보여준다.

Finetuning 시 Adafactor를 사용하면 발산하기 쉽고 lr과 batch size와 같은 하이퍼파라미터 설정에 어려움이 있다.

Lower-Precision Training with BF16

FP16 activations와 parameters으로 학습하면 FP16의 제한된 수치 범위로 인해 학습이 불안정하므로 FP32에 비해 학습 성능 저하가 없는 BF16(bfloat16)이 제안되었다. Gopher는 BF16 activations과 parameters로 학습하였지만, 작은 lr과 parameter updates의 사이즈로 인해 많은 parameters가 많은 학습 steps에서 업데이트 되지 않았다.

DeepMind는 Gopher 모델 중 417M 모델을 대상으로 BF16과 FP32의 영향도를 테스트한 후 다음과 같은 사실을 발견하였다.

• Optimizer update를 순수하게 FP32를 유지하는 것이 최선이다.
• Optimizer update에 사용된 FP32 parameters은 다시 forward pass 시 BF16으로 cast한다.
• FP32와 BF16의 mixed precision을 사용하여 BF16 학습에 비해 memory footprint는 약간 증가하지만 FP32 학습과 맞먹는 성능과 학습 속도를 개선할 수 있다.

4개의 테스트 구성 조건은 다음과 같다.

• (1) FP32 everywhere: parameters와 activations 둘다 FP32로 저장되며, 가장 많은 메모리를 사용하지만 가장 정확도가 높다.
• (2) BF16 parameters w/o Random Rounding: parameters와 activations이 BF16으로 cast 되며 parameter 업데이트 동안, random rounding이 사용되지 않는다.
• (3) BF16 parameters w/ Random Rounding: parameters와 activations이 BF16으로 cast 되며 parameter 업데이트 동안, random rounding이 사용된다. Parameters는 거리에 비례하여(in BF16 space) 랜덤하게 round-up 또는 round-down된다.
• (4) BF16 parameters with FP32 Copy in the partitioned Optimizer State: parameters와 activations이 BF16으로 cast 된다. 그러나 parameters의 사본이 optimizer state에 FP32로 저장되고 parameters 업데이트 시 사용된다. Parameters는 forward pass 시 BF16으로 랜덤하게 round된다.

모든 구성에서 loss 함수의 attention softmax와 softmax cross-entropy를 계산하기 위해 FP32를 사용하며 이는 low-precision training을 안정화한다. BF16을 사용하는 모든 방법은 FP32 everywhere 구성에 비해 1.4x의 속도 개선이 있다.

위의 결과에서 보듯, (2) 구성이 가장 최악의 성능을 보이며, (1)은 가장 많은 memory footprint를 사용하지만 베이스라인 성능 역할을 한다. (4)은 성능은 FP32와 거의 동일한 반면, 더 적은 memory footprint를 사용하면서 1.4x 속도 개선이 있다.

레퍼런스

[1] Scaling Language Models: Methods, Analysis & Insights from Training Gopher

Scaling Language Models: Methods, Analysis & Insights from Training Gopher