카카오브레인, PyTorch 호환 딥러닝 성능 라이브러리 'Trident' 공개

카카오브레인 ML Optimization팀이 딥러닝 모델의 훈련과 추론 속도를 높이기 위한 성능 라이브러리 Trident를 오픈소스로 공개했다. 이 라이브러리는 OpenAI Triton을 기반으로 작성되어 PyTorch와 동일한 추상화 계층을 제공한다.

PyTorch와의 호환성 및 사용법

Trident는 기존 PyTorch 코드와 쉽게 교체하거나 혼합하여 사용할 수 있도록 설계되었다. 사용자가 nn 객체를 trident 객체로 변경하는 것만으로 적용 가능하며, InstanceNorm1d, Dropout, Linear 등 주요 모듈을 대체할 수 있다. 이러한 구조 덕분에 기존 모델의 아키텍처를 크게 수정하지 않고도 성능 최적화를 수행할 수 있다.

성능 개선 효과

벤치마크 결과, Trident는 Linear, RMSNorm, ShiftGELU 등 주요 연산에서 PyTorch보다 낮은 소요 시간을 기록했다. 이를 통해 동일한 하드웨어 환경에서 모델의 훈련 및 추론 속도를 약 15% 수준으로 향상시킬 수 있다. 카카오브레인은 이 라이브러리를 자사의 Foundation 모델 개발 효율성을 높이기 위해 개발했으며, 기술 공유와 협력을 위해 오픈소스로 배포했다.