[논문 리뷰] Mixed-Precision Embedding Using a Cache

Jie Yang, Jianyu Huang외 3명|arXiv (Cornell University)|2020. 10. 21.

Topic Modeling참고 문헌 32인용 수 15

한 줄 요약

이 논문은 대규모 추천 시스템에서 자주 또는 최근에 액세스되는 임bedding 행을 저장하기 위해 작은 고정밀도 캐시를 사용하는 혼합 정밀도 임베딩 훈련 방법을 제안한다. 대부분의 임베딩 행은 저정밀도(INT8 또는 INT4)로 훈련되며, 이 방법은 GPU에서 호스트로의 전송을 최소화하고 효율적인 캐시 교체 정책을 활용함으로써 최대 7배의 메모리 절감과 정확도 유지, 산업 규모 모델에서 16%의 훈련 속도 향상을 달성한다.

ABSTRACT

In recommendation systems, practitioners observed that increase in the number of embedding tables and their sizes often leads to significant improvement in model performances. Given this and the business importance of these models to major internet companies, embedding tables for personalization tasks have grown to terabyte scale and continue to grow at a significant rate. Meanwhile, these large-scale models are often trained with GPUs where high-performance memory is a scarce resource, thus motivating numerous work on embedding table compression during training. We propose a novel change to embedding tables using a cache memory architecture, where the majority of rows in an embedding is trained in low precision, and the most frequently or recently accessed rows cached and trained in full precision. The proposed architectural change works in conjunction with standard precision reduction and computer arithmetic techniques such as quantization and stochastic rounding. For an open source deep learning recommendation model (DLRM) running with Criteo-Kaggle dataset, we achieve 3x memory reduction with INT8 precision embedding tables and full-precision cache whose size are 5% of the embedding tables, while maintaining accuracy. For an industrial scale model and dataset, we achieve even higher >7x memory reduction with INT4 precision and cache size 1% of embedding tables, while maintaining accuracy, and 16% end-to-end training speedup by reducing GPU-to-host data transfers.

연구 동기 및 목표

• 딥러닝 추천 시스템에서 대규모 임베딩 테이블의 증가하는 메모리 사용량을 해결하기 위해.
• 대규모 훈련 워크로드에서 모델 정확도를 훼손하지 않고 GPU 메모리 사용량을 줄이기 위해.
• 임베딩 테이블에 저정밀도 훈련과 고정밀도 캐시를 조합하는 효과를 탐색하기 위해.
• 캐시 크기, 교체 정책, 정밀도 수준이 모델 정확도와 훈련 효율성에 미치는 영향을 평가하기 위해.
• 더 스마트한 메모리 접근 패턴을 통해 PCIe 대역폭 사용을 줄임으로써 훈련 속도를 향상시키기 위해.

제안 방법

• 이 방법은 임베딩 테이블 크기의 1~5% 정도의 작은 고정밀도 캐시를 사용하여 자주 또는 최근에 액세스된 임베딩 행을 저장한다.
• 대부분의 임베딩 행은 저정밀도(INT8 또는 INT4)로 훈련되며, 캐시는 활성 행에 대한 전체 정밀도 기울기와 파라미터를 저장한다.
• LRU 및 LFU와 같은 캐시 교체 정책을 사용하여 액세스 빈도나 최종 액세스 시점에 따라 어느 행을 고정밀도로 승격시킬지 관리한다.
• 저정밀도 훈련 중에는 체계적인 반올림 편향을 완화하기 위해 확률적 반올림을 적용한다.
• 기존 딥러닝 프레임워크와 통합되며, 대규모 테이블에 대해 통합 메모리(UVM)를 사용하고, 캐시가 PCIe 트래픽을 줄인다.
• 이 방법은 오픈소스(DLRM with Criteo-Kaggle) 및 산업 규모 모델 모두에서 평가되었으며, 메모리 절감, 정확도, 훈련 속도를 측정하였다.

실험 결과

연구 질문

• RQ1 대부분의 임베딩 행이 저정밀도(INT8/INT4)로 훈련되더라도, 작은 고정밀도 캐시가 모델 정확도를 효과적으로 유지할 수 있는가?
• RQ2 다양한 캐시 교체 정책(LRU, LFU)이 저정밀도 임베딩 훈련에서 모델 정확도와 캐시 히트율에 어떤 영향을 미치는가?
• RQ3 메모리 절감과 정확도 복구 사이의 균형을 고려할 때, 임베딩 테이블에 비해 최적의 캐시 크기는 얼마인가?
• RQ4 고정밀도 캐시가 있는지 여부에 따라, 확률적 반올림과 반올림-가장가까운 방식의 성능 비교는 어떻게 되는가?
• RQ5 이 캐싱 메커니즘이 GPU에서 호스트로의 데이터 전송을 얼마나 줄일 수 있으며, 종합적인 훈련 속도 향상에 어떤 기여를 하는가?

주요 결과

• INT4 정밀도와 1% 고정밀도 캐시를 사용하여 산업 규모 모델에서 최대 7배의 메모리 절감을 달성했으며, 정확도는 중립적 유지되었다.
• 5% 고정밀도 캐시를 사용하여 오픈소스 DLRM 모델에서 INT8 정밀도로 3배의 메모리 절감을 달성했으며, 정확도는 완전히 유지되었다.
• 확률적 반올림은 특히 고정밀도 캐시와 조합되었을 때 체계적인 반올림 편향을 줄여주므로, 반올림-가장가까운 방식보다 항상 우수한 성능을 보였다.
• 캐시 히트율은 훈련 속도와 강한 상관관계를 보였다: 높은 히트율은 GPU 메모리 대역폭을 포화시키며, 낮은 히트율조차도 통합 메모리(UVM) 접근보다 우수한 성능을 보였다.
• 효율적인 캐싱을 통해 GPU에서 호스트로의 데이터 전송을 줄임으로써 산업 모델에서 종합적인 훈련 속도 향상 16%를 달성했다.
• 캐시 크기를 늘릴수록 정확도 향상 효과는 점점 감소하는 경향을 보였으며, 이는 작은 캐시(1~5%)로도 거의 최적의 성능을 달성할 수 있음을 시사한다.