BobMcDear/attorch

해결하는 문제

attorch는 PyTorch의 nn 모듈을 대체할 수 있는 읽기 쉽고, 독립적이며, 수정이 용이한 대안을 제공합니다. 개발자는 원시 CUDA 커널을 작성하는 데 필요한 기술적 전문 지식 없이도 Triton 기반 커널의 효율성을 갖춘 커스텀 딥러닝 연산을 생성할 수 있습니다.

작동 방식

이 프로젝트는 OpenAI의 Triton을 사용하여 순수 Python으로 작성되었습니다. 다양한 신경망 레이어(예: convolution, attention 및 다양한 활성화 함수)와 수학 함수를 구현합니다. 순전파(forward)와 역전파(backward)를 모두 지원하므로 학습과 추론 모두에 적합합니다. 통합을 단순화하기 위해 특정 attorch 모듈을 사용할 수 없거나 convolution 및 pooling과 같이 성능이 중요한 레이어가 PyTorch에 의해 더 잘 처리되는 경우 자동으로 PyTorch 레이어를 사용하는 폴백 메커니즘(attorch.nn)이 포함되어 있습니다.

대상

커스텀 딥러닝 연산을 개발하고자 하며, 순수 PyTorch 구현보다 성능 향상이 필요하지만 CUDA 커널을 작성할 수 없는 연구자와 개발자를 위해 설계되었습니다.

주요 특징

• Triton 기반: 효율적인 GPU 계산을 위해 OpenAI의 Triton을 사용합니다.
• Cuda 불필요: CUDA 코드를 작성하지 않고도 고성능 커널을 구현할 수 있습니다.
• 폭넓은 레이어 지원: MultiheadAttention, BatchNorm 및 선택적 커널 퓨전(kernel fusion) 기능이 있는 다양한 활성화 함수를 포함한 광범위한 레이어를 포함합니다.
• 학습 및 추론: 순전파와 역전파를 모두 완벽하게 지원합니다.
• PyTorch 통합: 기존 PyTorch 워크플로우와의 호환성을 보장하기 위한 폴백 인터페이스를 제공합니다.