[논문 리뷰] Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models

Yuliang Liu, Shenggui Li외 4명|arXiv (Cornell University)|2023. 02. 06.

Parallel Computing and Optimization Techniques인용 수 6

한 줄 요약

Colossal-Auto는 다중 GPU 시스템에서 대규모 모델 훈련을 위한 내부 연산자 병렬 처리와 활성화 체크포인팅을 종합적으로 최적화하는 통합적이고 자동화된 시스템을 소개한다. 하드웨어 인식 프로파일링과 이중 단계 솔버를 조합함으로써, 최소한의 코드 변경으로 근사 최적의 분산 실행 계획을 생성하며, 기존 방법에 비해 A100 GPU에서 뛰어난 약한 스케일링 성능을 달성한다.

ABSTRACT

In recent years, large-scale models have demonstrated state-of-the-art performance across various domains. However, training such models requires various techniques to address the problem of limited computing power and memory on devices such as GPUs. Some commonly used techniques include pipeline parallelism, tensor parallelism, and activation checkpointing. While existing works have focused on finding efficient distributed execution plans (Zheng et al. 2022) and activation checkpoint scheduling (Herrmann et al. 2019, Beaumont et al. 2021}, there has been no method proposed to optimize these two plans jointly. Moreover, ahead-of-time compilation relies heavily on accurate memory and computing overhead estimation, which is often time-consuming and misleading. Existing training systems and machine learning pipelines either physically execute each operand or estimate memory usage with a scaled input tensor. To address these challenges, we introduce a system that can jointly optimize distributed execution and gradient checkpointing plans. Additionally, we provide an easy-to-use symbolic profiler that generates memory and computing statistics for any PyTorch model with a minimal time cost. Our approach allows users to parallelize their model training on the given hardware with minimum code change based. The source code is publicly available at Colossal-AI GitHub or https://github.com/hpcaitech/ColossalAI

연구 동기 및 목표

• 대규모 모델을 위한 분산 훈련 전략을 수동으로 튜닝하는 데서 비롯하는 과제를 해결하기 위해, 깊은 전문 지식이 필요하고 실수가 쉬운 과정을 해결하고자 한다.
• 기존 시스템에서 별도로 다루는 경향이 있는 내부 연산자 병렬 처리와 활성화 체크포인팅을 종합적으로 최적화하고자 한다.
• 분산 훈련 경험 없이도 연구자와 실무자가 대규모 모델을 효율적으로 훈련할 수 있도록 접근 장벽을 낮추고자 한다.
• 복잡한 GPU 인터커넥트, 예를 들어 NVLink와 PCIe를 포함한 다양한 하드웨어 구성에 적응 가능한 하드웨어 인식 실행 계획 생성을 가능하게 하고자 한다.
• 임의의 차원의 장치 메시와 새로운 하드웨어 구성에 대응하는 유연하고 확장 가능한 병렬 처리 전략을 지원하고자 한다.

제안 방법

• 기호적 프로파일러는 계산 그래프에서 메타실행을 수행하여 최소한의 런타임 비용으로 세밀한 메모리 및 계산 오버헤드를 수집한다.
• 이중 단계 솔버는 먼저 최적의 내부 연산자 병렬 처리 전략을 식별한 후, 메모리 예산을 충족시키기 위해 활성화 체크포인팅을 종합적으로 스케줄링한다.
• 해밀토니안 텐서 레이아웃 변환 알고리즘은 임의의 N차원 장치 메시를 통해 효율적인 샤딩을 가능하게 하며, 복잡한 토폴로지를 지원한다.
• 시스템은 연산을 체크포인팅된, 병렬화된 함수로 대체함으로써 PyTorch 모델을 분산 실행 코드로 컴파일한다.
• 실행 계획 생성은 NVLink와 PCIe 대역폭의 차이를 고려한 하드웨어 토폴로지를 고려하여, 통신을 계산과 겹치도록 한다.
• 시스템은 PyTorch와 통합되며, 최소한의 사용자 코드 변경으로 자동으로 코드 재컴파일을 지원한다.

실험 결과

연구 질문

• RQ1 내부 연산자 병렬 처리와 활성화 체크포인팅을 종합적으로 최적화함으로써 대규모 모델 훈련에서 더 나은 메모리 및 계산 효율성을 달성할 수 있는가?
• RQ2 기호적 프로파일러는 낮은 오버헤드로 PyTorch 모델의 메모리 및 계산 비용을 정확하게 추정할 수 있는가?
• RQ3 통합 솔버는 NVLink와 PCIe와 같은 복잡한 GPU 인터커넥트 토폴로지에 적응 가능한 하드웨어 인식 실행 계획을 생성할 수 있는가?
• RQ4 병렬 처리와 체크포인팅의 종합 최적화는 순차적 또는 별도 최적화에 비해 훈련 성능과 메모리 사용 측면에서 어떻게 비교되는가?
• RQ5 수동 튜닝 없이도 다양한 모델 아키텍처와 하드웨어 구성에 대해 시스템이 얼마나 일반화될 수 있는가?

주요 결과

• Colossal-Auto는 평가된 모든 방법 중에서 가장 뛰어난 약한 스케일링 성능을 기록했으며, 14.55B 매개변수 모델에 대해 8개의 A100 GPU에서 0.824 PFLOPS에 도달했다.
• 모든 실험 설정에서 Megatron-LM, Optimus, 3D 텐서 병렬 처리보다 뛰어난 성능을 보였으며, 특히 더 큰 모델에서 두드러진 성능 향상을 보였다.
• 실행 계획은 인접한 GPU 간 고대역폭 NVLink 링크를 자동으로 활용하며, 저대역폭 PCIe 통신을 계산과 겹치도록 한다.
• 데이터 병렬 처리(DPD)와 달리, 대규모 스케일에서도 메모리 초과 오류를 피할 수 있다.
• 기호적 프로파일러는 계산 그래프 전반에 걸쳐 정확하고 저비용의 메모리 및 계산 사용량 프로파일링을 가능하게 하여 효과적인 최적화를 가능하게 한다.
• 병렬 처리와 체크포인팅의 종합 최적화는 각각 별도로 최적화하는 것보다 더 나은 메모리 활용도와 성능을 제공한다.