[리뷰] M6–10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
알리바바의 DAMO 아카데미는 최근 세계에서 가장 큰 10 Trillion 파라미터를 가진 M6–10T을 발표하였다. M6 (MultiModality-to-MultiModality Multitask Mega-transformer)모델은 그 이름에서 알 수 있듯 멀티모달, 멀티태스킹 능력에 중점을 두고 보편성을 목표로 하는 모델이다.
알리바바는 M6–10T를 NVIDIA-V100 512 GPUs로 10일 이내 학습하여 놀라운 성능에 도달하였다. (A100 GPU 1024개로 34 일이 걸린다.) 또한 GPT-3와 동일한 규모라면 M6는 GPT-3 학습에 필요한 에너지의 1% 만을 소모한다.
M6의 Model Architecture
M6는 self-attention과 FFN(Feed-Forward Neural Networks)로 구성된 transformer 레이어를 적층하여 만들며 dense 모델을 sparse 전문가 모델로 변환하기 위해 FFN 레이어를 MoE(Mixture of Expert)로 교체하였다. MoE는 다수의 전문가로 구성되어 있으며, 다른 디바이스상에 분산되어 있는 FFNs이다.
M6의 게이팅 네트워크(Gating Network)는 각 토큰의 디스패치 및 조합 행동을 결정하므로 토큰들은 다양한 디바이스에서 처리될 수 있다. 이러한 메커니즘은 data parallelism과 expert parallelism의 조합이므로 model capacity가 거대하더라도 매우 효율적이다.
M6 모델은 이해와 생성을 동시에 학습하기 위해 일반 텍스트 데이터에 대한 텍스트 노이즈 제거 및 언어 모델링으로 학습되며, 멀티 모달 데이터에 대한 이미지기반 텍스트 노이즈 제거 및 이미지 캡셔닝으로 학습된다.
2-Stage 학습 전략: “Pseudo-to-Real”
알리바바는 M6–10T을 학습하기 위해 “Pseudo-to-Real”이라는 2 stage 학습 전략을 이용한다. 첫번째 stage에서 더 적은 수의 파라미터를 사용하지만 큰 연산 그래프(“Pseudo Giant”)를사용하여 M6–10T 모델을 학습한다. 두번째 stage에서는 공유 레이어와 연결 해제된 weight로 초기화된 해당 큰 모델(“Real Giant”)을 학습한다.
1st Stage(“Pseudo” Stage): Layer-wise parameter sharing
“Pseudo” 단계의 핵심은 레이어 간 파라미터를 공유하는 Pseudo-Giant를 학습하는 것이다. 교차 레이어 파라미터 공유(Cross-layer parameter sharing)는 더 적은 양의 파라미터를 유지하는 동안 만족스러운 성능을 유지할 수 있다는 사실이 증명되었다.
교차 레이어 파라미터 공유는 말 그대로 transformer layer 간에 같은 파라미터를 공유하며 사용하는 것이다. 교차 레이어 파라미터 공유는 Universal Transformer 등에서 처음 언급되는데 아래 그림에서 보는 것과 같이 Transformer의 각 레이어의 Output이 다시 다음 레이어의 Input으로 들어가는 형태로 Recursive Transformer라고도 할 수 있다.
Pseudo 단계 학습은 교차 레이어 파라미터 공유을 통해 통신 오버헤드를 해결할 수 있고 훨씬 적은 메모리 footprint를 사용하기 때문에 이점을 얻을 수 있다. 또한 더 적은 파라미터를 사용하므로 메모리 용량에 의해 제한받지 않으므로 더 큰 batch size로 학습이 가능하다.
모든 레이어에 걸쳐 파라미터를 공유하는 L 개의 레이어를 가진 M6 모델를 만든다고 가정하면, Pseudo Giant의 파라미터와 optimizer states의 개수는 1/L이며, gradients의 개수는 2/L이 된다. 이러한 메모리 절감은 더 큰 batch size로 더 빠르게 학습할 수 있도록 한다.
Pseudo stage의 메모리 효율성을 이용하여, 더 많은 experts를 Single GPU에 분배하고 experts를 더 많은 GPU에 분할할 수 있다. 각 MoE 레이어에 512개의 experts를 가진 모델을 가정하면, 각 GPU에 2개의 experts가 있는 256개 GPUs만으로 Pseudo Giant를 학습시킨 후, 각 GPU에 1개의 expert만 있는 512개 GPUs에서 Real Giant를 학습할 수 있다.
2nd Stage(“Real” Stage): Delinking the shared parameters
Pseudo Giant와 비교하여 교차 레이어 파라미터 공유가 없는 large model의 이름은 “Real Giant”이다. Pseudo Giant와 Real Giant 모두 계산 그래프를 공유하지만 파라미터의 개수가 다르다. 이 논문에서는 두 모델간에 연결을 구축하는 방법을 발견하였다.
주어진 Pseudo Giant가 수렴될 때까지 완전히 학습되면, Real Giant 학습을 가속화하기 위해 교차 레이어 파라미터 공유의 연결을 해제한다. 이와 같이 “Pseudo-to-Real”를 사용하면 large model을 “from scratch”에서 학습할 필요없이 모델이 낮은 perplexity로 수렴을 시작할 수 있다.
Real Giant는 Pseudo Giant와 동일한 구조인 덕분에 각 Real Giant의 각 레이어는 Pseudo Giant의 파라미터로 초기화딜 수 있다. 즉 Real Giant의 초기 파라미터는 full-trained Pseudo Giant와 동일하다.
이와 같이 극도로 간단한 학습 전략은 M6–10T와 같은 extreme-scale 모델에 매우 유용하다. Pseudo 단계는 더 빠른 수렴을 많은 시간을 절약하고, Real 단계는 상당한 양의 계산 리소스를 사용할 수 있다.
Timing for Switching
Pseudo 단계 학습의 가장 큰 장점은 훨씬 더 빠른 수렴 속도이나 Pseudo Giant의 성능은 제한된 파라미터 수에 의해 제한된다. Pseudo Giant를 최종적으로 수렴될 때까지 훈련하는 것은 분명히 많은 시간 낭비를 초래할 것이다.
알리바바는 수렴 속도를 기반으로 Pseudo 단계에서 Real 단계로 전환하기 위한 학습 스텝을 결정하기 위해 간단한 전략을 제안한다. Pseudo 단계의 훈련 중 고정된 interval로 학습 단계를 Real 단계로 옮기고 잠시 학습(e.g., 30분)을 시도하여 Real Giant의 성능을 평가한다. 그런 다음, Real 단계와 동일한 학습 시간 동안 Pseudo 단계의 학습을 계속한다. Real 단계의 loss 감소 속도가 Pseudo 단계의 속도를 초과하면 평가된 학습 단계를 다음 단계 학습(여기서는 Real 단계)을 위한 최상의 switching point로 결정한다.
실험
M6 모델의 upstream과 downstream 성능을 관찰하여 학습 전략의 효율성을 평가한 결과이다.
실험 셋업
먼저 약 16GB의 일반 텍스트가 있는 말뭉치인 BookCorpus 및 English Wikipedia으로 M6–10T를 pretrain한다. Upstream 및 downstream 성능을 분석하기 위해 중간 크기의 M6 모델을 구현하여 “Pseudo-Real” 학습 전략의 효율성을 검증한다. 2 단계 학습이 학습 효율성에서 차이를 만들 수 있는 높은 메모리 footprint의 요구사항을 만족하기 위해 Single NVIDIA V100(32GB)에서 학습할 수 있는 가장 큰 모델인 1.4B 파라미터 모델을 실험에 사용한다.
- vocabulary of around 50K subword
- sequence length: 512
- N_layer = 36
- d_model = 1024
- d_ff = 16384
Real Giant는 1.4B의 파라미터 개수를 가지며 해당 Pseudo Giant는 90M의 파라미터 개수를 갖는다. M6의 Baseline 모델은 358M 파라미터, d_ff=4096, N_layer=24로 구성된다. Pseudo 단계와 Real 단계 양쪽으로 학습한 모델은 “Pseudo-to-Real”를 의미하는 “P2R”로 명명한다.
그밖의 학습을 위한 하이퍼파라미터와 조건은 다음과 같다.
- AdamW optimizer 사용
- Pseudo stage의 lr = 0.0002
- Real sate의 lr=0.00008
- Cosine decaying mechanism for learning rate scheduling & warm-up ratio of 0.001
- Total batch size=6114
- Micro-batch size = 32
- Gradient accumulation step = 4
- 48 NVIDIA V100 GPU
각 base, pseudo, real의 더 좋은 비교를 위해 2가지 방식을 사용한다.
1)Limited budget: 각 모델별 동일한 시간 학습했을 때 모델 성능을 측정. (Pseudo: 25K steps, Real: 4.7K steps, P2R: 12K steps)
2)Training for longer: 더 나은 비교를위해 M6–1B의 실제 성능을 보여주는 방식
Downstream Evaluation
Downstream task에 사용된 데이터셋은 다음과 같다.
- WikiText103
- Gigaword
두 downstream task에 대해 M6-base 모델은 Pretrained UniLM, Megatron-LM의 성능과 동등하다. 실험 결과는 Pseudo-to Real 학습이 효과적으로 학습 속도를 높일 수 있다라는 가정을 뒤받침해 준다. 또한 제한된 시간(limited budget)에서, P2R 모델이 두 downstream task에서 Pseudo 모델과 Real 모델의 성능을 능가한다.
Real Giant 모델은 from scratch로 학습하기 때문에 시간이 많이 소모되지만, Pseudo Giant 모델은 Real Giant 모델에 비해 수렴 속도가 5배정도 더 빠르다. From scratch로 학습된 P2R과 Real Giant의 loss 수렴속도를 관찰하였을 때 P2R의 log perplexity는 Real Giant보다 훨씬 더 빨리 감소한다.
Towards A 10 Trillion-Parameter Model
Model Step
알리바바는 기존 1 Trillion 모델보다 10배 큰 10 Trillion 모델인 M6-T10을 설계하면서 메모리를 효율적으로 활용하기 위해 MoE(Mixture of Experts)를 채택하고 모든 FFN 레이어를 메모리 효율적인 MoE 레이어로 교체하였다. 또한 모델 품질 개선과 학습 안정성을 위한 expert prototyping [2]을 적용하였다.
[M6–10T 모델의 상세한 정보]
- N_layer = 48
- d_model = 1024
- d_ff = 9984
- N_experts = 10240 (distributed on multiple devices)
- 80 prototypes of experts
[M6–10T 모델 학습을 위한 학습 시스템 정보]
- batch size per GPU = 8
- EFLOPS[3] 적용
- 8-GPU workers로 구성된 클러스터를 100 Gb RDMA 네트워크로 연결됨
- 각 worker당 CPU 메모리 750 GB
- Expert 분배와 granular offloading 전략은 Whale 프레임워크[4]을 이용함
Granular CPU offloading(세분화된 CPU 오프로딩)
M6–10T 모델 학습에 더 적은 리소스로 더 큰 모델에 CPU 메모리를 활용하기 위해 기존 CPU offloading보다 효율성이 높은 granular CPU offloading을 적용하였다.
학습 과정은 “Forward(Fn)”, “Backward(Bn)”,”Apply(An)” 단계로 구성된다. Fn 및 Bn에서 모든 모델 파라미터를 CPU로 offloading하려면 CPU에서 GPU 메모리로 파라미터를 두번 로드해야 한다. Activation의 재계산을 가져오는Activation checkpointing은 Bn에서 로드된 파라미터를 필요로 한다. An은 GPU 메모리로부터 CPU 메모리로 오프로드된 gradients를 필요로 한다. 모델 파라미터의 사이즈가 W이라면, 위의 과정에서 파라미터 이동 사이즈는 4*W이다.
Offloading에 있어서, PCIe는 전체 학습 프로세스의 병목점이다. 재계산으로 모든 매개변수를 오프로드할 때 GPU 메모리가 idle 상태인 것으로 관찰된다. 만일 GPU 메모리에 모델의 일부를 남겨두고 선택적 offloading을 통해 GPU 메모리를 채우는 방법을 사용하면 디바이스간 메모리 복사를 줄여 모델 학습을 가속화시킬 수 있다.
다음과 같은 실험 환경(48 layers, M6-78B on 8 NVIDIA-V100 GPUs)에서 CPU 메모리로 파라미터를 완전히 offloading 시키는데 걸리는 학습 step time은 89초이다. Granularly offloading를 적용하여 처음 24개 레이어는 CPU 메모리로 offloading하고 나머지 24개 레이어는 GPU 메모리에 남겨두면 학습 step time이 45초로 줄어든다.
따라서 Offloading의 granularity와 GPU의 util.은 extreme-scale pretraining에서 심각하게 고려되어야 한다. 추가적으로 모델이 극단적으로 크면, 전체 모델을 offloading은 CPU 메모리의 OOM 에러를 야기시킨다.
M6–10T는 Granular CPU offloading으로 512 NVIDIA-V100 GPUs 환경에서 학습 가능하며, 더나아가 Psuedo 단계에서 CPU 메모리로 offloading 없이 단지 256 GPUs 만으로 10T 계산 그래프로 Pseudo Giant를 학습할 수 있다. 그러므로 단지 256 GPUs로 Pseudo Giant를 학습하고, experts를 분할하여 512 GPUs로 재분배한다. 이러한 방법은 GPU 리소스 사용량을 절약하여 더 리소스 효율적이며 환경 친화적이다.
레퍼런스
[1] M6–10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
[2] M6-T: Exploring Sparse Expert Models and Beyond
[3] EFLOPS: Algorithm and System Co-design for a High Performance Distributed Training Platform
[4] Whale: Scaling Deep Learning Model Training to the Trillions
