[논문 리뷰] Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning

Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan외 3명|arXiv (Cornell University)|2019. 06. 20.

Ferroelectric and Negative Capacitance Devices참고 문헌 23인용 수 30

한 줄 요약

Placeto는 계산 그래프의 구조적 정보를 인코딩하기 위해 그래프 신경망을 사용하고, 배치를 반복적인 개선 단계로 모델링함으로써 분산 딥러닝을 위한 일반화 가능한 장치 배치 정책을 학습하는 강화학습 프레임워크를 제안한다. 이는 재학습 없이도 새로운 그래프에 일반화되며, 최대 6.1배 빠른 학습을 달성하며, 이전의 RNN 기반 방법보다 속도와 배치 품질 면에서 뛰어나다.

ABSTRACT

We present Placeto, a reinforcement learning (RL) approach to efficiently find device placements for distributed neural network training. Unlike prior approaches that only find a device placement for a specific computation graph, Placeto can learn generalizable device placement policies that can be applied to any graph. We propose two key ideas in our approach: (1) we represent the policy as performing iterative placement improvements, rather than outputting a placement in one shot; (2) we use graph embeddings to capture relevant information about the structure of the computation graph, without relying on node labels for indexing. These ideas allow Placeto to train efficiently and generalize to unseen graphs. Our experiments show that Placeto requires up to 6.1x fewer training steps to find placements that are on par with or better than the best placements found by prior approaches. Moreover, Placeto is able to learn a generalizable placement policy for any given family of graphs, which can then be used without any retraining to predict optimized placements for unseen graphs from the same family. This eliminates the large overhead incurred by prior RL approaches whose lack of generalizability necessitates re-training from scratch every time a new graph is to be placed.

연구 동기 및 목표

• 각 새로운 계산 그래프에 대해 다시 학습을 시작해야 하는 기존의 RL 기반 장치 배치 방법의 높은 학습 오버헤드 문제를 해결하기 위해.
• 새로운 그래프에 대해 재학습 없이도 동일한 가족에 属하는 새로운 그래프에 대해 학습된 배치 정책을 일반화할 수 있도록 하기 위해.
• 분산 학습에서 고성능 장치 배치를 찾기 위해 필요한 배치 평가 수를 줄이기 위해.
• 임의의 노드 인덱싱 및 순회 순서에 대해 불변인 정책을 설계하여 정책의 강건성과 이식 가능성을 향상시키기 위해.
• 속도와 성능 면에서 RNN 기반 기준선을 능가하는 확장 가능하고 효율적이며 일반화 가능한 장치 배치 방법을 개발하기 위해.

제안 방법

• Placeto는 장치 배치를 반복적인 개선 단계의 시퀀스로 모델링하며, 정책은 현재 부분적 배치를 기반으로 한 번에 한 노드의 장치를 예측한다.
• 계산 그래프의 구조적 정보를 캡처하기 위해 그래프 신경망(GNN)을 사용하며, 메시지 전달을 통해 부모-자식 및 병렬 노드 간의 관계를 노드 순서나 레이블과 무관하게 인코딩한다.
• 정책 네트워크는 노드의 국소적 주변부(부모, 자식, 병렬 노드)의 특징을 별도의 서브넷을 통해 집계한 후 장치 결정을 내린다.
• GNN 기반 아키텍처는 RNN 기반 방법이 순차적 처리에 의존하는 것과 달리, 다른 노드 인덱싱이나 순서에 대해 일반화할 수 있도록 한다.
• 정책 학습을 더 빠른 수렴으로 이끌기 위해 중간 보상 구조를 포함한 강화학습 목표를 사용한다.
• 프레임워크는 각 그래프 가족당 하나의 일반화 가능한 정책을 학습하며, 이 정책은 해당 가족의 새로운 그래프에 대해 재학습 없이 적용 가능하다.

실험 결과

연구 질문

• RQ1 분산 딥러닝에서 새로운 계산 그래프에 대해 재학습 없이 일반화 가능한 강화학습 정책을 학습시킬 수 있는가?
• RQ2 배치를 반복적인 개선으로 모델링함으로써 학습 복잡도를 줄이고 정책의 일반화 성능를 향상시킬 수 있는가?
• RQ3 노드 레이블이나 순서에 의존하지 않고 그래프 신경망이 계산 그래프의 구조적 정보를 효과적으로 인코딩할 수 있는가?
• RQ4 Placeto의 일반화 능력은 RNN 기반 접근법과 비교해 배치 품질과 학습 효율성 측면에서 어떻게 다른가?
• RQ5 정책 아키텍처의 선택이 다양한 그래프 가족 간의 일반화 성능에 얼마나 큰 영향을 미치는가?

주요 결과

• Placeto는 RNN 기반 방법 대비 최대 6.1배 적은 배치 평가 수를 기록하면서도, 최고의 이전 결과와 동등하거나 그 이상의 배치 품질을 달성한다.
• NMT 모델에서 Placeto는 RNN 기반 접근법보다 16.5% 더 빠른 실행 속도를 기록한 배치를 찾았다.
• 재학습 없이 일반화 가능한 정책을 사용하는 Placeto Zero-Shot은 수천 번의 평가가 필요한 Placeto Optimized와 거의 동일한 성능을 달성했다.
• RNN 기반 접근법은 제로샷 설정에서 성능이 열악했으며, RNN Zero-Shot은 랜덤 기반 성능과 거의 동일하여 일반화 능력 부족을 확인할 수 있었다.
• Placeto의 정책는 다양한 노드 순회 순서에 대해 강건하게 일반화되며, 평균적으로 실행 시간 변동이 5% 이내였지만, RNN 기반 방법은 평균 30%의 변동을 보였다.
• 제거 실험 결과, 부모, 자식, 병렬 노드의 국소적 주변부 집계 및 풀링이 일반화에 필수적임을 확인했으며, 이들 구성 요소를 제거할 경우 성능이 13~20% 저하되었다.