[논문 리뷰] Alpa: Automationg Inter-and Intra-Operator Parallelism for Distrivuted Deep Learning

 

Paper: Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning, Yonghao Zhuang, Zhuohan Li, Hao Zhang et al. , OSDI 2022

 

1.  Background

 딥러닝 모델의 규모는 최근 몇 년간 폭발적으로 증가하여, GPT-3와 같은 모델은 1,750억 개의 파라미터를 가지고 있고 350GB 이상의 GPU메모리가 필요합니다. 이러한 거대 모델들은 단일 디바이스의 메모리에 담을 수 없어, 효율적인 분산 학습 방법이 필수적입니다. 이러한 이유로 다양한 Parallelism방식이 제안되어 왔습니다.

 

2.  Inference

 분산 학습에는 모델 구조를 복제하고 데이터를 배치 단위로 나누어 처리하는 Data Parallelism 방식과, 모델을 여러 GPU에 분할하여 각 GPU가 모델의 일부만을 처리하는 Model Parallelism 방식이 있습니다. 논문에서는 Model Parallelism의 Inter-Operator Parallelism과 Intra-Operator Parallelism이라는 두 가지 방식으로 문제를 해결합니다. Inter-Operator Parallelism은 모델의 레이어를 여러 GPU에 나누어 할당하는 것으로, 레이어 단위로 작업을 분할하는 방식입니다. Intra-Operator Parallelism은 개별 연산자나 레이어 내부의 연산을 여러 GPU에 분할하는 방식입니다. 두 방식은 계산 속도를 향상시킬 수 있지만, 디바이스 간 통신 오버헤드와 복잡한 동기화로 인해 병렬화 효율이 떨어지는 문제가 발생합니다. Alpa는 계산 그래프를 분석하여 각 연산자나 레이어를 어떻게 분할하고 어느 디바이스에 배치할지를 자동으로 결정합니다. 이를 통해 다양한 모델과 하드웨어 환경에서 효율적인 분산 학습을 지원하며, 개발자가 복잡한 병렬화 전략을 수동으로 설계하지 않아도 된다고 설명합니다.

 

3.  Main Stretgy

 Hierarchical Parallelism Strategy는 Alpa에서 병렬화 공간을 체계적으로 탐색하는 전략입니다. Hierarchical Parallelism Strategy와 Hierarchical Parallelism Strategy를 실행하는 런타임 아키텍처가 Alpa의 핵심 기술이라고 판단되어 다음과 같이 두 방식을 설명하겠습니다.

 

3-1. Hierarchical Parallelism Strategy

 

Alpa는 Inter-Operator Parallelism와 Intra-Operator Parallelism를 계층적으로 결합하여 광범위한 병렬화 공간에서 최적의 Strategy를 탐색합니다. 전통적인 Parallelism Strategy는 특정 모델이나 클러스터 환경에 맞춰 수작업으로 설계되는데, 이러한 고민을 알고리즘으로서 해결했습니다.
그림 1에서 (a), (b), (c), (d)는 연산자와 계산 그래프를 단순히 연속적인 방식으로 operator를 분할하여 GPU에 할당하는 구조입니다. 하지만 논문에 적용된 방식인 (e)는 Intra-operator와 Inter-operator을 동적으로 최적의 Parallelism으로 탐색합니다. Stage Partitioning 단계에서는 Dynamic-Programming(DP)을 통해 계산 그래프를 Stage로 나누고, 디바이스에 적절히 할당하는 Stage-Device Mapping 방식을 자동으로 탐색합니다. Operator Sharing 단계에서는 Integer Linear Programming(ILP)을 통해 연산자를 세분화하고, GPU 간의 부하 균형과 통신 비용을 최소화하는 방향으로 작업을 분배합니다. 이 두 단계는 서로 보완적으로 동작합니다. Stage Partitioning은 계산 그래프의 상위 구조에서 작업을 분리하며, GPU 간 작업 부하를 큰 단위로 조정하여 전반적인 실행 시간을 최적화합니다. 반면, Operator Shading은 각 스테이지 내부의 세부 연산자 단위까지 조정을 수행하여 GPU 자원의 활용도를 극대화 하고 통신 병목을 줄입니다. 결과적으로 그림 1의 (e)방식은 (a), (b), (c), (d)에서 볼 수 있는 각각의 단점을 보완하여 계산효율성을 최적화하는 병렬화 Strategy를 구현합니다.

