[논문 리뷰] Kernel Operations on the GPU, with Autodiff, without Memory Overflows

Benjamin Charlier, Jean Feydy외 3명|arXiv (Cornell University)|2020. 03. 27.

Parallel Computing and Optimization Techniques참고 문헌 8인용 수 80

한 줄 요약

KeOps는 대규모 문제에서 큰 커널 버퍼를 피하고 전통적인 GPU 백엔드보다 우수하게, 자동 미분을 갖춘 커널 및 거리 계산을 위한 빠르고 메모리 효율적인 GPU 프레임워크를 semi-symbolic LazyTensor 추상화를 사용하여 제공한다.

ABSTRACT

The KeOps library provides a fast and memory-efficient GPU support for tensors whose entries are given by a mathematical formula, such as kernel and distance matrices. KeOps alleviates the major bottleneck of tensor-centric libraries for kernel and geometric applications: memory consumption. It also supports automatic differentiation and outperforms standard GPU baselines, including PyTorch CUDA tensors or the Halide and TVM libraries. KeOps combines optimized C++/CUDA schemes with binders for high-level languages: Python (Numpy and PyTorch), Matlab and GNU R. As a result, high-level "quadratic" codes can now scale up to large data sets with millions of samples processed in seconds. KeOps brings graphics-like performances for kernel methods and is freely available on standard repositories (PyPi, CRAN). To showcase its versatility, we provide tutorials in a wide range of settings online at \url{www.kernel-operations.io}.

연구 동기 및 목표

• 대규모 데이터에 대한 GPU의 메모리 효율적인 커널 계산 필요성의 동기를 부여한다.
• 전체 커널 행렬을 구현하는 것을 피하는 반-상징적 감소 프레임워크를 소개한다.
• LazyTensor 추상화와 그것이 Python, NumPy, PyTorch와 어떻게 통합되는지 설명한다.
• 성능과 재현성을 위한 자동 미분 지원과 컴파일된 C++ 백엔드를 제시한다.
• 표준 GPU 라이브러리와의 성능 비교를 평가하고 의도된 사용법과 한계를 논의한다.

제안 방법

• 기호 행(entry) F(p, x_i, y_j)에서 작동하는 일반적인 감소 프레임워크를 정의한다.
• 의미 계산을 구성하고 감소를 통해 평가를 트리거하는 LazyTensor 래퍼를 설명한다.
• 각 감소식/공식에 대해 최적화된 CUDA 코드를 컴파일하고 재사용을 위한 이진 파일을 저장하는 내부 C++ 엔진을 설명한다.
• 임의 차수의 그래디언트를 위한 PyTorch autograd와의 자동 미분 통합을 설명한다.
• kernel 행렬-벡터 곱을 사용하여 KeOps의 성능을 PyTorch, TensorFlow, Halide, TVM, JAX 및 도메인 특화 라이브러리와 비교한다.
• Nyström 및 FFM 근사와 같은 향후 확장 및 CUDA Tensor Core 지원을 포함한 의도된 사용 사례와 향후 개선점을 강조한다.

실험 결과

연구 질문

• RQ1 GPU/CPU에서 선형 메모리 풋프린트로 커널/거리 행렬 연산을 어떻게 수행할 수 있는가?
• RQ2 대규모 커널 연산에 대해 사용하기 쉽고 높은 성능을 동시에 제공하는 반-상징적이고 고수준의 인터페이스가 가능할까?
• RQ3 특화된 커널 연산 엔진과 자동 미분이 어떻게 통합되는가?
• RQ4 다양한 데이터 크기에서 KeOps가 주류 GPU 라이브러리와 비교하여 실제로 얻는 성능 이점은 무엇인가?
• RQ5 이 프레임워크의 근사 및 희소성에 대한 한계와 향후 방향은 무엇인가?

주요 결과

• KeOps는 데이터 배열에서 기호 수식 F를 평가하면서 선형 메모리 풋프린트를 달성하고 전체 커널 행렬 저장을 피한다.
• LazyTensor 추상화를 통해 사용자는 NumPy/PyTorch와 유사한 코드를 작성하고 기호 계산을 느리게 구성하며 최종 감소를 실제로 발생시키는 것만 트리거한다.
• KeOps는 증가하는 데이터 크기에서 가우시안 커널 벡터 곱에서 PyTorch, Halide, TVM 및 기타 기준선보다 뛰어나며 훨씬 작아진 메모리 사용으로 경쟁력 있는 또는 우수한 성능을 제공한다.
• 벤치마킹 표는 KeOps가 특정 구성에서 0.4 ms(N=10k) 및 1.38 s(N=1M)와 같은 결과를 달성하는 반면 일부 경쟁자는 대형 N에서 메모리 오류를 발생시키는 것을 보여준다.
• KeOps는 torch.autograd와 통합된 자동 미분을 지원하여 KeOps 연산을 통한 역전파가 가능하다.
• 논문은 KeOps가 차원 최대 100에서 10k 포인트 이상에 가장 유용하며 Nyström/FFM 근사 및 CUDA Tensor Core 지원에 대한 향후 작업을 제시한다고 주장한다.