[논문 리뷰] Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training

Li, Shenggui, Hongxin Liu외 6명|arXiv (Cornell University)|2021. 10. 28.

Topic Modeling참고 문헌 15인용 수 21

한 줄 요약

Colossal-AI는 데이터, 파이프라인, 텐서, 시퀀스 병렬 처리를 통합 지원함으로써 효율적인 대규모 분산 학습을 가능하게 하는 통합된 딥러닝 시스템입니다. 동적 메모리 관리와 사용자 친화적인 병렬 처리 기능을 통해 순차 모델의 유연한 병렬 처리를 가능하게 하며, 기준 시스템 대비 최대 2.76배의 학습 속도 향상을 달성합니다.

ABSTRACT

The success of Transformer models has pushed the deep learning model scale to billions of parameters. Due to the limited memory resource of a single GPU, However, the best practice for choosing the optimal parallel strategy is still lacking, since it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism, as well as heterogeneous training methods integrated with zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.

연구 동기 및 목표

• 1. 수십억 파라미터 모델의 대규모 분산 학습을 위한 통합적이고 사용자 친화적인 시스템의 부족을 해결하기 위해.
• 2. 연구자와 엔지니어가 최적의 병렬 처리 전략을 선택하는 데 필요한 전문 지식의 장벽을 낮추기 위해.
• 3. 시퀀스 병렬 처리와 다차원 텐서 병렬 처리를 포함한 포괄적인 병렬 처리 방법 세트를 지원하기 위해.
• 4. 동적 스키밍과 오프로딩을 통해 메모리 효율성과 학습 처리량을 향상시키기 위해.
• 5. 아키텍처 재작성 없이 순차 모델 코드를 다중 GPU 클러스터로 원활하게 확장할 수 있도록 하기 위해.

제안 방법

• 1. 사용자가 순차 모델 코드를 최소한의 코드 변경으로 분산 환경으로 확장할 수 있도록 통합된 API를 제공합니다.
• 2. 데이터, 파이프라인, 텐서, 시퀀스 병렬 처리를 통합하여 최적의 성능을 위한 다양한 조합을 가능하게 합니다.
• 3. 활성화 메모리 병목 현상을 줄이고 확장성을 향상시키기 위해 다차원 텐서 병렬 처리를 적용합니다.
• 4. 시퀀스 병렬 처리는 입력 및 활성화 텐서를 시퀀스 차원에 따라 분할하여 통신 및 메모리 오버헤드를 감소시킵니다.
• 5. 동적 텐서 배치를 구현하여 필요할 때만 오프로딩하여 GPU 메모리 활용도를 향상시킵니다.
• 6. ZeRO-3와 이종 학습을 지원하여 중복된 최적화기 상태를 최소화하고 혼합 하드웨어 환경에서의 학습을 가속화합니다.

실험 결과

연구 질문

• RQ1 1. 통합된 딥러닝 시스템은 최적의 학습 성능을 위해 다수의 병렬 처리 기법을 어떻게 효과적으로 조합할 수 있는가?
• RQ2 2. 1D 텐서 병렬 처리와 비교했을 때 시퀀스 병렬 처리는 메모리 효율성과 확장성에 어떤 영향을 미치는가?
• RQ3 3. 낮은 배치 크기 환경에서 동적 스키밍과 오프로딩은 처리량과 메모리 활용도를 향상시킬 수 있는가?
• RQ4 4. Megatron-LM 및 DeepSpeed와 같은 최첨단 시스템과 비교했을 때 Colossal-AI의 학습 속도 향상은 어떠한가?
• RQ5 5. 파이프라인 스테이지 수와 GPU 수가 증가함에 따라 Colossal-AI는 학습 처리량을 얼마나 잘 확장할 수 있는가?

주요 결과

• 1. Colossal-AI는 대규모 모델에서 기준 시스템 대비 최대 2.76배의 학습 속도 향상을 달성합니다.
• 2. 12개 GPU에서 시퀀스 병렬 처리는 1D 텐서 병렬 처리 대비 4.44배 큰 배치 크기와 1.18배 긴 시퀀스 길이를 가능하게 합니다.
• 3. 4개의 파이프라인 스테이지에서 시퀀스 병렬 처리는 1D 텐서 병렬 처리 대비 1.55배 더 빠른 학습을 제공합니다.
• 4. 낮은 배치 크기 환경에서 Colossal-AI의 동적 오프로딩은 불필요한 CPU 전송을 방지하여 DeepSpeed의 정적 정책을 초월하는 처리량 향상을 달성합니다.
• 5. 8개 GPU와 13B 파라미터를 가진 OPT 모델, 배치 크기 32에서 Colossal-AI는 DeepSpeed 대비 1.33배의 속도 향상을 달성합니다.
• 6. 선형 복잡도 어텐션을 사용할 경우 GPU 수에 비례해 최대 시퀀스 길이를 선형 확장할 수 있어 문서 수준의 텍스트 이해가 가능해집니다.