3-2.  Runtime Architecture

Parallelism Strategy의 런타임은 MPMD(Multiple Program Multiple Data) 아키텍처를 기반으로 동작합니다.
해당 아키텍처는 각 GPU가 독립적인 작업을 수행할 수 있도록 관리합니다. 특히, 대규모 학습 환경에서 효율적인 데이터 이동을 보장하기 위해 Cross-Mesh Resharding 기술을 이용해 Mesh간 데이터를 재배치합니다. 이 과정에서 All-Reduce는 Gradient를 동기화하는데 사용되며, 이를 통해 모델 업데이트의 일관성을 유지합니다. 또한, All-Gather는 Mesh간의 데이터를 수집하고, 이를 모든 GPU에 복제하여 다음 연산을 준비하는데 사용됩니다. 이 과정을 통해 대규모 모델에서 Mesh 간의 데이터 이동을 최적화합니다. 그림 2에서 (a)는 Megatron-LM 방식은 Mesh 간 데이터를 직접 전송하는 구조를 보여줍니다.이 구조는 구현은 간단하지만 통신 경로가 비효율적이며 네트워크 부하가 증가하는 문제가 있습니다. (b)는 GPU 간 직접 데이터를 송수신을 수행하는 방식으로, 타일 데이터를 세밀하게 처리할 수 있으나, 네트워크 트래픽이 급격히 증가합니다. Alpa에 적용된 (c)는 Mesh 내부에서 데이터를 먼저 집계한 후, 이를 Mesh 간 효율적으로 전송하는 방식으로 통신 오버헤드를 최소화합니다.

 

 

4. Results

 

 Alpa와 타 방법론과의 성능을 비교하는 그림 3을 살펴보면 MoE와 Wide-ResNet 모델에서 Alpa를 적용한 결과, 다른 방법론 대비 압도적인 Throughput을 보입니다. GPU가 많아질수록 기존의 방법론들은 네트워크 병목과 부하 불균형으로 인해 성능이 정체되거나 감소하는 모습을 보이는 반면, Alpa는 Linear Scaling에 가까운 성능을 유지하며, GPU 추가에 따라 지속적으로 Throughput이 증가하는 모습을 보여줍니다. 다만 GPT에서는 Alpa와 Megatron-LM의 성능 차이가 상대적으로 크지 않다는 점이 눈에 띕니다. GPU 수가 64개에 도달했을 때, 두 방법론의 Throughput은 약 3.8 PFLOPS로 비슷한 수준을 보입니다.

 

5. Analysis and Improvement Suggestions

 최종적으로 GPT 모델에서 Alpa와 Megatron-LM의 Throughput 성능이 비슷하게 나타난 이유를 분석해보고자 합니다. GPT모델은 Transformer의 Decoder Only 구조를 채택합니다. 이 구조는 레이어간 의존성이 상대적으로 단순하며, 순차적으로 연산을 진행하는 과정에서도 Inter-operator parallelism만으로 충분히 높은 Parallelism 효율을 발휘할 수 있습니다. 논문의 의도는 Alpa가 단순한 모델 환경에서도 기존 방법론과 유사한 수준의 성능을 제공하면서, 더 복잡한 구조에서는 더 뛰어난 성능 향상을 발휘하고 다양한 딥러닝 모델에서의 학습이 안정적이라는 것을 보여주는 것을 의도한 듯 합니다. 그렇지만 최근 트렌드의 GPT-4, Llama, Claude, Mistral 등의 LLM의 모델 구조는 모두 Transformer Decoder Only 구조라는 점에서 단순한 구조를 가진 모델에서도 Alpa의 기술적 개선 가능성을 높이고, 통신 효율을 더욱 높이기 위한 방향성을 제시가 필요해 보입니다.

 

6. Reference

유튜브 Neural_Acceleration_Study 채널의 윤유경님 논문 리뷰