[논문 리뷰] TF-Replicator: Distributed Machine Learning for Researchers

Peter Buchlovsky, David Budden외 10명|arXiv (Cornell University)|2019. 02. 01.

Adversarial Robustness in Machine Learning참고 문헌 47인용 수 21

한 줄 요약

TF-Replicator는 TensorFlow 기반 프레임워크로, 최소한의 코드 변경으로 데이터 병렬 또는 모델 병렬 학습을 통해 CPU, GPU, TPU로 구성된 분산 클러스터에서 딥러닝 모델을 쉽게 확장할 수 있도록 한다. 이 프레임워크는 분산 시스템 전문 지식이 없이도 다양한 아키텍처—예를 들어 ResNet-50, SN-GAN, D4PG 모델—에서 뛰어난 확장성을 달성하며, Tensorflow 2.0의 일부로 오픈소스화될 예정이다.

ABSTRACT

We describe TF-Replicator, a framework for distributed machine learning designed for DeepMind researchers and implemented as an abstraction over TensorFlow. TF-Replicator simplifies writing data-parallel and model-parallel research code. The same models can be effortlessly deployed to different cluster architectures (i.e. one or many machines containing CPUs, GPUs or TPU accelerators) using synchronous or asynchronous training regimes. To demonstrate the generality and scalability of TF-Replicator, we implement and benchmark three very different models: (1) A ResNet-50 for ImageNet classification, (2) a SN-GAN for class-conditional ImageNet image generation, and (3) a D4PG reinforcement learning agent for continuous control. Our results show strong scalability performance without demanding any distributed systems expertise of the user. The TF-Replicator programming model will be open-sourced as part of TensorFlow 2.0 (see https://github.com/tensorflow/community/pull/25).

연구 동기 및 목표

• 딥러닝 연구 분야에서 점점 증가하는 확장성 있고 일반적인 목적의 분산 학습 요구에 부응하기 위해.
• 저수준의 분산 시스템 관련 고려사항을 추상화함으로써 분산 학습을 구현하는 복잡성을 줄이기 위해.
• 동일한 모델을 다양한 하드웨어(CPU, GPU, TPU)와 학습 방식(동기/비동기) 간에 원활하게 배포할 수 있도록 하기 위해.
• 다중 손실 최적화나 강화학습 에이전트와 같은 복잡한 비표준 학습 패턴을 지원하기 위해.
• 연구자들이 기존의 단일 머신 모델을 새로운 도구를 배우지 않고도 최소한의 수정으로 분산 클러스터로 확장할 수 있도록 하여 연구 반복 속도를 가속화하기 위해.

제안 방법

• TF-Replicator는 TensorFlow의 분산 실행을 추상화하는 고수준 API를 제공하여, 사용자가 한 번 모델을 정의하고도 이질적인 클러스터에서 배포할 수 있도록 한다.
• 레플리카 기반 프로그래밍 모델을 사용하여 각 워커에서 모델의 복제본을 실행하고, 파라미터 서버 또는 집합적 통신(e.g., NCCL)을 통해 자동으로 기울기 동기화를 수행한다.
• 데이터 병렬 및 모델 병렬 학습을 모두 지원하며, 기계 간 장치 간 자동 통신 삽입과 자동 장치 배치를 통해 투명하게 처리한다.
• TensorFlow 2.0의 즉각 실행(eager execution)과 Keras 스타일 API와 통합되어 자연스러운 제어 흐름과 융통성 있는 학습 루프를 가능하게 한다.
• 사용자가 장치 배치 및 통신 전략을 구성 가능하게 함으로써 데이터 병렬과 모델 병렬을 조합할 수 있는 하이브리드 병렬 처리를 지원한다.
• 클러스터 설정은 ClusterSpec을 통해 자동으로 처리되며, 논리적 장치를 물리적 하드웨어에 매핑하여 단일 머신과 다중 머신 환경 간 이식성을 보장한다.

실험 결과

연구 질문

• RQ1 저수준의 분산 시스템 전문 지식이 없이도 연구자들이 분산 학습을 구현하는 데 있어 고수준 추상화가 얼마나 단순화할 수 있는가?
• RQ2 이 프레임워크는 다양한 모델 아키텍처(CNN, GAN, RL 에이전트)와 하드웨어(GPU, TPU) 간에 얼마나 잘 확장되는가?
• RQ3 TF-Replicator는 다중 손실 최적화나 계층적 강화학습과 같은 복잡한 학습 방식을 지원할 수 있는가?
• RQ4 스케일링 능력과 학습 속도 측면에서 수동 최적화된 분산 시스템과 비교해 성능가 얼마나 우수한가?
• RQ5 연구자들이 단일 머신 코드를 얼마나 최소한의 수정으로 분산 학습에 재사용할 수 있는가?

주요 결과

• TF-Replicator는 ImageNet에서 ResNet-50 모델에 대해 강력한 확장성을 확보하여 GPU 및 TPU 클러스터에서 경쟁 가능한 학습 속도를 달성했다.
• 클래스 조건부 이미지 생성을 위한 SN-GAN 모델은 TF-Replicator를 통해 효과적으로 확장되었으며, 이는 복잡한 생성 모델 지원 능력을 입증한다.
• D4PG 강화학습 에이전트는 훈련 시간을 크게 가속화했으며, 8개의 V100 GPU는 8코어 TPUv2 장치의 성능을 동등하게 구현했다.
• 단일 TPUv2 장치가 NVLink 연결된 8개의 V100 GPU와 유사한 성능을 보여주었으며, 이는 프레임워크의 통신 및 장치 관리 효율성을 입증한다.
• 프레임워크는 다양한 클러스터 아키텍처에서 최소한의 코드 변경으로 모델을 확장시켜 높은 처리량과 낮은 통신 오버헤드를 유지했다.
• TF-Replicator는 다중 손실 및 계층적 RL을 포함한 복잡한 학습 루프를 가진 모델을 학습하는 데 성공했으며, 저수준의 분산 시스템 프로그래밍이 필요 없었다.