Kraken: Inherently Parallel Transformers For Efficient Multi-Device Inference
This work addresses the challenge of efficient multi-device inference for large Transformers, which is crucial for improving end-user experience and enabling new applications, though it is incremental as it builds on existing tensor parallelism schemes.
The paper tackles the problem of high inference latency in large Transformer networks by introducing Kraken, an architecture designed to overlap collective communication with compute, which reduces Time To First Token by a mean of 35.6% on multi-GPU systems while maintaining similar perplexity and language modeling capabilities as standard Transformers.
Large Transformer networks are increasingly used in settings where low inference latency can improve the end-user experience and enable new applications. However, autoregressive inference is resource intensive and requires parallelism for efficiency. Parallelism introduces collective communication that is both expensive and represents a phase when hardware resources are underutilized. Towards mitigating this, Kraken is an evolution of the standard Transformer architecture that is designed to complement existing tensor parallelism schemes for efficient inference on multi-device systems. By introducing a fixed degree of intra-layer model parallelism, the architecture allows collective operations to be overlapped with compute, decreasing latency and increasing hardware utilization. When trained on OpenWebText, Kraken models reach a similar perplexity as standard Transformers while also preserving their language modeling capabilities when evaluated on the SuperGLUE benchmark. Importantly, when tested on multi-GPU systems using TensorRT-LLM engines, Kraken speeds up Time To First Token by a mean of 35.6% across a range of model sizes, context lengths, and degrees of tensor parallelism.