The Future of AI Chips: Insights from Google TPU v4 and Chinese AI Chip Unicorn

The Future of AI Chips: Insights from Google TPU v4 and Chinese AI Chip Unicorn

Introduction:
In recent years, advancements in artificial intelligence (AI) have led to the development of specialized AI chips that are revolutionizing various industries. In this article, we will explore the future of AI chips by examining the features and optimizations of the Google TPU v4 and a Chinese AI chip unicorn. We will also discuss the common points between these two chips and the implications they have for the AI industry.

The Optimization of Google TPU v4:
Google's TPU v4 is designed to optimize embedded layers, which are crucial for recommendation systems. To achieve this, Google has incorporated a dedicated acceleration module called SparseCore (SC) in each TPU v4. Each SC consists of its own vector processing unit (scVPU), 2.5 MB of local SRAM, and a memory access interface that can access up to 128TB of shared 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 significant number of SCs are deployed in each TPU v4, occupying around 5% of the overall chip area and power consumption. Comparisons between running embedding layers on CPUs and TPU v4 SCs have shown speed improvements of over 6 times when utilizing the SCs. This highlights the significant performance gains achieved through domain-specific design, even with minimal chip area and power consumption.

Optimizing Data Flow with Reconfigurable Optical Interconnects:
Different machine learning models have varying data flow requirements, broadly categorized into data parallelism, model parallelism, and pipeline parallelism. Each of these data flows corresponds to different TPU interconnection topologies. With reconfigurable optical interconnects, the interconnection topology between TPUs can be adjusted based on the specific model's data flow, resulting in optimal performance gains of over 2 times. This flexibility in data flow optimization is a significant advantage provided by reconfigurable optical interconnects.

Enhanced Reliability through Reconfigurable Optical Interconnects:
In large-scale supercomputers composed of a massive number of chips, ensuring high performance even in the presence of a small percentage of faulty chips is crucial. Conventional fixed interconnect architectures can lead to system-wide performance degradation if a single chip fails. However, with reconfigurable optical interconnects, faulty chips can be bypassed, maintaining overall system performance with only a slight sacrifice. Google's research indicates that by employing reconfigurable optical interconnects and optical circuit switches, an average performance improvement of up to 6 times can be achieved, assuming a chip reliability rate of 99%. This showcases the importance of optical interconnect switches in maintaining reliability and performance in large-scale AI chip systems.

The Focus on Large Models and Embedding Layers:
While the initial focus of AI chips like the TPU v1 was on accelerating convolutional neural networks (CNNs), the current emphasis has shifted towards addressing the challenges posed by large models. For Google, the primary large model of utmost importance is the recommendation system, which determines the company's main source of revenue. Consequently, the TPU v4 has been optimized to accelerate the embedding layer, a bottleneck in recommendation system models. The embedding layer maps high-dimensional sparse features to low-dimensional dense features, facilitating further processing by neural networks. Typically implemented using lookup tables, these tables can be massive, reaching sizes of up to 100GB. Distributed computing is required to distribute each embedding layer's lookup table across multiple TPU v4 chips. Google's research highlights that embedding layer computations primarily involve 1D vector computations, which are often sparse and distributed across different chips. This necessitates efficient optimization of shared memory for effective data exchange between chips.

Commonalities and the Future of AI Chip Design:
Different chip companies, such as AMD, Nvidia, and Google, focus on various aspects of AI chip design. However, they share a common goal of supporting AI chip scalability to meet the demands of large models. The future of AI chip design may involve scalability metrics, such as data interconnect bandwidth, becoming as important as peak computing power. This trend indicates that AI chip design requires expertise not only in digital logic and computer architecture but also in packaging and data interconnect domains.

Conclusion:
The future of AI chips is evolving rapidly, driven by advancements in optimization techniques, scalability, and reliability. The Google TPU v4 and a Chinese AI chip unicorn exemplify the innovative approaches taken to enhance performance and address the challenges posed by large models. As the demand for AI continues to grow, it is crucial for chip manufacturers to prioritize optimization, scalability, and reliability. To capitalize on this, here are three actionable advice for AI chip designers:

1. Emphasize domain-specific design: Identifying critical bottlenecks in specific applications, such as recommendation systems, and optimizing chip architectures accordingly can yield significant performance gains.

2. Invest in reconfigurable interconnect technologies: Reconfigurable optical interconnects enable flexible data flow optimization and enhanced reliability, allowing for efficient scaling of AI chip systems.

3. Foster cross-domain expertise: AI chip design requires expertise not only in digital logic and computer architecture but also in packaging and data interconnect domains. Encouraging interdisciplinary collaboration can lead to comprehensive solutions that address the challenges of large models and performance scalability.

By leveraging these insights and embracing future advancements, AI chip manufacturers can shape the future of AI technology and drive innovation across industries.