The Future of AI Chips: Insights from Google TPU v4 and the Next Phase of GPU Chip Startups

The Future of AI Chips: Insights from Google TPU v4 and the Next Phase of GPU Chip Startups

Introduction:
The advancement of artificial intelligence (AI) has led to the development of specialized chips like the Google TPU v4. These chips are designed to optimize various aspects of AI models, with a particular focus on embedding layers. This article explores the key features of TPU v4, its impact on recommendation systems, and the future of GPU chip startups.

Google TPU v4 and SparseCore:
To optimize embedding layers, Google incorporated a dedicated acceleration module called SparseCore (SC) in TPU v4. Each SC consists of its vector processing unit (scVPU), 2.5 MB local SRAM, and a memory access interface that can access up to 128TB of shared high-bandwidth memory (HBM). Additionally, SC includes specialized acceleration logic for embedding layer operations such as sorting, reducing, and concatenating. Despite the simplicity of each SC's structure, a large number of SCs are deployed in each TPU v4, accounting for approximately 5% of the TPU's overall area and power consumption. According to Google's research, utilizing TPU v4 SCs for running embedding layers in recommendation systems resulted in over a 6x improvement in overall system performance compared to running the layers on a CPU.

Different Data Flow Requirements and TPU Topologies:
Machine learning models have varying data flow requirements, broadly classified into three categories: data parallelism, model parallelism, and pipeline parallelism. Each data flow corresponds to a different TPU interconnect topology. With reconfigurable optical interconnects, the interconnect topology between TPUs can be adjusted based on the specific model's data flow, resulting in optimal performance improvements of over 2x.

Advantages of Reconfigurable Optical Interconnects:
Reconfigurable optical interconnects offer enhanced reliability in supercomputers composed of a massive number of chips. In traditional fixed interconnect architectures, a single faulty chip can impact the entire system's functionality. However, with reconfigurable optical interconnects, the faulty chip can be bypassed, ensuring the overall system continues to function with minimal performance degradation. Google's research demonstrates that with a chip reliability rate of 99%, the average system performance improvement achieved using reconfigurable optical interconnects and optical circuit switches (OCS) can reach up to 6x.

The Importance of Large Model Optimization:
While the initial focus of TPU v1 in 2017 was on accelerating convolutional neural networks (CNNs), the landscape has shifted towards optimizing large models. For Google, the most critical large model is the recommendation system, which determines the company's primary revenue stream. Therefore, TPU v4 is designed to address the acceleration bottleneck in the embedding layer of recommendation system models. The embedding layer's purpose is to map high-dimensional sparse features to low-dimensional dense features, enabling further processing by neural networks. Typically, the implementation of the embedding layer involves a lookup table, which can be extremely large, ranging up to the order of 100GB. In a recommendation system model, multiple such lookup tables can be distributed across multiple TPU v4 chips for computation. Google's research highlights that the computation in the embedding layer primarily involves 1D vector operations, sparse calculations, and efficient data exchange optimization for shared storage among different chips.

Common Focus of Chip Companies:
Different chip companies, such as AMD, Nvidia, and Google, have their own unique approaches to AI chip scalability. AMD focuses on micro-level chiplet usage for package-level scalability, Nvidia utilizes technologies like NvLink for scalability and performance improvement in multi-GPU setups, while Google directly designed an optical switch chip for massive TPU interconnection. However, a shared focus among these companies is the support for AI chip scalability to meet the demands of large models. In the future, the scalability of AI chips, such as data interconnect bandwidth, may become a primary indicator, similar to peak computing power. This cross-domain requirement in AI chip design necessitates expertise in digital logic, computer architecture, packaging, and data interconnect.

Conclusion:
The future of AI chips, exemplified by Google TPU v4, showcases the importance of domain-specific optimization, reconfigurable interconnects, and scalability for large models. As the AI landscape progresses, GPU chip startups must navigate the challenging second half of the game, where winning against competition becomes crucial. To thrive in this space, startups must focus on outperforming their competitors in terms of innovation, scalability, and meeting the demands of large AI models.

Actionable Advice:

1. Emphasize domain-specific optimization: Identify critical bottlenecks in specific AI applications and design specialized acceleration modules to address them effectively.
2. Invest in reconfigurable interconnects: Develop interconnect architectures that can be adjusted based on the data flow requirements of different AI models, enabling optimal performance improvements.
3. Prioritize scalability for large models: Recognize the increasing importance of handling large models in AI applications, and invest in solutions that can efficiently scale to meet the demands of these models.

In summary, the future of AI chips lies in their ability to optimize domain-specific tasks, incorporate reconfigurable interconnects, and provide scalability for large models. GPU chip startups must navigate the competitive landscape by focusing on innovation, scalability, and meeting the demands of the AI industry. By leveraging the insights from Google TPU v4 and understanding the evolving requirements, these startups can position themselves for success in the second half of the game.