CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

MIT, Princeton, Meta, Together AI의 연구진은 메모리 병목이 발생하는 Transformer 연산을 일반 행렬 곱셈(GEMM)의 에필로그 단계 안에서 실행되도록 재매개변수화하는 방법인 CODA를 개발했다.

Transformer 학습에서는 정규화와 활성화, 잔차 갱신처럼 계산은 적지만 큰 중간 텐서를 메모리로 옮기는 연산이 적지 않은 시간을 차지한다. CODA는 이런 연산을 별도 커널로 실행하지 않고 행렬 곱셈 결과 타일이 칩 안에 남아 있을 때 에필로그 단계에서 처리하도록 대수적으로 다시 표현한다. 행렬 곱셈 본체는 고정하고 크기 조절, 축약, 쌍별 변환, 누적을 조합하는 작은 에필로그 명령 집합을 제공한다. 이 제한된 구조로 표준 Transformer 블록의 순방향과 역방향에서 어텐션을 제외한 거의 모든 계산을 표현할 수 있다. 대표 작업에서 사람이 작성한 커널과 언어 모델이 작성한 커널 모두 높은 성능을 보였지만, 논문의 적용 범위는 어텐션 자체보다 행렬 곱셈 주변의 메모리 병목 연산에 맞춰져 있다.

핵심 요약

• 상태 유지 변환: 교차 엔트로피 손실의 온라인 log-sum-exp 같은 연산을 위해 로컬 상태를 유지한다.
• 효율성: 데이터 이동으로 손실되는 GPU 시간의 상당 부분을 회복한다.
• 프로그래밍 용이성: 구조화된 DSL을 제공해 고성능 커널 작성을 간소화한다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.