TensorDIMM 논문 리뷰
안녕하세요, 오늘 리뷰해드릴 TensorDIMM 논문은 PIM(Process In Memory)을 공부하던 중 찾았는데, 좋은 내용이어서 공유하고자 합니다.
이 논문은 Facebook, Netflix, Youtube 등에서 사용되는 추천 모델(Recommendation model)에 주목하고 있습니다.
이러한 추천 모델에 쓰이는 DNN 기반의 머신 러닝에서(흔히 딥러닝으로 알려져 있죠), 사용되는 데이터양이 시간이 지날수록 급격하게 증가하고 있습니다. 그런데 이는 엄청난 "memory wall"에 직면하게 됩니다.
즉, 딥러닝 알고리즘은 엄청난 양의 memory capacity와 memory bandwidth를 동시에 요구하게 됩니다.
쉽게 말씀드리면 사용되는 데이터양도 많지만, 그 많은 데이터가 연산 되고 load/store 되는 시간당 처리량이 매우 많다는 뜻입니다.
사람들의 선호도에 대한 정보, 예를 들어 영화라면 어떤 장르를 좋아하는지, 어떤 배우를 좋아하는지 등의 정보는 Dense 또는 Sparse 한 features로 저장됩니다.
이것에 대해 조금 더 자세히 말씀드리면, Sparse feature의 대표적인 예시로는 one-hot encoding이 있습니다.
영화라는 아이템에서, 임베딩 dimension이 [ 액션, 드라마, 범죄, 로맨스, 코미디, 공포 ]라면, 어떤 로맨스 영화에 대한 정보를 [0, 0, 0, 1, 0, 0]로 표현하는 식입니다.
물론, 장르가 2개 이상일 수도 있겠지만, 원핫인코딩에서 1은 단 한 개의 값만 갖게 되어있습니다.
원솟값이 대부분이 0으로 이루어져 있고 딱 하나만 1로 이루어져 있어서 정보가 매우 희소하죠. 그래서 Sparse feature라고 부릅니다.
하지만 보시는 것처럼 아이템의 특성을 자세히 표현하려면 벡터의 차원을 계속 늘려야 하고, 단어 간의 유사성을 표현하기에도 어렵습니다.
그래서 dense 하게 벡터를 표현하는 워드 임베딩이라는 기법이 있는데, 대표적으로 구글에서 개발한 word2vec 학습 방식이 있습니다.
W(서울) = [ 0.2, -0.3, -0.1…]
W(도쿄) = [ 0.2, 0.3, -0.7…]
위와 같이 단어의 다양한 추상적 특성을 다차원 배열로 표현하고 아래와 같이 단어 간 유사도를 구현하는 방식입니다.
W(이모) - W(이모부) = W(여성) - W(남성)
W(서울) - W(한국) + W(일본) = W(도쿄)
이러한 방식의 학습은 인공지능이 문장을 지어낼 때, 문맥상 자연스러운 단어를 찾아내는 데 도움을 줄 수 있겠죠.
이러한 표현법은 원핫인코딩에 비해서 배열 안의 모든 원소가 다른 벡터와 관계도를 나타내는 정보를 밀도 있게 가지고 있어서 dense features라고 할 수 있습니다.
그렇지만 단정적으로 dense feature가 더 좋고 sparse는 쓰이지 않는다는 말은 아닙니다.
실제 DNN 트레이닝에서 sparse 한 정보를 이용하는 경우도 많이 있을 겁니다.
추천 알고리즘에서는 각자의 활용처가 있으리라 생각합니다.
이러한 feature 정보들은 데이터양 자체가 매우 많아서, 즉, 성능이 메모리 용량에 의해 주로 결정 나서 Memory capacity dominated라고 부릅니다.
다음 단계는 임베딩 작업인데, 사용자들이 가지고 있는 정보와 아이템이 가진 정보를 모으는 단계(Gathering)입니다.
여기서는 해당 데이터를 찾고, 모으는 작업이 주되기 때문에 고 대역폭이어야 원활한 데이터 처리를 할 수 있겠죠.
이를 Memory bandwidth dominated라고 합니다.
다음 그러한 feature들의 상호작용 단계의 경우, 각 임베딩을 거친 하드웨어 사이에서의 통신이 빨라야 수월하게 처리할 수 있는 Communication dominated 영역이 됩니다.
마지막으로 가장 상위의 다층 퍼셉트론(MLP, Multi-Layer Perceptron)에서는 연산량이 많아 Compute dominated 한 영역이라고 합니다.
Motivation 파트에서는, CPU와 GPU 구조를 나누어서 현재의 한계를 설명합니다.
GPU 또는 NPU의 경우 병렬 작업처리에 매우 유리한데, 한 번에 이동하는 데이터양이 많다 보니 bandwidth에 유리한 3D stacked memory(HBM, HMC)를 사용합니다.
이러한 메모리 구조는 인터포저 위에 TSV(Through Silicon Via) 기술을 이용해 저전력으로 많은 데이터를 한 번에 전송할 수 있는 장점을 이루었지만, 구조상 한 층당 크기를 무한히 넓힐 수는 없는 제약이 있습니다.
논문에서는 이러한 적층형 구조의 메모리는 수십 GB의 용량밖에 가지지 못한다고 합니다.
CPU 메모리의 경우에는 상대적으로 고용량의 DRAM을 채택할 수 있지만, sequential 하게 연산을 수행하는 CPU 특성상, 엄청난 양의 임베딩 레이어를 연산하는 능력은 GPU 또는 NPU에 밀릴 수밖에 없습니다.
그렇기 때문에 많은 양의 feature data는 CPU 메모리에 저장해두고, 임베딩 룩업 단계는 GPU로 가지고 와서 계산하는 기법을 생각해볼 수 있습니다.
이렇게 되면 CPU만 사용했을 때보다 DNN 연산을 훨씬 빨리 끝낼 수 있는 장점이 있지만, cudaMemcpy 함수를 이용해서 PCIe 통신을 이용해서 CPU->GPU로 데이터를 넘겨주는 동안의 communication latency가 있습니다.
GPU만 사용해서 연산한다면 사실 제일 좋지만, 앞서 말씀드린 대로 GPU 메모리 capacity의 한계로 feature들을 다 담기 어렵다는 문제가 있습니다.
GPU 메모리가 무한하다는 가정하에, 전체 작업 속도는 아래와 같이 차이가 납니다.
B(N)이란 N size를 갖는 batch를 의미하는데요, N값이 1일 땐 CPU만 가지고 연산하는 성능도 꽤 좋지만, batch size가 커지면 GPU로만 연산하는 게 제일 유리하다는 것을 알 수 있습니다.
그리고 이 논문에서는 다층 퍼셉트론의 dimension 증가보다 Embedding dimension의 증가가 model size를 훨씬 dramatic 하게 증가시킨다고 설명하고 있습니다.
그만큼 임베딩 하려는 개수가 모델 연산량을 좌지우지한다고 볼 수 있겠죠.
따라서 이 논문에서는 memory capacity와 bandwidth, 두 마리의 토끼를 잡기 위해 TensorDIMM과 TensorNode를 제안하게 됩니다.
이 논문에서 제안한 TensorDIMM은 메모리 제약을 가진 추천 시스템에서 시스템 레벨의 보틀넥에 의문을 제기하는 다음 두 가지 키 포인트를 가지고 있습니다.
첫째, 임베딩 feature에 사용되는 중요 텐서 연산은 element-wise한 연산입니다.
모인 N 임베딩을 개별적으로 GPU 메모리에서 복사해와서 GPU에서 연산을 시작하기보다, N 임베딩을 모으고 near-memory에서 축약시켜 압축된 싱글 텐서를 GPU 메모리에서 복사해오자는 논리입니다.
이러한 Near-memory processing(NMP)는 임베딩을 모아오는 latency를 줄여줄뿐 아니라, N배 만큼의 데이터 전송 사이즈를 줄여주고 이에 따라 CPU->GPU cudaMemcpy의 통신 보틀넥을 덜어주는 효과도 있습니다.
둘째, 딥러닝 시스템 공급업자들은 GPU/NPU와 중앙 처리 장치와의 연결에 legacy host device PCIe와는 독립적인 기술을 사용하고 있습니다.
그러한 기술들은 대표적으로 NVLINK가 있으며, GPU/NPU간 통신에서 PCIe보다 대략 9배 빠릅니다.
이 기술의 강력한 장점으로는 비-가속기 부분들도 아주 강력하게 통합시킬 수 있다는 부분에 있습니다.
출처 : Youngeun Kwon, Yunjae Lee, Minsoo Rhu. 2019. TensorDIMM: A Practical Near-Memory Processing Architecture for Embeddings and Tensor Operations in Deep Learning. In Proceedings of the International Symposium on Microarchitecture(MICRO).