Kakao Brain Releases 'Trident', a PyTorch-Compatible Deep Learning Performance Library

The ML Optimization team at Kakao Brain has open-sourced Trident, a performance library designed to accelerate the training and inference of deep learning models. Built on OpenAI Triton, this library provides an abstraction layer identical to PyTorch.

Compatibility and Usage with PyTorch

Trident is designed to be easily swapped with or mixed into existing PyTorch code. It can be applied simply by changing nn objects to trident objects, allowing it to replace key modules such as InstanceNorm1d, Dropout, and Linear. Thanks to this structure, performance optimization can be performed without significantly modifying the architecture of existing models.

Performance Improvements

Benchmark results show that Trident recorded lower execution times than PyTorch for key operations such as Linear, RMSNorm, and ShiftGELU. This enables an improvement in model training and inference speeds of approximately 15% on the same hardware environment. Kakao Brain developed this library to enhance the efficiency of its Foundation model development and distributed it as open source for technology sharing and collaboration.