[논문 리뷰] TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

Zhuohan Li, Siyuan Zhuang외 5명|arXiv (Cornell University)|2021. 02. 16.

Topic Modeling참고 문헌 31인용 수 33

한 줄 요약

TeraPipe는 Transformer LMs에 대한 토큰 수준 파이프라인 병렬화를 도입하고, AWS에서의 이전 동기식 모델-병렬 방법 대비 GPT-3-175B에서 최대 5.0배의 학습 속도 향상을 달성한다.

ABSTRACT

Model parallelism has become a necessity for training modern large-scale deep language models. In this work, we identify a new and orthogonal dimension from existing model parallel approaches: it is possible to perform pipeline parallelism within a single training sequence for Transformer-based language models thanks to its autoregressive property. This enables a more fine-grained pipeline compared with previous work. With this key idea, we design TeraPipe, a high-performance token-level pipeline parallel algorithm for synchronous model-parallel training of Transformer-based language models. We develop a novel dynamic programming-based algorithm to calculate the optimal pipelining execution scheme given a specific model and cluster configuration. We show that TeraPipe can speed up the training by 5.0x for the largest GPT-3 model with 175 billion parameters on an AWS cluster with 48 p3.16xlarge instances compared with state-of-the-art model-parallel methods. The code for reproduction can be found at https://github.com/zhuohan123/terapipe

연구 동기 및 목표

• 단일 디바이스 메모리 한계를 넘는 매우 큰 Transformer LMs를 학습하기 위한 더 심도 있는 모델 병렬화의 필요성을 제시한다.
• 토큰 시퀀스에 따라 오토회귀 의존성을 활용하는 새로운 세밀한 파이프라인 차원을 식별한다.
• 최대 파이프라인 효율을 위한 최적 토큰 조각(partition) 분할을 계산하기 위해 동적 프로그래밍 기반 알고리즘을 개발한다.

제안 방법

• 단일 입력 시퀀스 내에서 토큰 차원을 가로지르며 파이프라인을 구성하는 토큰 수준 파이프라인 병렬화를 제안한다.
• 토큰 조각 크기와 클러스터 특성의 함수로 순전파/역전파 지연을 모델링한다.
• 학습 지연을 최소화하기 위해 토큰 차원에서 최적의 슬라이싱 방식을 찾는 동적 프로그래밍 알고리즘을 개발한다.
• 간단한 성능 모델로 순전파 시간을 추정하고 이를 통해 DP 최적화를 안내한다.
• 직교성 보여주기: TeraPipe는 기존의 데이터/모델 병렬 방법(microbatching, operation partitioning, data parallelism)과 직교적으로 결합될 수 있다.

실험 결과

연구 질문

• RQ1 자가회귀 Transformer에서 파이프라인 병렬화를 계층 차원에서 토큰 차원으로 확장하는 방법은 무엇인가?
• RQ2 주어진 LM과 클러스터에 대해 토큰 차원에서 어떤 슬라이싱 방식이 총 학습 지연을 최소화하는가?
• RQ3 토큰 수준 파이프라인이 다른 모델 병렬 기법 및 데이터 병렬성과 어떻게 상호 작용하는가?
• RQ4 토큰 수준 파이프라인 병렬화를 사용해 대형 GPT-3 규모 모델에서 어떤 성능 향상을 얻을 수 있는가?
• RQ5 시퀀스 길이가 토큰 수준 파이프라인 병렬화의 효과에 어떠한 영향을 미치는가?

주요 결과

• TeraPipe는 대형 LM에 대해 상당한 속도 향상을 제공하며, 48 AWS p3.16xlarge GPU에서 GPT-3-175B에 대해 기존의 동기식 모델-병렬 방법 대비 최대 5.0배 빠른 학습을 달성한다.
• 동적 프로그래밍 접근법은 파이프라인 효율성을 최대화하기 위한 최적 토큰 슬라이싱 방식을 효과적으로 결정하며, 조사된 경우에 균일 슬라이싱보다 약 1.04x–1.12x 우수하다.
• 메서드는 메모리 제약으로 배치 크기를 줄이고 파이프라인 단계를 늘리기 때문에 더 큰 모델에서 더 큰 이점을 제공하며, 토큰 수준 파이프라이닝이 더 많은 포화 기회를 제공한다.
• 더 긴 입력 시퀀스 길이는 토큰 수준 파이프라이닝의 잠재적 이점을 크게 높여 주며, 관찰된 속도 향상은 시퀀스 길이가 커질수록 증가한다.
• TeraPipe는 기존의 마이크로배치 기반 파이프라인, 연산 분할, 데이터 병렬성 등 기존 병렬 학습 방법과 직교하며 함께 사용할 수 있다.