[논문 리뷰] High-performance, Distributed Training of Large-scale Deep Learning Recommendation Models.

Dheevatsa Mudigere, Yuchen Hao외 49명|arXiv (Cornell University)|2021. 04. 12.

Stochastic Gradient Optimization Techniques참고 문헌 56인용 수 18

한 줄 요약

이 논문은 대규모 딥러닝 추천 모델(DLRMs)의 고성능 분산 학습을 위한 소프트웨어-하드웨어 공동 설계 시스템을 제안하며, 이전 시스템 대비 40배 빠른 속도 향상을 달성한다. 이는 파이토치 기반 학습 스택과 ZionEX 플랫폼을 결합하여, 최적화된 데이터/모델 병렬 처리, 계층적 임bedding 분할, 감소된 정밀도 통신, 고대역폭·저지연성 네트워크 인fra를 통해 최대 12조 파라미터의 모델 학습을 가능하게 한다.

ABSTRACT

Deep learning recommendation models (DLRMs) are used across many business-critical services at Facebook and are the single largest AI application in terms of infrastructure demand in its data-centers. In this paper we discuss the SW/HW co-designed solution for high-performance distributed training of large-scale DLRMs. We introduce a high-performance scalable software stack based on PyTorch and pair it with the new evolution of Zion platform, namely ZionEX. We demonstrate the capability to train very large DLRMs with up to 12 Trillion parameters and show that we can attain 40X speedup in terms of time to solution over previous systems. We achieve this by (i) designing the ZionEX platform with dedicated scale-out network, provisioned with high bandwidth, optimal topology and efficient transport (ii) implementing an optimized PyTorch-based training stack supporting both model and data parallelism (iii) developing sharding algorithms capable of hierarchical partitioning of the embedding tables along row, column dimensions and load balancing them across multiple workers; (iv) adding high-performance core operators while retaining flexibility to support optimizers with fully deterministic updates (v) leveraging reduced precision communications, multi-level memory hierarchy (HBM+DDR+SSD) and pipelining. Furthermore, we develop and briefly comment on distributed data ingestion and other supporting services that are required for the robust and efficient end-to-end training in production environments.

연구 동기 및 목표

• 생산 데이터 센터에서 대규모 딥러닝 추천 모델(DLRMs)의 증가하는 인프라 수요를 해결하기 위해.
• 수십 조 개의 파라미터를 가진 DLRMs의 분산 학습에서 발생하는 성능 저하 문제를 극복하기 위해.
• 결정론적 옵티마이저 업데이트를 유지하면서 하드웨어 활용도를 극대화하는 확장성 있고 고처리량의 학습 시스템을 개발하기 위해.
• 생산 환경에서의 통합된 데이터 인gest 및 지원 서비스를 통해 엔드 투 엔드 학습 효율성을 보장하기 위해.
• 모델 정확도나 학습 안정성에 손상 없이 이전 시스템 대비 뚜렷한 속도 향상을 달성하기 위해.

제안 방법

• 고대역폭, 최적의 토폴로지, 효율적인 전송 프로토콜을 갖춘 전용 확장 네트워크를 갖춘 ZionEX 플랫폼을 설계 및 구현하기 위해.
• 분산 워커 간 모델 병렬 처리와 데이터 병렬 처리를 모두 네이티브로 지원하는 파이토치 기반 학습 스택을 구현하기 위해.
• 임베딩 테이블을 행과 열 차원에 따라 분할하고, 여러 워커 간 부하를 균형 있게 분배하는 계층적 분할 알고리즘을 개발하기 위해.
• 결정론적 옵티마이저 업데이트를 완전히 지원하는 고성능 코어 연산자 통합을 통해 학습 재현 가능성을 유지하기 위해.
• 감소된 정밀도 통신, 다수준 메모리 하이어라키(HBM, DDR, SSD), 통신 피이프라인 기법을 활용하여 지연 시간과 대역폭 부담을 줄이기 위해.
• 신뢰성 있고 확장성 있으며 효율적인 엔드 투 엔드 학습을 보장하기 위해 분산 데이터 인게스트 파이프라인 및 지원 서비스를 구축하기 위해.

실험 결과

연구 질문

• RQ1 생산 데이터 센터 환경에서 최대 12조 파라미터를 가진 DLRMs의 분산 학습을 어떻게 확장성 있고 효율적으로 구현할 수 있는가?
• RQ2 대규모 DLRMs의 이전 학습 시스템 대비 40배 빠른 속도 향상을 달성하기 위해 필요한 시스템 수준 최적화는 무엇인가?
• RQ3 임베딩 테이블의 분할을 어떻게 계층적으로 워커 간에 분할하여 부하 균형을 맞추면서도 통신 오버헤드를 최소화할 수 있는가?
• RQ4 감소된 정밀도 통신과 메모리 하이어라키 최적화가 학습 스루풋을 얼마나 향상시킬 수 있으며, 학습 안정성에 영향을 주지 않는가?
• RQ5 소프트웨어와 하드웨어 컴포넌트의 공동 설계가 결정론적이고 고성능의 대규모 모델 학습을 어떻게 가능하게 하는가?

주요 결과

• 이 시스템은 대규모 DLRMs의 이전 학습 시스템 대비 솔루션 도달 시간에서 40배 빠른 성능 향상을 달성한다.
• 이 플랫폼은 최대 12조 파라미터의 DLRMs 학습을 지원하여, 알려진 최대 규모의 추천 모델에까지 확장 가능성을 입증한다.
• 계층적 임베딩 테이블 분할은 분산 워커 간 효과적인 부하 균형을 맞추며 통신 병목 현상을 최소화한다.
• 감소된 정밀도 통신과 피이프라인 처리 기법은 학습 지연 시간을 크게 줄이고 하드웨어 활용도를 향상시킨다.
• HBM, DDR, SSD 메모리 계층의 통합은 큰 모델 상태와 기울기 처리를 효율적으로 가능하게 한다.
• 시스템은 강력한 최적화에도 불구하고 완전히 결정론적 옵티마이저 업데이트를 유지하여 학습 재현 가능성을 보장한다.