nvmath-python 소개
NVIDIA가 공개한 nvmath-python은 기존의 Python 환경에 NVIDIA의 고성능 수학 라이브러리를 직접 연결해 주는 프로젝트입니다. 지금까지 NumPy, CuPy, PyTorch와 같은 배열·텐서 연산 프레임워크는 GPU 연산을 지원하더라도 NVIDIA의 모든 수학 라이브러리 기능을 완전히 활용하기에는 한계가 있었습니다. nvmath-python은 이러한 공백을 채우며, CUDA 기반의 cuBLASLt, cuFFT 등과 같은 NVIDIA 라이브러리를 Python 사용자들이 친숙한 문법으로 다룰 수 있도록 돕습니다.
또한 nvmath-python은 단순히 기존 API를 Python으로 감싸는 수준이 아니라, 파이썬 친화적인 인터페이스를 제공하고, 기존 프레임워크에서 부족했던 고급 기능들까지 노출합니다. 예를 들어, 행렬 곱셈에서의 mixed precision 연산 선택이나 FFT에서 사용자 정의 epilog 삽입 같은 기능을 통해 세밀한 제어가 가능합니다. 이 덕분에 연구자나 개발자는 고성능 연산을 위한 GPU 활용도를 극대화할 수 있으며, 기존 코드와도 원활히 통합됩니다.
현재 nvmath-python은 베타 버전이므로 완벽한 안정성을 보장하지는 않지만, NVIDIA는 개발자 커뮤니티의 피드백을 통해 빠르게 개선해 나갈 계획임을 밝혔습니다. 따라서 GPU 가속 기반의 수학적 연산이 중요한 분야, 예를 들어 머신러닝, 과학 연산, 신호 처리 등을 연구하거나 개발하는 분들에게 특히 주목할 가치가 있습니다.
nvmath-python와 CuPy, PyTorch와의 비교
nvmath-python은 CuPy, PyTorch와 같은 기존 Python 기반의 GPU 연산 프레임워크와 직접 경쟁하기보다는, 그들의 부족한 부분을 보완하는 성격을 가지고 있습니다.
예를 들어 CuPy는 NumPy 호환 배열 연산을 GPU에서 빠르게 수행할 수 있게 해주지만, NVIDIA의 cuBLASLt 라이브러리에서 제공하는 모든 세부 파라미터에 접근할 수는 없습니다. PyTorch 역시 강력한 텐서 연산 기능과 딥러닝 생태계를 갖추고 있지만, GPU 라이브러리의 저수준 기능을 모두 노출하지는 않습니다. 반면 nvmath-python은 이러한 세밀한 제어를 가능하게 하여, 연구자들이 라이브러리의 성능을 극한까지 끌어낼 수 있게 돕습니다.
또한 nvmath-python은 기존 프레임워크(CuPy, PyTorch, NumPy)에서 생성된 텐서를 그대로 받아들일 수 있어, 완전히 새로운 데이터 구조를 배워야 하거나 코드를 크게 수정할 필요가 없습니다. 즉, 기존의 워크플로우는 그대로 유지하면서, 더 높은 수준의 제어와 성능 최적화를 가능하게 합니다.
nvmath-python의 주요 기능
- 고급 행렬 연산(Matmul API): nvmath-python은 상태 기반(Stateful) Matmul 객체를 제공합니다. 이를 활용하면 같은 연산을 여러 번 수행할 때 매번 새로운 계획을 세우지 않고, 한 번의 최적화된 계획을 재사용할 수 있습니다. 예를 들어 mixed precision을 활용한 행렬 곱셈을 설정하고, epilog 단계에서 bias를 바로 추가하는 식으로, 기존보다 훨씬 세밀한 연산 구성이 가능합니다.
이 과정은 기존의 Python API에서 접근할 수 없는 cuBLASLt의 고급 기능을 활용할 수 있게 해주며, 특히 딥러닝 모델 학습·추론 과정에서 GPU 활용도를 크게 높일 수 있습니다.
-
디바이스 사이드(Device-side) API: nvmath-python은 GPU 디바이스 코드 내부에서 직접 NVIDIA 라이브러리 호출을 할 수 있는 기능도 제공합니다. 예를 들어, numba의 CUDA JIT 함수 안에서 cuFFT 연산을 호출하여 FFT 기반 컨볼루션을 직접 구현할 수 있습니다. 이는 GPU 상에서 실행되는 커널 내부에 고성능 연산을 삽입할 수 있음을 의미하며, 고급 사용자에게는 매우 유용한 기능입니다.
-
FFT 프로로그/에필로그 커스터마이징: 기존의 FFT 연산에 단순히 변환만 하는 것이 아니라, 사용자 정의 epilog를 작성하여 LTO-IR로 컴파일한 후 바로 실행 결과에 적용할 수 있습니다. 예를 들어, 단위 정규화(unitary normalization)를 위해 FFT 결과를 자동으로 1/sqrt(N)으로 스케일링하는 함수를 정의할 수 있습니다. 이러한 기능은 신호 처리나 푸리에 기반 계산에서 매우 자주 필요한 옵션을 손쉽게 통합할 수 있게 합니다.
라이선스
nvmath-python 프로젝트는 Apache 2.0 라이선스로 공개 및 배포되고 있습니다. 따라서 상업적 사용도 허용되며, 일부 조건(저작권 고지 유지 등)을 지키는 한 자유롭게 사용할 수 있습니다.
nvmath-python 공식 문서
https://docs.nvidia.com/cuda/nvmath-python/
nvmath-python 프로젝트 GitHub 저장소
더 읽어보기
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. 
파이토치 한국 사용자 모임
이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일
로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)
아래
쪽에 좋아요
를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ 