[논문 리뷰] Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads

Siyu Wang, Yi Rong외 7명|arXiv (Cornell University)|2020. 07. 08.

Parallel Computing and Optimization Techniques참고 문헌 31인용 수 5

한 줄 요약

Auto-MAP는 DQN 기반 강화학습 프레임워크를 제안하며, HLO IR에서 작동함으로써 DNN 워크로드에 대한 최적의 분산 실행 계획을 자동으로 탐색한다. 이는 데이터 병렬성, 모델 병렬성, 파이프라인 병렬성을 통합적으로 탐색할 수 있게 하여, 두 시간 이내에 NLP 및 CNN 모델에서 최적의 처리량을 달성한다. 수작업 및 히ュ리스틱 방법보다 뛰어난 성능을 보인다.

ABSTRACT

The last decade has witnessed growth in the computational requirements for training deep neural networks. Current approaches (e.g., data/model parallelism, pipeline parallelism) parallelize training tasks onto multiple devices. However, these approaches always rely on specific deep learning frameworks and requires elaborate manual design, which make it difficult to maintain and share between different type of models. In this paper, we propose Auto-MAP, a framework for exploring distributed execution plans for DNN workloads, which can automatically discovering fast parallelization strategies through reinforcement learning on IR level of deep learning models. Efficient exploration remains a major challenge for reinforcement learning. We leverage DQN with task-specific pruning strategies to help efficiently explore the search space including optimized strategies. Our evaluation shows that Auto-MAP can find the optimal solution in two hours, while achieving better throughput on several NLP and convolution models.

연구 동기 및 목표

• 다양한 DNN 워크로드에 대해 효율적인 분산 훈련 전략을 수작업으로 설계하는 도전 과제를 해결하기 위해.
• 단일 프레임워크 내에서 데이터, 모델, 파이프라인 병렬성을 통합적으로 탐색하여 프레임워크 또는 모델에 특화된 솔루션을 피하기 위해.
• 저수준 분산 시스템 설정을 제거함으로써 사용자 부담을 줄이고, 프레임워크에 종속되지 않고 워크로드 중립적인 병렬화 계획을 가능하게 하기 위해.
• 작업에 특화된 프루닝과 DQN 기반 강화학습을 통해 대규모이고 복잡한 병렬성 공간에서의 탐색 효율성을 향상시키기 위해.
• BERT, T5, ResNet과 같은 다양한 모델에서 자동으로 고성능이며 일반화 가능한 병렬화 전략을 제공하기 위해.

제안 방법

• Auto-MAP는 XLA에서 생성되는 저수준이고 하드웨어에 종속되지 않은 중간 표현인 HLO IR에서 작동하여, 운영자 또는 레이어 수준의 표현보다 더 세밀하고 표현력 있는 모델 분석을 가능하게 한다.
• 데이터 병렬성, 모델 병렬성, 파이프라인 병렬성을 포함한 광범위한 분산 실행 계획의 탐색을 위해 작업에 특화된 프루닝 전략을 적용한 딥 퀀트 네트워크(DQN)를 활용한다.
• 훈련 처리량을 기반으로 한 보상 함수를 사용하여 DQN 에이전트가 최적 또는 근사 최적의 병렬화 전략으로 유도한다.
• XLA 컴파일러 스택을 활용하여 TensorFlow, JAX, Trax 등의 다양한 프론트엔드와의 호환성을 확보하며, 사용자 모델에 대한 코드 침입을 방지한다.
• 실제 모델 분포에서의 추론을 위해, DQN 에이전트는 먼저 무작위 분포를 가진 합성 환경에서 훈련된 후, 실제 모델 분포에서 미세조정된다.
• NLP 모델의 경우, 에이전트는 사전 훈련 후에 잘 일반화되며, CNN의 경우 정확한 파이프라인 분할을 달성하기 위해 모델 전용 분포에서의 미세조정이 필요하다.

실험 결과

연구 질문

• RQ1 수작업 설계 없이 다양한 DNN 모델에서 최적의 분산 훈련 전략을 자동으로 탐색할 수 있는가?
• RQ2 단일 확장 가능한 탐색 공간 내에서 데이터, 모델, 파이프라인 병렬화 전략을 동시에 탐색할 수 있는 통합 프레임워크가 가능한가?
• RQ3 HLO IR에서 작동함으로써 레이어 또는 운영자 수준의 추상화보다 더 나은 성능과 일반화 능력을 확보할 수 있는가?
• RQ4 DQN에서의 작업에 특화된 프루닝이 탐색 공간을 크게 줄이고 최적 해에의 수렴 속도를 빠르게 할 수 있는가?
• RQ5 사전 훈련된 DQN 에이전트가 다양한 모델 아키텍처 간에 얼마나 잘 일반화되는가? 그리고 언제 미세조정이 필요한가?

주요 결과

• Auto-MAP는 NLP 및 컨volutional 모델 모두에서 두 시간 이내에 최적의 분산 실행 계획을 탐색하여 수작업 튜닝 시간을 크게 줄였다.
• BERT 및 T5 모델의 경우, 사전 훈련된 합성 데이터에서 DQN 에이전트가 미세조정 없이도 최적의 파이프라인 분할 전략을 성공적으로 유추하였다.
• T5-11B와 같은 더 큰 모델의 경우, 10,000 에피소드 동안 모델 전용 분포에서의 미세조정을 통해 정확한 최적의 파이프라인 경계를 추론할 수 있었다.
• 여러 NLP 및 CNN 모델에서 베이스라인 히ュ리스틱 및 수작업 방법보다 더 높은 훈련 처리량을 달성하였다.
• HLO IR의 사용은 TensorFlow, JAX, Trax 간에서 프레임워크에 종속되지 않고 침입 최소화된 병렬화 계획을 가능하게 하였다.
• 작업에 특화된 프루닝을 DQN에 통합함으로써 하이브리드 병렬화 전략의 복잡한 공간에서 탐색 효율성과 수렴 속도가 향상되었다.