GFormer: Accelerating Large Language Models with Optimized Transformers on Gaudi Processors (2412.19829v1)

Abstract: Heterogeneous hardware like Gaudi processor has been developed to enhance computations, especially matrix operations for Transformer-based LLMs for generative AI tasks. However, our analysis indicates that Transformers are not fully optimized on such emerging hardware, primarily due to inadequate optimizations in non-matrix computational kernels like Softmax and in heterogeneous resource utilization, particularly when processing long sequences. To address these issues, we propose an integrated approach (called GFormer) that merges sparse and linear attention mechanisms. GFormer aims to maximize the computational capabilities of the Gaudi processor's Matrix Multiplication Engine (MME) and Tensor Processing Cores (TPC) without compromising model quality. GFormer includes a windowed self-attention kernel and an efficient outer product kernel for causal linear attention, aiming to optimize LLM inference on Gaudi processors. Evaluation shows that GFormer significantly improves efficiency and model performance across various tasks on the Gaudi processor and outperforms state-of-the-art GPUs.

• The paper introduces GFormer, a framework optimizing Transformer attention by integrating sparse and linear mechanisms to leverage Gaudi processor architecture.
• GFormer achieves significant speedups for GPT and ViT models on Gaudi processors, outperforming state-of-the-art GPU implementations for longer sequences.
• The work suggests GFormer enables cost-effective LLM deployment on Gaudi hardware and is potentially applicable to other accelerators.

Accelerating LLMs with GFormer on Gaudi Processors

The paper "GFormer: Accelerating LLMs with Optimized Transformers on Gaudi Processors" provides a detailed exploration of advancing the efficiency and performance of Transformer-based LLMs using specialized hardware. The core innovation, named GFormer, introduces a framework that optimizes the interplay between sparse and linear attention mechanisms to maximize utilization of the Gaudi processor's heterogeneous compute architecture.

Context and Challenges

Transformers have become pivotal in numerous NLP tasks, yet their computational efficiency limits broader adoption, primarily due to the self-attention mechanism's quadratic complexity in both time and memory. The pursuit of more comprehensive models demands processing of longer sequences, exacerbating these constraints. Specialized accelerators such as Intel's Gaudi processors offer a promising solution, leveraging a hybrid architecture combining Matrix Multiplication Engines (MMEs) and Tensor Processing Cores (TPCs). Despite notably accelerating deep learning workloads through dense matrix operations, these processors face bottlenecks with non-matrix operations like Softmax when handling extended sequences.

GFormer Methodology

GFormer addresses these challenges by integrating sparse and linear attention strategies to enhance computational efficiency while maintaining model integrity.

• Sparse and Linear Attention Integration: GFormer splits attention heads between sparse and linear mechanisms. Sparse attention is optimized for TPC through windowed local-context kernels, facilitating efficient data access patterns. Linear attention, through orthogonal projection of Softmax, aligns with MME's capabilities, utilizing dense matrix operations efficiently.
• TPC Kernels Optimization: Specific TPC kernels are developed, including a windowed self-attention kernel and an efficient outer product kernel. These are designed to optimally balance MME and TPC workloads, with benchmarks showing effective parallel processing capabilities and significant speedup relative to Softmax operations.
• Optimal Workload Partitioning: A performance modeling technique ensures balanced TPC and MME resource allocation. By estimating and equalizing runtime workloads, GFormer effectively leverages the heterogeneous architecture.

Experimental Evaluation

The paper demonstrates GFormer’s effectiveness through evaluations on GPT and Vision Transformer (ViT) models, achieving up to 2x and 2.2x speedups respectively over baseline Softmax attention models. Moreover, GFormer on Gaudi processors outperforms state-of-the-art GPU implementations in handling larger sequence lengths, providing up to a 1.5x speedup in such scenarios.

Implications and Future Directions

Practically, this work implies a pathway for employing Gaudi processors in large-scale LLM deployments with cost-effective computation compared to GPUs. Theoretically, it underscores the importance of tailoring deep learning architectures to hardware capabilities, particularly in addressing bottlenecks inherent to transformers like Softmax processing.

Future research could focus on scaling this architecture across multiple Gaudi units to explore distributed acceleration of LLMs. Additionally, there is potential to extend the mixed attention mechanism across other heterogeneous AI accelerators, potentially broadening the applicability of these optimizations beyond the current context.

In conclusion, GFormer demonstrates a significant advancement in optimally marrying algorithmic improvements with hardware-specific capabilities, paving the way for more scalable and efficient transformer models on emerging computational platforms.