Compute Budget 제한을 고려한 Google의 최신 Language Model 연구 동향

‘20년 GPT-3의 등장 이후 모델 사이즈와 데이터셋 사이즈를 늘려 모델 정확도를 늘리는 추세가 보편화되었다. 하지만 모델 사이즈와 데이터셋 사이즈의 급격한 증가는 Compute량의 급격한 증가를 초래하였다.

예를들어 NVIDIA와 MS는 530B MT-NLG 모델을 학습하기 위해 최대 420 대의 DGX-A100을 사용하였다. 즉, DGX-A100의 최대 파워 소모량은 6.5KW이므로 MT-NLG 모델을 학습하기 위해 1시간당 최대 2.7MW가 필요하다.

만일 목표한 성능(정확도)을 갖는 모델을 학습하는데 1개월이 걸린다면 약 2 GW의 파워가 필요하다.

• 6.5KW/1hour(DGX-A100의 최대 전력)* 420(DGX-A100 대수)* 720 hours (학습 시간. 1개월) = 1.966 GW

따라서 최근 LLM(Large Language Model) 학습은 결국 Compute Budget의 제한을 받기 시작하였다. 또한 ESG의 영향으로 인해 ML 학습 시 배출되는 탄소 량을 고려하기 시작하였다.

최근 구글은 에너지와 탄소 footprint를 줄이기 위한 모범사례 4 가지(4Ms)를 발표하였다. [1] 이를 통해 구글의 총 에너지 사용량 중 ML의 비중을 15% 미만으로 유지하고 있다고 한다.

• Model: Sparse 모델 사용 → 연산량을 3~10배 줄임
• Machine: ML 학습에 최적화된 프로세서와 시스템을 사용 → 성능 및 에너지 효율을 2x~5x을 개선
• Mechanization: 에너지 효율성이 뛰어난 클라우드를 이용하면 에너지 사용과 탄소 배출을 1.4x~2x까지 줄일 수 있음
• Map Optimization: 클라우드를 사용하는 소비자는 가장 깨끗한 에너지를 사용하는 위치를 선택할 수 있어 5x~10x의 탄소 footprint를 줄일 수 있음

이처럼 LLM 학습 시 에너지 사용량/탄소 배출이 고려되기 시작하면서 에너지 소비와 직접적인 연관이 된 Compute Budget을 제한하는 것이 중요한 과제가 되었다. 이와 같은 계기로 AI 업체들은 제한된 Compute Budget에서 가장 가성비가 높은 모델을 확보하는 방법을 찾기 시작하였다.

하지만 Compute Budget을 줄이기 위해선 학습 시간을 희생하는 일을 없어야한다. 다음 식은 학습 시간과 학습에 필요한 total flops의 관계를 간단하게 표현하였다. ( 주의. 분자/분모의 약어가 다르다.)

• Total FLOPs: 학습에 필요한 총 FLOPs로 Compute Budget과 직접적으로 연관되어 있으며 모델 파라미터와 데이터셋의 크기와 관계가 있다.
• Total FLOPS: 시간당 제공할 수 있는 총 FLOPS로 학습에 사용되는 총 가속기 성능과 관계가 있다.

위 수식에서 현실적으로 Total FLOPs(Floating Point Operations)를 줄이는 것이 Compute Budget과 학습 시간을 동시에 줄일 수 있음을 알 수 있다. 학습 시간을 유지한다면서 Compute Budget을 줄이면 학습 시스템의 Total FLOPS 요구량 또한 동시에 줄일 수 있기 때문에 에너지 소비를 줄일 수 있다.

결국 학습 시간을 유지하면서 Compute Budget을 줄이기 위해선 모델 학습에 필요한 Total FLOPs를 줄이는 것이 중요하다. 최근 구글 DeepMind와 Google Brain은 다음과 같이 Compute Budget을 고려한 LLM 학습과 관련된 연구 결과를 발표하였다. 이 논문들을 통해 앞으로 LLM의 발전 방향을 전망해본다.

• Training Compute-Optimal Large Language Model (‘22.04. DeepMind)
• ST-MoE: Designing Stable and Transferable Sparse Expert Models(‘22.02 → ‘22.04(revision) Google Brain)

1. Training Compute-Optimal Large Language Model(DeepMind)

