[논문 리뷰] PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers

Chaoyang He, Shen Li외 2명|arXiv (Cornell University)|2021. 02. 05.

Neural Networks and Applications참고 문헌 22인용 수 16

한 줄 요약

PipeTransformer는 딥러닝 모델의 분산 학습을 위한 자동화되고 탄력적인 파ip라인 프레임워크를 제안하며, 학습 중에 층을 동적으로 동결하고 활성 층을 적은 수의 GPU에 압축함으로써 데이터 병렬 폭을 늘리는 방식을 통해 ViT 및 BERT 모델에서 정확도 손실 없이 최대 2.83배의 속도 향상을 달성한다.

ABSTRACT

The size of Transformer models is growing at an unprecedented pace. It has only taken less than one year to reach trillion-level parameters after the release of GPT-3 (175B). Training such models requires both substantial engineering efforts and enormous computing resources, which are luxuries most research teams cannot afford. In this paper, we propose PipeTransformer, which leverages automated and elastic pipelining and data parallelism for efficient distributed training of Transformer models. PipeTransformer automatically adjusts the pipelining and data parallelism by identifying and freezing some layers during the training, and instead allocates resources for training of the remaining active layers. More specifically, PipeTransformer dynamically excludes converged layers from the pipeline, packs active layers into fewer GPUs, and forks more replicas to increase data-parallel width. We evaluate PipeTransformer using Vision Transformer (ViT) on ImageNet and BERT on GLUE and SQuAD datasets. Our results show that PipeTransformer attains a 2.4 fold speedup compared to the state-of-the-art baseline. We also provide various performance analyses for a more comprehensive understanding of our algorithmic and system-wise design. We also develop open-sourced flexible APIs for PipeTransformer, which offer a clean separation among the freeze algorithm, model definitions, and training accelerations, hence allowing it to be applied to other algorithms that require similar freezing strategies.

연구 동기 및 목표

• 대규모 Transformer 모델을 학습하는 데 소요되는 높은 계산 비용과 공학적 비용을 해결하기 위해.
• 깊이 있는 네트워크에서 층들이 하향식으로 수렴한다는 관찰을 바탕으로 학습 시간과 자원 사용량을 줄이기 위해.
• 동결된 층에 기반해 파이프라인 구성 및 데이터 병렬 복제를 동적으로 조정할 수 있는 시스템을 설계하기 위해.
• 실시간 동결, 동적 파이프라인, 변화하는 프로세스 그룹 내 집합적 통신, 그리고 복제 간 활성화 캐시 공유 등의 과제를 극복하기 위해.

제안 방법

• 학습 중 활성화 안정성에 기반해 층을 식별하고 동결하기 위해 조정 가능한 임계값 α를 사용하는 적응형 동결 알고리즘을 도입한다.
• AutoPipe를 활용해 활성 층을 적은 수의 GPU에 압축함으로써 통신 비용과 파이프라인 빈도를 최소화한다.
• 해당 GPU에 할당된 자원을 기반으로 동적으로 새로운 파이프라인 복제본을 생성하는 AutoDP를 사용하며, 효율적인 집합적 연산을 위해 계층적 통신 그룹을 유지한다.
• 기존 및 새로운 파이프라인 복제본 간에 캐시된 활성화를 공유하기 위해 AutoCache를 구현하며, 프로파일러를 통해 최적의 활성화 캐시 시작 시점을 결정한다.
• 이러한 구성 요소들을 통합하여 실시간 층 동결에 기반해 모델 분할 및 복제를 자동으로 조정하는 종합적인 시스템을 설계한다.
• PyTorch의 파이프라인 및 데이터 병렬 API를 활용하여 호환성과 효율적 실행을 보장한다.

실험 결과

연구 질문

• RQ1 학습 중 동적 층 동결 전략을 효과적으로 활용하여 Transformer의 분산 학습을 가속화할 수 있는가?
• RQ2 층 동결에 대응해 파이프라인과 데이터 병렬성을 공동 최적화할 수 있는가?
• RQ3 탄력적 파이프라인 환경에서 동결 임계값 α와 마이크로배치 청크 분할 수의 성능 상충 관계는 어떠한가?
• RQ4 변동하는 파이프라인 복제 수를 가진 환경에서 집합적 통신을 효율적으로 관리할 수 있는가?
• RQ5 동적 학습 환경에서 활성화 캐시를 최적의 시점에 활성화할 수 있는가?

주요 결과

• PipeTransformer는 ViT(ImageNet) 및 BERT(SQuAD/GLUE)에서 최신 기준 대비 최대 2.83배의 속도 향상을 달성했으며, 정확도 저하 없이 성능을 유지한다.
• 더 큰 동결 임계값 α(예: α = 1/5)는 더 큰 속도 향상을 가져오지만 약간의 정확도 손실을 수반함으로써 속도와 성능 사이의 조정 가능한 트레이드오프를 보여준다.
• 최적의 마이크로배치 수 M은 파이프라인 길이 K에 따라 크게 달라지므로, 성능 저하를 방지하기 위해 사전 프로파일링가 필요하다.
• 활성화 캐시가 유의미한 이득을 보이기 위해서는 일정 수 이상의 층이 동결되어야 하며, ViT의 경우 3개, BERT의 경우 5개의 층이 동결된 이후 AutoCache가 캐시를 시작한다.
• ImageNet 기준 총 학습 시간의 5.9%, SQuAD 기준 8.8%가 통신 비용을 차지함을 확인하여, 모델 크기를 추가로 줄이면 효율성이 더욱 향상됨을 시사한다.
• 감소된 파이프라인 빈도와 증가된 데이터 병렬 폭의 상호작용으로 인해 상당한 비가산적 속도 향상이 발생한다.