Apple has decided to employ Google’s Tensor Processing Units (TPUs) rather than NVIDIA GPUs for training its key AI language models designed to power innovative AI functionalities. The integration was highlighted in a detailed technical paper from Apple, which provided insight into the models and the development process involved.
Infrastructure and Technical Specifics
At WWDC 2024, Apple introduced the Apple Intelligence suite, featuring AI enhancements across iOS, iPadOS, and macOS. These enhancements leverage two main models: a smaller model with 3 billion parameters optimized for performance on individual devices, and a much larger model intended for server-side processing in Private Cloud Compute. Both of these models were trained using Google Cloud TPU v4 and v5p clusters, diverging from the typically favored NVIDIA H100 GPUs.
According to the technical report, the server model training utilized 8192 TPUv4 chips arranged in 8 clusters of 1024 chips each, linked via the data-center network (DCN). In contrast, the on-device model was trained on one cluster of 2048 TPUv5p chips. The Cloud TPU v4, launched in 2023, delivers almost ten times the performance of its predecessor, the TPU v3. Meanwhile, the newer TPU v5p offers more than double the FLOPS and three times the high-bandwidth memory (HBM) in comparison to the TPU v4.
Factors Influencing Apple’s Decision
The switch to Google TPUs could be attributed to the scarcity of NVIDIA H100s and the cost benefits of Google’s TPUs. This shift not only indicates the robustness of Google’s TPU technology but also acts as a marketing booster for Google Cloud. The move reflects the competitive dynamics in the AI hardware sector, where performance, cost, and availability play pivotal roles.
Apple’s foundational models are pre-trained using the AXLearn framework and a JAX-based deep learning library to enhance user interactions, refine text, and summarize notifications, among other functionalities. Incorporating these models into Apple’s software ecosystem indicates a significant step forward in the company’s AI advancements.
Last Updated on November 7, 2024 3:27 pm CET