지난해 DeepMind는 Gopher (280B. Dense Model)를 발표한데 이어 ‘22년도 4월달에 “Training Compute-Optimal Large Language Model” 논문을 통해 제한된 Compute Budge(FLOPs) 내에서 최고의 성능을 달성하기 위한 최적 모델 사이즈 및 학습 tokens를 trade-off하는 방식을 제안하였다. 이를 검증하기 위해 DeepMind는 Chinchilla라는 70B 모델을 소개하였다.

결과적으로 제안된 70B Chinchilla가 Gopher(280B), GPT-3(175B), Jurassic-1(178B) 및 MT-NLG(530B)보다 더 우수한 성능을 보인다. 이 논문이 제안한 optimal model scaling은 다음과 같다.

• Compute Budget를 증가시키면 모델 사이즈와 학습 데이터의 양은 비슷한 비율로 증가해야 한다. (Parameters가 x2 증가한다면 Tokens 또한 x2 증가해야 한다.)
• 더 큰 Compute budget에서 더 작은 모델이 더 최적이다.

아래 표는 주어진 크기의 모델를 위해 필요한 최적 학습 FLOPs와 학습 token을 평가한 것이다.

Chinchilla

ChinChilla는 Gopher와 동일한 compute budget을 사용하고 4배 이상의 데이터셋을 이용하여 Gopher보다 더 우수한 성능을 거두었다. 아래 표는 Gopher와 Chinchilla의 모델 구조를 비교하였다.

본 논문은 또한 다음과 같은 사실을 시사한다.

• 지금까지 대규모 언어모델 학습이 학습 tokens 수를 늘리지 않고 모델 크기만을 늘려왔지만 학습 데이터셋 확장이 더 중요하다.
• 데이터가 고품질일 때만 더 큰 데이터셋으로 확장하는 것이 이득이 된다.
• 큰 모델을 학습시키는 것은 비용이 많이 들어가기 때문에 최적의 모델 크기와 학습 단계를 미리 선택하는것이 필수 적이다.

2. ST-MoE: Designing Stable and Transferable Sparse Expert Models (Google Brain)

이 논문은 Google AI의 수장인 Jeff Dean을 비롯하여 Switch Transformer의 저자인 Barret Zoph, Noam Shazeer가 참여하였다. 이 논문에서 제안한 ST-MoE-32 모델(Stable Transferable MoE. 269B 모델이나 32B Dense Model과 동일한 FLOPS를 가짐)은 MoE 구조를 활용하여 학습 리소스를 절약하면서 PaLM 이전에 SOTA를 달성한 모델이다.

Sparse 모델은 매우 큰 파라미터에도 불구하고, 더 적은 compute 리소스를 사용하는 장점으로 인해 1 Trillion 이상의 모델은 Sparse 모델로 디자인하는 것이 대세가 되고 있다.

• e.g., Switch Transformer(1.6T), Glam(1.2T) 등

하지만 Switch Transformer는 SuperGLUE와 같은 common benchmark로 fine-tuning 되었을 때 더 작은 모델보다 성능이 떨어지는 경우가 발견된다. GlaM의 경우, 학습 시 불안정성이 문제가 발견되어 Sparse 모델 학습은 parameters와 computation의 균형이 필요함을 시사하고 있다.

본 논문의 목표는 sparse 모델의 “실용성”과 “안정성”을 증가시키는 것이다. 이를 위해 위 두 가지 이슈에 대해서 연구하고 SuperGLUE를 포함한 경쟁력 있는 NLP benchmark로 fine-tuning 했을 때 SOTA를 달성한 ST-MoE-32 모델을 pre-training 하였다. 또한 sparse 모델에 대한 추가적인 분석과 디자인을 제시하였다.

Sparse Models 안정화하기

(1) Multiplicative interactions 제거

• Multiplicative interactions은 모델의 품질을 개선시키지만, 학습을 불안정화시킴
• Multiplicative interactions를 가진 동작인 GEGLU(GELU Gated Linear Unit)와 RMS(Root Mean Square) normalization을 개별적으로 제거하는 것은 학습을 안정화시키지만 품질을 저하시킴
• GEGLU 레이어를 제거할 때, FLOPs 및 파라미터가 일치하도록 GEGLU 레이어와 동등한 Dense-ReLU-Dense 레이어로 대체함

(2) 노이즈 주입

• 입력 지터와 Dropout 모두 안정성을 개선하지만 모델 품질을 상당히 저하시킨다.
• 입력 지터는 라우터의 입력 logits에 [1–10^-2, 1+10^-2] 사이의 균일한 랜덤 변수를 곱함
• Dropout을 Transformer 전체에 적용함

