facebookincubator/AITemplate

해결하는 문제

AITemplate은 NVIDIA 및 AMD GPU에서 깊은 신경망의 추론 속도를 최대한 높이기 위해 설계되었습니다. 일반적인 런타임의 성능 병목을 해결하기 위해 모델을 고도로 최적화된 자체 완전한 C++ 코드로 변환하여, 하드웨어의 이론적 최대 성능(루프라인)에 근접한 추론 서빙을 가능하게 합니다.

작동 방식

이 프레임워크는 신경망 그래프를 CUDA 또는 HIP C++ 코드로 변환합니다. 메모리 오버헤드를 줄이고 GPU 활용도를 높이기 위해 여러 고급 융합 기술을 사용합니다:

• 수평 융합: 서로 다른 입력 형태를 가진 병렬 연산자(GEMM 또는 LayerNorm 등)를 하나의 커널로 통합합니다.
• 수직 융합: 요소 단위 연산 및 축소 연산과 같은 범위의 연산을 직접 TensorCore 또는 MatrixCore 연산으로 통합합니다.
• 메모리 융합: 연결, 분할, 슬라이스와 같은 메모리 연산을 그 전에 오는 연산자에 통합합니다.

FX2AIT 라는 도구를 포함하며, PyTorch 모델을 AITemplate 엔진으로 변환할 수 있습니다. 모델에 AITemplate에서 아직 지원되지 않는 연산자가 포함된 경우, FX2AIT는 모델을 분할하여 부분적인 가속을 수행할 수 있습니다.

대상 사용자

Ampere 세대 NVIDIA GPU(SM80+) 또는 CDNA2 AMD GPU(MI-210/250)에서 극한의 추론 성능이 필요한 개발자 및 엔지니어. 특히 BERT, Stable Diffusion, Vision Transformers와 같은 모델에 적합합니다.

주요 특징

• 하드웨어 독립성: NVIDIA 및 AMD GPU 플랫폼 모두에서 작동합니다.
• 의존성 없음: 컴파일된 모델은 자체 완전한 바이너리이며, 실행 시 cuBLAS, cuDNN, TensorRT와 같은 제3자 라이브러리가 필요하지 않습니다.
• PyTorch 통합: 추가 메모리 복사 없이 PyTorch 텐서를 입력 및 출력으로 사용할 수 있지만, PyTorch가 완전히 없는 환경에서도 실행 가능합니다.
• 확장성: Python에서 그래프 노드와 백엔드 코드 생성을 정의함으로써 새로운 연산자 또는 융합 커널을 추가할 수 있습니다.