tensorgi/TPA

해결하는 문제

TPA (Tensor ProducT ATTenTion)는 표준 트랜스포머 모델의 효율성과 성능 저하 문제를 해결합니다. 구체적으로, 추론 중 주요 메모리 병목 현상인 KV(키-값) 캐시의 크기를 줄이면서 모델 성능을 향상시키는 것을 목표로 합니다.

작동 방식

이 프로젝트는 표준 어텐션 메커니즘을 텐서 곱 어텐션으로 대체하는 T6 아키텍처를 구현합니다. 생성 과정을 최적화하기 위해 특수한 디코딩 알고리즘인 FlashTPA를 포함합니다. 구현은 nanoGPT 코드베이스를 기반으로 하며, PyTorch를 사용하여 다중 GPU에 걸쳐 분산 학습을 지원합니다.

대상 사용자

어텐션 메커니즘 개선과 LLM 추론 시 메모리 오버헤드 감소에 초점을 맞춘 트랜스포머 아키텍처를 연구하거나 개발하는 연구자 및 개발자들입니다.

주요 특징

• 텐서 곱 어텐션: 더 나은 성능과 효율성을 위한 새로운 어텐션 메커니즘.
• FlashTPA 디코딩: 더 빠르고 메모리 효율적인 디코딩을 위한 최적화된 Python 및 Triton 구현.
• 확장성: Fineweb-Edu-100B 및 OpenWebText과 같은 대규모 데이터셋에 최적화되어 있습니다.
• 표준화된 평가: 성능을 벤치마킹하기 위해 lm-evaluation-harness와 통합되어 있습니다.