[논문 리뷰] MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems

Tianqi Chen, Mu Li외 8명|arXiv (Cornell University)|2015. 12. 03.

Advanced Memory and Neural Computing참고 문헌 9인용 수 1,917

한 줄 요약

MXNet는 이질적 시스템, 모바일 기기에서 분산 GPU 클러스터에 이르기까지 효율적인 실행을 위한 상징적 계산과 명령형 텐서 연산을 통합하는 경량 다중 언어 기계학습 라이브러리이다. 통합 엔진 스케줄링, 인라인 연산 및 최적화된 통신을 통해 고성능과 메모리 효율성을 달성하였으며, ImageNet에서 10배의 속도 향상을 보이며 초선형 확장성을 입증하였다.

ABSTRACT

MXNet is a multi-language machine learning (ML) library to ease the development of ML algorithms, especially for deep neural networks. Embedded in the host language, it blends declarative symbolic expression with imperative tensor computation. It offers auto differentiation to derive gradients. MXNet is computation and memory efficient and runs on various heterogeneous systems, ranging from mobile devices to distributed GPU clusters. This paper describes both the API design and the system implementation of MXNet, and explains how embedding of both symbolic expression and tensor operation is handled in a unified fashion. Our preliminary experiments reveal promising results on large scale deep neural network applications using multiple GPU machines.

연구 동기 및 목표

• 다양한 하드웨어에서 효율적이고 확장 가능한 시스템이 필요한 딥러닝 워크로드의 복잡성 증가에 대응한다.
• 최적화 잠재력과 구현 유연성의 균형을 이루기 위해 상징적 및 명령형 프로그래밍 패러다임을 통합한다.
• CPU, GPU 및 분산 클러스터를 포함한 이질적 시스템에서 최소한의 성능 오버헤드로 효율적인 실행을 가능하게 한다.
• 정확성과 재현 가능성을 유지하면서도 인라인 연산 및 메모리 재사용을 통해 메모리 포트폴리오를 줄인다.
• 최소한의 코드 변경으로 다수의 머신 간 원활한 분산 학습을 지원하며 고도의 확장성을 달성한다.

제안 방법

• 모든 연산 간 데이터 및 의존성 흐름을 추적하기 위해 상징적 계산 그래프와 명령형 텐서 연산을 모두 스케줄링하는 통합 엔진을 사용한다.
• 모든 계산 단위—NDArray, 난수 생성기, 메모리 공간—을 태그된 자원으로 표현하여 세밀한 의존성 추적 및 스케줄링을 가능하게 한다.
• 변경 연산을 일급 자원으로 지원하여 인라인 업데이트를 허용하고 효율적인 매개변수 업데이트 및 메모리 재사용을 가능하게 한다.
• 분산 학습을 위한 이중 레벨 파라미터 서버 아키텍처(KVStore)를 구현하며, 내부 머신 및 간섭 머신 간 동기화 레이어를 통해 대역폭을 줄이고 다양한 일致성 모델을 지원한다.
• 엔진을 통해 통신 및 계산 스케줄링을 통합하여 장치 및 머신 간의 데이터 일관성과 원활한 조율을 보장한다.
• 인라인 연산 및 메모리 할당 통합과 같은 공격적인 메모리 최적화 기법을 적용하여 학습 및 추론 중 메모리 사용을 최소화한다.

실험 결과

연구 질문

• RQ1 어떻게 한 기계학습 시스템 내에서 상징적 및 명령형 프로그래밍 패러다임을 효과적으로 통합하여 최적화 잠재력과 개발자 유연성의 최고 조합을 달성할 수 있는가?
• RQ2 어떤 시스템 수준의 설계 패턴이 모바일 기기에서 다중 GPU 클러스터에 이르기까지 이질적 하드웨어에서 딥러닝 모델의 효율적 실행을 가능하게 하는가?
• RQ3 성능나 정확성을 희생시키지 않고 딥러닝 학습에서 메모리 포트폴리오를 어느 정도 줄일 수 있는가?
• RQ4 계산 및 통신 스케줄링 통합이 분산 딥러닝의 확장성과 수렴성에 어떤 영향을 미치는가?
• RQ5 통합 엔진이 고성능과 낮은 오버헤드를 유지하면서도 상징적 계산 그래프와 명령형 텐서 연산을 동시에 관리할 수 있는가?

주요 결과

• MXNet는 단일 GPU에서 Torch7 및 Caffe와 유사한 원시 성능를 달성하였으며, 오래된 CUDNN 버전을 사용함으로써 TensorFlow는 약 2배 느린 성능을 보였다.
• 인라인 연산 및 통합 메모리 할당 덕분에 추론 시 최대 4배, 학습 시 최대 2배의 메모리 사용 감소를 기록하였으며, VGGNet의 경우 학습에 16MB 미만의 추가 메모리만 필요로 하였다.
• 10台의 머신을 사용한 분산 학습에서 ILSVRC12 데이터셋에서 평균 데이터 패assing 시간이 단일 머신의 14,000초에서 1,400초로 감소하여 10배의 속도 향상을 기록하였으며, 초선형 확장성을 입증하였다.
• 분산 환경에서 10번의 데이터 패assing 이후 단일 머신 학습보다 빠르게 수렴함으로써 효과적인 통신 및 동기화 전략을 확인하였다.
• 통합 엔진은 명령형 연산과 상징적 계산의 원활한 통합을 가능하게 하여, 성능 손실 없이 동일한 프로그램 내에서 두 패러다임을 모두 사용할 수 있도록 하였다.
• 이중 레벨 KVStore 아키텍처는 간섭 머신 간 통신 대역폭을 감소시키고 내부 및 간섭 머신 동기화를 위한 다양한 일치성 모델을 지원하여 시스템의 효율성과 신뢰성을 향상시켰다.