[논문 리뷰] DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models

Sam Adé Jacobs, Masahiro Tanaka외 5명|arXiv (Cornell University)|2023. 09. 25.

Topic Modeling인용 수 5

한 줄 요약

DeepSpeed-Ulysses는 입력 시퀀스를 GPU 간에 분할하고 효율적인 all-to-all 통신을 사용하여 초장기 시퀀스(최대 100만 토큰)를 가진 트랜스포머 모델의 훈련을 가능하게 하는 새로운 시퀀스 병렬 처리 방법을 제안한다. 이는 기존 시스템 대비 2.5배 빠른 훈련 속도와 10배 이상의 통신 감소를 달성하면서도 높은 처리량과 FlashAttention 및 ZeRO-3와의 호환성을 유지한다.

ABSTRACT

Computation in a typical Transformer-based large language model (LLM) can be characterized by batch size, hidden dimension, number of layers, and sequence length. Until now, system works for accelerating LLM training have focused on the first three dimensions: data parallelism for batch size, tensor parallelism for hidden size and pipeline parallelism for model depth or layers. These widely studied forms of parallelism are not targeted or optimized for long sequence Transformer models. Given practical application needs for long sequence LLM, renewed attentions are being drawn to sequence parallelism. However, existing works in sequence parallelism are constrained by memory-communication inefficiency, limiting their scalability to long sequence large models. In this work, we introduce DeepSpeed-Ulysses, a novel, portable and effective methodology for enabling highly efficient and scalable LLM training with extremely long sequence length. DeepSpeed-Ulysses at its core partitions input data along the sequence dimension and employs an efficient all-to-all collective communication for attention computation. Theoretical communication analysis shows that whereas other methods incur communication overhead as sequence length increases, DeepSpeed-Ulysses maintains constant communication volume when sequence length and compute devices are increased proportionally. Furthermore, experimental evaluations show that DeepSpeed-Ulysses trains 2.5x faster with 4x longer sequence length than the existing method SOTA baseline.

연구 동기 및 목표

• 기존 병렬 처리 기법(데이터, 텐서, 파이프라인 병렬 처리)이 트랜스포머 모델의 초장기 시퀀스 길이를 처리하는 데 가지는 확장성 한계를 해결하기 위해.
• 기존 시퀀스 병렬 처리 접근 방식에서 발생하는 메모리-통신 비효율성을 해결하여 장기 시퀀스에 대한 확장성을 향상시키기 위해.
• 다양한 하드웨어와 어텐션 메커니즘에서 장기 시퀀스 LLM의 효율적이고 확장 가능한 훈련을 가능하게 하는 일반적이고 이식 가능하며 코드 수정 최소화가 가능한 시스템 최적화를 설계하기 위해.
• 하드웨어 피크 성능의 54% 이상을 활용하고, 밀도 있는 어텐션과 희소 어텐션 모두를 지원하면서 성능 저하를 최소화하기 위해.

제안 방법

• 입력 시퀀스를 여러 GPU에 걸쳐 시퀀스 차원에 따라 분할하여 상호 겹치지 않는 시퀀스 세그먼트를 병렬로 처리할 수 있도록 한다.
• 어텐션 계산 전후에 all-to-all 통신 프리미티브를 사용하여 쿼리, 키, 밸류를 어텐션 헤드 간에 재분배하여 각 GPU가 고유한 헤드 세트에 대해 계산하도록 보장한다.
• 두 단계로 구성된 all-to-all 통신을 구현한다: 첫 번째 단계는 각 GPU가 할당된 어텐션 헤드에 맞는 전체 시퀀스를 수신할 수 있도록 데이터를 재구성하는 것이며, 두 번째 단계는 결과를 취합하면서 시퀀스 차원에 따라 재분할하는 것이다.
• 모델 병렬 처리와 메모리 최적화를 위해 ZeRO-3와 원활하게 통합되어 대규모 모델 크기와 장기 시퀀스 길이 모두를 지원한다.
• FlashAttention v2 및 희소 어텐션(예: 블록 기반 희소 어텐션)과 같은 효율적인 어텐션 메커니즘을 지원하며, 낮은 통신 오버헤드를 유지한다.
• 실행 환경에 종속되지 않는 설계를 하여 기존 훈련 프레임워크에 최소한의 코드 수정으로 통합할 수 있도록 한다.

Figure 1: Multi-head attention Transformer

실험 결과

연구 질문

• RQ1 장기 시퀀스 길이가 증가함에 따라, 시퀀스 병렬 처리가 확장성 있는 통신 효율성을 확보할 수 있는가?
• RQ2 시퀀스 길이와 장치 수가 증가함에 따라 통신 부담이 증가하는 상황에서, all-to-all 통신을 어떻게 활용하여 통신 부담을 일정하게 유지할 수 있는가?
• RQ3 DeepSpeed-Ulysses와 같은 시스템 수준 최적화가, 뚜렷한 코드 리팩터링 없이도 장기 시퀀스 훈련에서 높은 처리량과 강력한 확장성을 달성할 수 있는가?
• RQ4 기존의 시퀀스 병렬 처리 접근 방식(예: Megatron-LM)에 비해, 이 방법이 처리량, 메모리 효율성, 확장성 측면에서 얼마나 뛰어나게 성능을 발휘하는가?
• RQ5 ZeRO-3 및 FlashAttention와의 통합이 훈련 효율성과 모델 크기 확장성에 어떤 영향을 미치는가?

주요 결과

• DeepSpeed-Ulysses는 최대 100만 토큰의 시퀀스 길이를 가진 모델의 훈련을 가능하게 하여, 기존 시스템 대비 4배 긴 시퀀스 길이를 지원한다.
• 밀도 있는 어텐션을 사용하는 7B 및 30B 파라미터 모델에서 SOTA 기준(Megatron-LM) 대비 2.5배 높은 훈련 처리량을 달성한다.
• 시퀀스 길이와 GPU 수가 비례적으로 증가함에 따라 통신 부담이 일정하게 유지되며, 다른 방법들은 증가하는 통신 오버헤드로 인해 성능 저하를 겪는다.
• 실험 전반에서 GPU당 175 TFlops 이상을 유지하며, 하드웨어 피크 성능의 54% 이상을 달성한다.
• ZeRO-3 통합으로 인한 메모리 절감 덕분에 동일한 하드웨어 조건에서도 Megatron-LM 대비 4배 긴 시퀀스 길이를 지원할 수 있다.
• 강한 약한 확장성(weak scaling)을 보이며, 시퀀스 길이와 GPU 수가 비례적으로 증가하더라도 처리량이 높게 유지되며, 통신 오버헤드로 인한 성능 저하가 미미하다.

Figure 2: DeepSpeed sequence parallelism (DeepSpeed-Ulysses) design