[논문 리뷰] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Deepak Narayanan, Mohammad Shoeybi외 10명|arXiv (Cornell University)|2021. 04. 09.

Topic Modeling인용 수 7

한 줄 요약

이 논문은 메가트론-LM에서 파ip라인, 텐서, 데이터 병렬화의 새로운 조합인 PTD-P를 제안하여 3072개의 GPU에서 트리리온 파rameter 언어 모델의 효율적인 훈련을 가능하게 한다. 교차 배치 파이ipel라인 스케줄링을 도입함으로써 502 페타플롭스/초의 처리량과 이론적 최대 성능의 52%에 해당하는 GPU 활용률을 달성하여 처리량을 크게 향상시키면서도 메모리 효율성을 유지한다.

ABSTRACT

Large language models have led to state-of-the-art accuracies across a range of tasks. However, training these models efficiently is challenging for two reasons: a) GPU memory capacity is limited, making it impossible to fit large models on even a multi-GPU server, and b) the number of compute operations required to train these models can result in unrealistically long training times. Consequently, new methods of model parallelism such as tensor and pipeline parallelism have been proposed. Unfortunately, naive usage of these methods leads to fundamental scaling issues at thousands of GPUs, e.g., due to expensive cross-node communication or devices spending significant time waiting on other devices to make progress. In this paper, we show how different types of parallelism methods (tensor, pipeline, and data parallelism) can be composed to scale to thousands of GPUs and models with trillions of parameters. We survey techniques for pipeline parallelism and propose a novel interleaved pipeline parallelism schedule that can improve throughput by 10+% with memory footprint comparable to existing approaches. We quantitatively study the trade-offs between tensor, pipeline, and data parallelism, and provide intuition as to how to configure distributed training of a large model. Our approach allows us to perform training iterations on a model with 1 trillion parameters at 502 petaFLOP/s on 3072 GPUs with achieved per-GPU throughput of 52% of theoretical peak. Our code is open sourced at https://github.com/nvidia/megatron-lm.

연구 동기 및 목표

• GPU 메모리 용량을 초과하거나 훈련 시간이 지나치게 오래 걸리는 대규모 언어 모델을 훈련하는 데 도전하는 것.
• 수천 개의 GPU를 통해 1조 개 이상의 파라미터를 가진 모델의 확장 가능한 훈련을 가능하게 하는 것.
• 파이프라인, 텐서, 데이터 병렬화의 조합을 최적화하면서 엄격한 옵티마이저 의미 체계를 유지하는 것.
• 새로운 스케줄링 전략을 통해 파이프라인 플러시 오버헤드를 줄이고 소규모 배치 크기에서 처리량을 향상시키는 것.
• 대규모 GPU 클러스터에서 높은 계산 효율성(이론적 최대 성능의 52%)을 달성하는 것.

제안 방법

• 노드 간 파이프라인 병렬화, 노드 내 텐서 병렬화, 데이터 병렬화를 조합한 PTD-P를 사용하여 수천 개의 GPU에 걸쳐 훈련을 확장하는 것.
• 파이프라인 플러시 시간을 줄이고 기준 방법 대비 10% 이상 처리량을 향상시키기 위해 교차 배치 파이프라인 스케줄링을 도입하는 것.
• 활성화 재계산을 적용하여 메모리 사용량을 줄이고, 더 큰 모델이 GPU 메모리 제약 조건에 맞출 수 있도록 하는 것.
• 연산자 융합과 최적화된 데이터 레이아웃을 적용하여 메모리 제약이 큰 커널과 통신 오버헤드를 최소화하는 것.
• 모델 장치 간 분할을 안내하기 위해 비용 모델을 활용하여 통신을 최소화하면서도 장치 활용률을 극대화하는 것.
• 서버 내 고대역폭 인터커넥트(NVLink 등)와 최적화된 노드 간 통신을 활용하여 다중 노드 훈련에서 지연 시간을 줄이는 것.

실험 결과

연구 질문

• RQ1 파이프라인, 텐서, 데이터 병렬화를 어떻게 효과적으로 조합하여 수천 개의 GPU에 걸쳐 훈련을 확장할 수 있는가?
• RQ2 소규모 배치 크기에서 파이프라인 플러시 오버헤드를 최소화하고 처리량을 극대화하는 데 가장 효과적인 스케줄링 전략은 무엇인가?
• RQ3 GPU 클러스터에서 트리리온 파라미터 모델을 훈련할 때 메모리 제약 조건을 어떻게 극복할 수 있는가?
• RQ4 활성화 재계산과 통신 최적화가 훈련 효율성에 미치는 영향는 무엇인가?
• RQ5 제안된 방법이 대규모 GPU 클러스터에서 높은 활용도(예: 최대 성능의 52%)를 얼마나 달성할 수 있는가?

주요 결과

• 제안된 교차 배치 파이프라인 스케줄링은 표준 파이프라인 접근 방식 대비 훈련 처리량을 10% 이상 향상시킨다.
• 3072개의 GPU에서 최대 훈련 처리량 502 페타플롭스/초를 달성하여 이론적 GPU 최대 성능의 52%에 도달한다.
• 제안된 PTD-P 프레임워크를 사용하면 1조 파라미터 모델의 훈련이 약 3개월 내에 가능하다.
• 텐서, 파이프라인, 데이터 병렬화의 조합은 단일 또는 다중 GPU 서버의 메모리 및 계산 한계를 초월한 확장 가능한 훈련을 가능하게 한다.
• 활성화 재계산은 FLOP 수를 크게 증가시키지 않으면서도 메모리 사용량을 줄여 더 큰 모델 훈련을 가능하게 한다.
• 파이프라인 동기화로 인한 유휴 시간을 최소화하면서도 엄격한 옵티마이저 의미 체계를 유지한다.