[논문 리뷰] Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training

Can Karakus, Rahul Huilgol외 8명|arXiv (Cornell University)|2021. 11. 10.

Topic Modeling참고 문헌 28인용 수 8

한 줄 요약

이 논문은 PyTorch에서 자동 파이프라인 및 텐서 병렬 처리를 통해 효율적인 대규모 모델 훈련을 가능하게 하는 유연하고 일반적인 프레임워크인 Amazon SageMaker Model Parallelism을 제시한다. 코드 변경 없이 다양한 아키텍처—비-Transformer, 이종적, MoE 모델—을 지원하며, 네이티브 PyTorch 기능을 유지하면서 DeepSpeed와 경쟁 가능한 성능을 달성한다.

ABSTRACT

With deep learning models rapidly growing in size, systems-level solutions for large-model training are required. We present Amazon SageMaker model parallelism, a software library that integrates with PyTorch, and enables easy training of large models using model parallelism and other memory-saving features. In contrast to existing solutions, the implementation of the SageMaker library is much more generic and flexible, in that it can automatically partition and run pipeline parallelism over arbitrary model architectures with minimal code change, and also offers a general and extensible framework for tensor parallelism, which supports a wider range of use cases, and is modular enough to be easily applied to new training scripts. The library also preserves the native PyTorch user experience to a much larger degree, supporting module re-use and dynamic graphs, while giving the user full control over the details of the training step. We evaluate performance over GPT-3, RoBERTa, BERT, and neural collaborative filtering, and demonstrate competitive performance over existing solutions.

연구 동기 및 목표

• 다양한 모델 아키텍처와 사용 사례를 지원하는 일반적이고 유연한 대규모 모델 훈련 프레임워크의 부족을 해결하기 위해.
• 새로운 훈련 스크립트의 통합 부담을 줄이기 위해 자동 분할 및 최소한의 코드 변경을 가능하게 하기 위해.
• 동적 계산 그래프 및 모듈 재사용과 같은 네이티브 PyTorch 기능과의 완전한 호환성을 유지하기 위해.
• 다중 GPU 및 다중 노드 환경에서 파이프라인 및 텐서 병렬 처리에 대해 확장 가능하고 스케일러블한 시스템을 제공하기 위해.
• DeepSpeed와 같은 기존 솔루션과 비교해도 경쟁 가능한 훈련 성능을 제공하기 위해.

제안 방법

• 모듈-서버 아키텍처를 사용하여 임의의 모델 아키텍처를 장치 간에 자동으로 로드 밸런싱 분할할 수 있도록 한다.
• 균일한 대규모 모델, 고립된 큰 구성 요소를 가진 모델, 그리고 믹스처 오브 응용자(MoE) 아키텍처를 지원하는 일반적인 텐서 병렬 처리 프레임워크를 도입한다.
• 장치 간 및 노드 간에 동적으로 생성된 통신 요청을 처리하기 위한 전용 장치 간(D2D) 통신 백엔드를 제공한다.
• 원래 코드를 수정하지 않고 네이티브 PyTorch 모듈을 분산 대응 모듈로 매핑하기 위한 유연한 API와 후크(init_hook, forward_hook, return_hook)를 제공한다.
• 낮은 수준의 텐서 병렬 처리 제어를 위한 유틸리티로 parameter_creation_scope 및 분산 텐서 연산(e.g., allgather, allreduce, reduce-scatter)을 포함한다.
• 모델 배포 시까지 매개변수 초기화를 연기함으로써 CPU 메모리 압박을 줄이는 smp.delay_param_initialization라는 컨텍스트 매니저를 제공한다.

실험 결과

연구 질문

• RQ1 최소한의 코드 변경으로 임의의 모델 아키텍처를 지원하는 일반적인 모델 병렬 처리 프레임워크를 구축할 수 있는가?
• RQ2 텐서 병렬 처리를 비균일한 모델, 큰 임bedding 레이어, 믹스처 오브 응용자(MoE) 아키텍처에 일반화할 수 있는가?
• RQ3 동적 그래프 및 모듈 재사용과 같은 네이티브 PyTorch 기능이 분산 훈련 시스템에서 얼마나 잘 유지되는가?
• RQ4 실제 훈련 워크로드에서 제안된 프레임워크의 성능이 DeepSpeed와 같은 기존 솔루션과 비교해 어떻게 되는가?
• RQ5 자동 로드 밸런싱 분할이 복잡하고 이종적인 모델 아키텍처에 효과적으로 적용될 수 있는가?

주요 결과

• 프레임워크는 GPT-3, RoBERTa, BERT 및 신경 협업 필터링 모델을 DeepSpeed와 비교해 경쟁 가능한 훈련 스루풋으로 성공적으로 훈련시켰다.
• 수동적인 아키텍처 분석이나 재구성 없이도 자동 분할이 장치 간 효과적인 로드 밸런싱을 달성했다.
• 시스템은 PyTorch의 동적 계산 그래프와 조건부 실행 플로우, 매개변수 재사용을 모두 유지했다.
• D2D 통신 백엔드는 동적으로 생성된 통신 패턴을 효율적으로 처리하여 복잡한 훈련 워크플로우에서 오버헤드를 줄였다.
• smp.delay_param_initialization의 사용으로 CPU 메모리가 제한된 환경에서도 수십억 개의 매개변수를 가진 모델을 훈련시킬 수 있었다.
• 텐서 병렬 처리 프레임워크는 이전 솔루션보다 더 넓은 범위의 용도를 지원했으며, MoE 및 비-Transformer 아키텍처를 포함했다.