(3) Activations 와 Gradients 제한하기

• 뉴럴 네트워크를 안정화시키기 위한 가장 성공적인 방법은 activations와 gradients를 제한하는 것임
• 이 논문은 메모리 효율성을 때문에 Adafactor optimizer를 사용하며 gradient clipping 대신 update clipping을 사용함
• Update clipping과 router z-loss 둘다 모두 모델을 안정화시키나 update clipping은 모델 품질을 상당히 해침.
• 따라서 모델 품질과 안정성을 개선하기 위해 router z-loss을 사용함

• 낮은 precision format을 사용하면 a) 프로세서와 메모리 간의 통신 비용, b) 계산 비용, c) 텐서를 저장하기 위한 메모리(e.g. activation)를 줄임으로써 보다 효율적인 모델을 만들 수 있음
• Sparse expert 모델은 router로 인해 더 많은 지수 함수를 가지고 있기 때문에 반올림 에러에 민감함
• Router z-loss는 모델이 작은 값의 logits를 생산하도록 하여 더 정확한 모델을 만들 수 있음. → 모든 지수화된 텐서들을 float32로 캐스트됨

Sparse Models 디자인하기

(1) Expert의 수

• Expert의 개수가 많을수록 유리하나, expert의 수가 256 이상이면 그 이점이 빠르게 감소하고, 매우 sparse한 모델(expert의 1%만 활성화하는 경우)인 경우 또한 그 이점이 빠르게 감소됨
• 즉, Switch Transformer(512개의 expert)와 같이 다수의 expert 보다는 크고 256 개 이하의 expert가 유리하다 라는 것으로 해석이 가능함

(2) Core 당 expert 수

• 텐서를 컴퓨팅 코어에 로드하는 시간이 텐서를 연산하는 시간을 초과하면 그 모델은 메모리 bound가 되므로 OI(Operational Intensity) 비율을 증가시키는 것이 효율성을 개선함.
• Expert 수를 늘리면 수행되는 계산량은 그대로지만 메모리 전송 요구사항이 증가하므로 OI 비율을 낮추는 결과를 초래함
• 이 논문은 Core 당 2개 이상의 expert는 효율성을 해치므로 TPU에서는 Core당 ≤ 1개 expert를 추천함

(3) Capacity Factor(CF)

• 학습 및 평가 Capacity Factor를 모두 증가시키면 품질이 향상됨
• 하지만 CF가 증가하면 학습 속도가 느려짐. 파레토-효율적 CF는 특정 하드웨어 시스템에 의해 결정됨.
• 이 논문은 CF=1.25를 추천함

(4) Routing Algorithm

• 동일한 CF에서 top-(n+1)과 top-n routing을 비교하였을 때 매우 작은 이득만 있는 것으로 확인되었다. 이 논문에서는 top-2 routing을 추천한다.

특정 HW/SW 시스템은 최적 top-n routing과 CF를 결정한다. 예를들어 시스템이 빠른 all2all과 allreduce 통신을 지원하면, 더 큰 capacity factors와 더 큰 top-n routing이 최적이다. (반대로 all2all & allreduce 통신이 느리면 더 작은 CF가 유리하다.)

향후 전망

• LLM 모델의 방향은 당분간 1) 극한의 모델 성능을 실험하기 위한 모델(e.g. PaLM)과 2) Compute Budget을 고려한 가성비 모델(e.g. chinchilla, MoE기반 모델)으로 양분되어 발전될 가능성이 크다.
• 실제 추론 서비스를 하기 위해 상대적으로 모델 파라미터 사이즈가 작은 가성비 모델이 사용될 것이다.
• Sparse MoE 모델은 Compute Budget 제한된 상황에서 선택할 수 있는 모델 구조이나, 모델 성능이 학습 시스템의 all2all, allreduce와 같은 collective comm.의 속도에 의존할 것이다.

레퍼런스

[1] Good News About the Carbon Footprint of Machine Learning Training

Good News About the Carbon Footprint of Machine Learning Training

[2] Training Compute-Optimal Large Language Model

Training Compute-Optimal Large Language Models

[3] ST-MoE: Designing Stable and Transferable Sparse Expert Models

ST-MoE: Designing Stable and Transferable Sparse Expert Models