[논문 리뷰] Sequence Parallelism: Making 4D Parallelism Possible

Shenggui Li, Fuzhao Xue외 2명|arXiv (Cornell University)|2021. 05. 26.

Advanced Neural Network Applications참고 문헌 23인용 수 12

한 줄 요약

이 논문은 시퀀스 병렬 처리를 소개한다. 이는 장기 시퀀스를 다중 GPU에 나누어 저장함으로써 메모리 효율적인 학습 방법으로, 주로 린 스타일의 전역 감소를 통해 어텐션 임베딩을 통신한다. PyTorch에서 데이터 병렬 및 파이프라인 병렬과의 완전한 호환성을 유지하면서, 64개의 P100 GPU에서 텐서 병렬보다 13.7배 큰 배치 크기와 3.0배 긴 시퀀스 길이를 달성한다.

ABSTRACT

Within Transformer, self-attention is the key module to learn powerful context-aware representations. However, self-attention suffers from quadratic memory requirements with respect to the sequence length, which limits us to process longer sequence on GPU. In this work, we propose sequence parallelism, a memory efficient parallelism method to help us break input sequence length limitation and train with longer sequence on GPUs. Compared with existing parallelism, our approach no longer requires a single device to hold the whole sequence. Specifically, we split the input sequence into multiple chunks and feed each chunk into its corresponding device (i.e. GPU). To compute the attention output, we communicate attention embeddings among GPUs. Inspired by ring all-reduce, we integrated ring-style communication with self-attention calculation and proposed Ring Self-Attention (RSA). Our implementation is fully based on PyTorch. Without extra compiler or library changes, our approach is compatible with data parallelism and pipeline parallelism. Experiments show that sequence parallelism performs well when scaling with batch size and sequence length. Compared with tensor parallelism, our approach achieved $13.7 imes$ and $3.0 imes$ maximum batch size and sequence length respectively when scaling up to 64 NVIDIA P100 GPUs. We plan to integrate our sequence parallelism with data, pipeline and tensor parallelism to further train large-scale models with 4D parallelism in our future work.

연구 동기 및 목표

• Transformer의 자기 어텐션에서 발생하는 제곱형 메모리 증가 문제를 해결하여 GPU에서의 시퀀스 길이 제한을 완화한다.
• 자기 어텐션 계산 중 전체 입력 시퀀스를 하나의 장치가 보유해야 하는 문제를 해결한다.
• 모델 또는 하드웨어 수정 없이 다중 GPU에 걸쳐 시퀀스 조각을 분산하여 장기 시퀀스 학습을 효율적으로 가능하게 한다.
• 기존의 병렬 처리 전략(예: 데이터 병렬 및 파이프라인 병렬)과 호환되는 통신 효율적인 어텐션 메커니즘을 설계한다.
• 4D 병렬 처리 프레임워크에 시퀀스 병렬 처리를 통합하여 배치 크기와 시퀀스 길이 양쪽 모두에서 확장 가능한 학습을 달성한다.

제안 방법

• 입력 시퀀스를 여러 조각으로 나누고 각 조각을 별도의 GPU에 할당하여, 어느 한 장치가 전체 시퀀스를 저장할 필요가 없도록 한다.
• 자기 어텐션 계산 중에 GPU 간 어텐션 임베딩을 교환하기 위해 린 스타일의 전역 감소 통신을 사용하여 분산 어텐션 계산을 가능하게 한다.
• 린 통신을 어텐션 계산 파이프라인에 직접 통합한 새로운 어텐션 메커니즘인 린 자기 어텐션(Ring Self-Attention, RSA)을 도입한다.
• 컴파일러나 라이브러리 변경 없이 PyTorch 내부에서 직접 구현하여 기존 학습 프레임워크와 광범위한 호환성을 확보한다.
• 데이터 병렬 및 파이프라인 병렬과 결합하여 대규모 모델 학습을 위한 4D 병렬 처리를 가능하게 한다.
• 역전파 단계에서 린 전역 감소를 사용해 장치 간 어텐션 출력을 동기화함으로써 기울기 일致성과 모델 정확도를 유지한다.

실험 결과

연구 질문

• RQ1 시퀀스 병렬 처리는 다중 GPU에 걸쳐 시퀀스 조각을 분산함으로써 Transformer의 시퀀스 길이 제한을 극복할 수 있는가?
• RQ2 텐서 병렬과 비교했을 때 시퀀스 병렬 처리는 최대 배치 크기와 시퀀스 길이 스케일링에서 어떤 성능을 보이는가?
• RQ3 데이터 병렬 및 파이프라인 병렬과 결합했을 때 시퀀스 병렬 처리가 배치 크기와 시퀀스 길이 양쪽 모두에서 얼마나 잘 확장되는가?
• RQ4 아키텍처나 컴파일러 수준의 변경 없이도 기존의 병렬 처리 전략(예: 데이터 병렬 및 파이프라인 병렬)과 시퀀스 병렬 처리를 통합할 수 있는가?
• RQ5 표준 전역 감소나 파arameter 공유와 비교했을 때, 시퀀스 병렬 처리에서 린 스타일의 통신은 어텐션 계산에 얼마나 효율적인가?

주요 결과

• 64개의 NVIDIA P100 GPU로 확장했을 때, 텐서 병렬 대비 13.7배 큰 배치 크기로 학습이 가능하다.
• 동일한 하드웨어 환경에서 텐서 병렬 대비 3.0배 긴 입력 시퀀스를 지원한다.
• 배치 크기와 시퀀스 길이 양쪽 모두에서 높은 확장성을 보이며, 강력한 선형 스케일링 효율성을 입증한다.
• 데이터 병렬 및 파이프라인 병렬과 완전히 호환되어 기존 학습 파이프라인에 원활하게 통합된다.
• 링 전역 감소를 활용함으로써 린 자기 어텐션(RSA)이 통신 오버헤드를 효과적으로 줄여 학습 중 높은 처리량을 유지한다.
• PyTorch나 외부 라이브러리의 수정 없이도 구현 가능하여 기존 딥러닝 워크플로우에 쉽게 배포할 수 있다.