콘텐츠로 이동

PyTorch RBLN: RBLN NPU를 위한 PyTorch 확장

Warning

PyTorch RBLN은 현재 beta 상태이고 변경될 수 있습니다.

PyTorch RBLN(torch-rbln)은 Rebellions NPU의 연산 성능을 PyTorch에서 자연스럽게 사용할 수 있도록 해주는 PyTorch Extension입니다. 'define-by-run'으로 동작하는 PyTorch의 기본 실행 모드인 eager mode를 구현하여, PyTorch 기반 생태계에서 모델 개발, 배포, 서빙의 전 단계를 아우르는 엔드 투 엔드(End-to-End) 워크플로우에서 필수적인 사용자 경험을 제공합니다. 또한 다양한 디버깅 상황에서도 편리하게 사용할 수 있습니다.

PyTorch RBLN은 GPU 및 CPU와 동일한 인터페이스를 제공하며 PyTorch 환경에서 Rebellions NPU를 자연스럽고 일관되게 활용할 수 있도록 지원합니다.

설치 방법

torch-rbln을 설치하기 위해 RBLN Portal 계정이 필요합니다. 다음은 pip를 통해 최신 릴리즈를 설치하는 명령어입니다.

1
$ pip3 install -i https://pypi.rbln.ai/simple/ torch-rbln

최신 버전 및 변경 사항에 대해 자세한 내용은 릴리즈 노트를 참조하세요.

Note

torch-rbln은 out-of-tree extension입니다. 그러므로 torch-rbln을 설치하면 torch가 자동으로 종속성으로 함께 설치됩니다.

디자인 개요

PyTorch의 torch.add 또는 torch.mul과 같은 연산을 실행할 때, 내부적으로는 aten::add나 aten::mul과 같은 PyTorch ATen 연산자에 바인딩됩니다. 디스패처를 사용하여 이러한 ATen 연산은 CPU, CUDA 및 XLA 등 다양한 백엔드의 구현으로 연결됩니다. 이를 위해 PyTorch 디스패처는 주로 입력 텐서의 장치 속성을 이용해 어떤 백엔드가 해당 연산을 제공하는지 찾습니다. PyTorch의 백엔드 중 하나인 PrivateUse1 디스패치 키는 다양한 장치를 지원하기 위해 Out-of-Tree (OOT) 확장에 의해 사용됩니다. PyTorch RBLN은 OOT를 사용하여 RBLN 장치별 런타임과 연산 구현을 PyTorch에 등록합니다. 따라서 rbln 장치 속성을 지정한 입력 텐서에 연산자를 적용하면 PrivateUse1 디스패치 키를 통해 RBLN 연산 구현으로 디스패치됩니다. 이 연산 구현은 주로 torch.compile에서 제공되며, 이는 RBLN 컴파일러를 통해 RBLN 장치용 실행 가능 파일을 생성하고 실행하기 위한 PyTorch 인터페이스입니다.

PyTorch RBLN의 구성 요소

PyTorch RBLN은 다음 그림과 같이 크게 PyTorch, torch-rbln, RBLN 컴파일러 및 런타임, RBLN 드라이버의 네 가지 요소로 구성되어 있습니다.

image

  • PyTorch (torch)
    • libtorch_python.so: Python과 PyTorch C++ backend 사이의 인터페이스
    • libtorch.so: PyTorch의 메인 라이브러리로 PyTorch C++ API를 제공
    • libtorch_cpu.so: CPU 전용으로 설계된 LibTorch backend 라이브러리
    • libc10.so: PyTorch의 저수준 유틸리티 및 기본 구조를 제공하는 라이브러리
      • 텐서 관리, 디바이스 추상화, 메모리 할당 등
  • PyTorch RBLN (torch-rbln)
    • libtorch_rbln.so: ATen op C++ 라이브러리 (RBLN 런타임 API를 이용하는 copy, resize등)
    • register_ops.py: 리벨리온 NPU에서 동작하는 python op 구현체 (torch.compile 이용)
    • libc10_rbln.so: libc10.so를 기반으로 하는 리벨리온 NPU의 PyTorch runtime components 라이브러리
  • RBLN 컴파일러 (rebel-compiler)
    • librbln.so: 리벨리온 컴파일러와 런타임 라이브러리
  • RBLN 드라이버
    • librbln-thunk.so: 리벨리온 NPU 디바이스 드라이버

지원하는 연산

지원하는 연산은 여기에서 확인할 수 있습니다. PyTorch RBLN에 의해서 지원되는 연산은 RBLN SDK가 업데이트됨에 따라서 점차 늘어날 예정입니다.

튜토리얼

PyTorch RBLN을 쉽게 시작할 수 있도록, 기본적인 기능을 보여주는 튜토리얼을 제공합니다.