苹果公司最近在训练 AI 模型时，选择了谷歌的 TPU（Tensor Processing Unit）而不是传统的英伟达 GPU。自 2015 年推出以来，谷歌的 TPU 在性能和成本效益上进行了优化，现在已经在 AI 训练领域占据了重要地位。根据苹果发布的技术论文，他们采用了谷歌的 TPUv4 和 TPUv5 芯片进行 AI 模型训练。具体来说，苹果在服务器端 AI 模型训练中使用了 8192 个 TPUv4 芯片，而在设备端 AI 模型训练中则用了 2048 个 TPUv5 芯片。
TPU 在处理大规模数据训练时可以显著缩短训练时间，提高整体模型性能。TPU 的设计初衷是在能效和性能之间找到最佳平衡，使其在大规模神经网络计算任务中表现得游刃有余。根据苹果的论文，TPU 的高效计算能力在生成式 AI 模型的训练和推理方面显著提升了训练速度和推理性能。TPU v4 和 TPU v5 在计算能力和能效上都有提升，支持更复杂和大规模的 AI 模型训练。
TPU 的核心组件是矩阵乘法器单元（MXU），能够高效执行矩阵运算，提升计算吞吐量。和传统的 CPU、GPU 相比，TPU 在处理特定类型的计算任务时显得更为高效。它的设计目标是提供高效的计算能力，以支持大规模的机器学习模型。在苹果的 AI 模型训练中，TPU 发挥了关键作用，尤其是在 AFM-server 和 AFM-on-device 模型的训练中。服务器端 AI 模型 AFM-server 用了 8192 个 TPU v4 芯片进行训练，而设备端 AI 模型 AFM-on-device 则用了 2048 个 TPU v5 芯片。通过使用 TPU，苹果能够更快地完成 AI 模型的训练，提高模型的性能和效率。在训练过程中，AFM-server 和 AFM-on-device 模型分别处理了 6.3 万亿 tokens 的数据。
谷歌通过其 Cloud TPU 服务，将这种能力提供给更多的企业和开发者，进一步扩大了 TPU 的市场影响力。不仅是苹果，谷歌、Meta 等科技公司也在广泛采用 TPU 进行 AI 模型训练。谷歌自身在其 AI 服务中大量使用 TPU，并通过 Google Cloud 向外部客户提供 TPU 服务。Meta 等公司也看中了 TPU 的高效计算能力，纷纷将其应用于自家的 AI 项目中。