[논문 리뷰] DreamShard: Generalizable Embedding Table Placement for Recommender Systems

Daochen Zha, Louis Feng외 7명|arXiv (Cornell University)|2022. 10. 05.

Recommender Systems and Techniques인용 수 7

한 줄 요약

DreamShard는 분산 추천 시스템에서 일반화 가능한 임bedding 테이블 배치를 위한 강화학습 접근법을 제안한다. 비용 네트워크를 통해 융합된 연산 비용을 예측하고, 추정된 MDP에 기반한 정책 네트워크를 훈련시켜, 테이블 수나 장치 수가 다른 새로운 환경에서도 피팅 없이도 일반화를 달성한다. 이는 합성 및 실사용 환경 워크로드에서 최신 기준 대비 최대 19% 향상된 학습 속도를 달성한다.

ABSTRACT

We study embedding table placement for distributed recommender systems, which aims to partition and place the tables on multiple hardware devices (e.g., GPUs) to balance the computation and communication costs. Although prior work has explored learning-based approaches for the device placement of computational graphs, embedding table placement remains to be a challenging problem because of 1) the operation fusion of embedding tables, and 2) the generalizability requirement on unseen placement tasks with different numbers of tables and/or devices. To this end, we present DreamShard, a reinforcement learning (RL) approach for embedding table placement. DreamShard achieves the reasoning of operation fusion and generalizability with 1) a cost network to directly predict the costs of the fused operation, and 2) a policy network that is efficiently trained on an estimated Markov decision process (MDP) without real GPU execution, where the states and the rewards are estimated with the cost network. Equipped with sum and max representation reductions, the two networks can directly generalize to any unseen tasks with different numbers of tables and/or devices without fine-tuning. Extensive experiments show that DreamShard substantially outperforms the existing human expert and RNN-based strategies with up to 19% speedup over the strongest baseline on large-scale synthetic tables and our production tables. The code is available at https://github.com/daochenzha/dreamshard

연구 동기 및 목표

• 분산 추천 시스템에서 임베딩 테이블 배치 문제를 해결하기 위해, 특히 새로운 구성에 대한 일반화 가능성과 융합 연산 비용 고려가 핵심임을 인지한다.
• 기존 강화학습 기반 장치 배치 방법의 한계를 극복하기 위해, 새로운 테이블 수나 장치 수에 대해 일반화하지 못하고, 융합된 연산 비용을 고려하지 않는 점을 해결한다.
• 피팅 없이도 어떤 새로운 배치 작업에도 직접 일반화할 수 있는 시스템을 설계하기 위해 표현 축소를 활용해 확장성을 확보한다.
• 지능적인 테이블 융합과 배치를 통해 계산 비용과 통신 비용을 균형 잡음으로써 학습 지연을 줄인다.
• 합성 환경과 실사용 환경 모두에서 인간이 설계한 히ュ리스틱과 RNN 기반 기준 대비 뛰어난 성능을 내는 것을 목표로 한다.

제안 방법

• 테이블의 차원, 해시 크기, 분포 등의 특성 기반으로 융합된 임베딩 연산의 계산 비용과 통신 비용을 예측하기 위해 비용 네트워크를 훈련한다.
• 상태와 보상이 비용 네트워크의 예측 결과에서 유도되는 추정된 마르코프 결정 과정(MDP)에 기반해 정책 네트워크를 강화학습으로 훈련한다. 이로써 실제 GPU 실행이 필요 없어진다.
• 테이블 표현에는 합산 축소를, 장치 표현에는 최대값 축소를 적용하여, 다른 수의 테이블과 장치를 가진 새로운 작업에 대해 제로샷 일반화를 가능하게 한다.
• 정책 네트워크는 장치 간 최대 지연을 최소화하는 방식으로 테이블을 배치함으로써 계산 비용과 통신 비용을 균형 잡는다.
• 합성 데이터에서 엔드 투 엔드로 훈련하고, 오픈소스 합성 벤치마크와 실제 산업용 생산 테이블에서 평가한다.
• 초대규모 산업용 추천 모델(약 1,000개의 임베딩 테이블 포함)과 128개 GPU 클러스터를 사용해 128-GPU 환경에서 접근법을 평가한다.

실험 결과

연구 질문

• RQ1 강화학습 기반 접근법이 피팅 없이도 다양한 수의 테이블과 장치를 가진 새로운 임베딩 테이블 배치 작업에 일반화할 수 있는가?
• RQ2 학습된 비용 네트워크가 다양한 테이블 조합에서 융합된 임베딩 연산의 성능을 얼마나 정확하게 예측할 수 있는가?
• RQ3 DreamShard가 계산 비용과 통신 비용을 균형 잡는 데서 인간 히ュ리스틱과 RNN 기반 기준 대비 얼마나 뛰어난가?
• RQ4 테이블 특성(예: 차원, 크기, 분포)이 비용 예측 정확도와 최종 배치 성능에 어떤 영향을 미치는가?
• RQ5 다양한 테라바이트 크기의 임베딩 테이블을 가진 실사용 산업 워크로드에서 의미 있는 속도 향상을 달성할 수 있는가?

주요 결과

• DreamShard는 대규모 합성 테이블과 실사용 워크로드에서 가장 강력한 기준 대비 최대 19% 향상된 학습 속도를 달성한다.
• 약 1,000개의 테이블과 다량 테라바이트 메모리가 포함된 초대규모 산업용 추천 모델에서 DreamShard는 학습 스루풋을 45.3% 향상시켰으며, 다음으로 좋은 기준 대비 30 포인트 높은 성능을 기록했다.
• 생산 데이터셋에서 비용 네트워크는 테스트 MSE 0.303을 기록했으며, 테이블 차원, 풀링 인자, 분포 특성이 정확한 예측에 가장 큰 영향을 미쳤다.
• 제거 실험 결과, 개별 특성을 제거할 경우 테스트 MSE가 증가함을 확인하여, 모든 특성이 비용 예측에 중요하다는 것을 입증했다.
• 균형과 비균형 상황을 포함한 여러 테스트 케이스에서 DreamShard는 무작위 배치와 최상의 인간 전문가 전략 모두를 일관되게 뛰어넘었다.
• 드문 경우이지만 DreamShard의 배치가 최상의 히ュ리스틱보다 略로 나쁜 경우가 있어도, 전체적으로는 경쟁력 있거나 더 우수한 성능을 유지함으로써 강건성과 일반화 능력을 입증했다.