跳转至

国外AI芯片

学习目标

  1. 了解Google TPU的发展历程、架构特点和脉动阵列原理
  2. 掌握Google TPU各代产品的技术规格和应用场景
  3. 理解Sparse Core的作用和工作原理
  4. 了解Tesla DOJO的架构特点
  5. 了解英特尔Gaudi芯片的设计理念

章节导言

在全球AI芯片市场中，除了NVIDIA这一主导者外，众多科技巨头和芯片厂商也在积极布局专用AI芯片。Google的TPU（张量处理单元）是其中的标杆产品，已经历四代演进，在Google内部大规模部署并对外提供服务。Tesla的DOJO超级计算机代表了自动驾驶AI计算的新方向。Intel的Gaudi系列则展现了传统芯片巨头在AI领域的布局。

本章将深入介绍这些国外AI芯片的架构设计、技术特点和应用场景。通过分析这些芯片的设计理念，读者将更全面地理解AI芯片的多样性和发展趋势。

4.1 Google TPU概述

4.1.1 TPU的诞生背景

Google TPU的诞生源于2006年的一个内部讨论。当时Google工程师开始思考：能否用GPU、FPGA或ASIC来加速数据中心的机器学习工作负载？然而，彼时Google数据中心的计算资源相对充裕，专门开发芯片的成本似乎不够经济。

2013年，情况发生了变化。Google的研究人员预测，如果用户每天通过语音搜索进行3分钟的DNN语音识别，Google数据中心的算力需要翻倍才能满足需求。仅仅依靠通用CPU来满足这种爆发式增长的成本极其昂贵。

在这种背景下，Google于2013年正式启动了TPU项目，目标是开发一款专门加速深度学习推理的ASIC芯片。从立项到大规模部署，Google仅用了15个月，这在芯片开发领域堪称奇迹。TPU项目的负责人Norm Jouppi表示："我们的首批交付的硅片无需进行任何错误修正或掩模的更改，这本身就是一个非凡的成就。"

4.1.2 TPU发展历程

Google TPU已经发展了四代主要产品：

型号 发布时间 主要特点 应用场景
TPU v1 2016 92 TOPS, INT8推理 数据中心推理
TPU v2 2017 180 TFLOPS, HBM, BF16 训练+推理
TPU v3 2018 420 TFLOPS, 128GB HBM 训练+推理
TPU v4 2021 275 TFLOPS, 3D Torus互联 训练+推理

此外，Google还推出了面向边缘和移动设备的Edge TPU和Pixel Neural Core。

4.1.3 TPU的定位

与NVIDIA GPU的通用并行计算不同，Google TPU从一开始就被定位为神经网络推理专用加速器。这种专注使得TPU能够在特定工作负载下实现极高的性能和能效。

TPU的核心设计哲学：

  1. 专用性：专为神经网络计算优化
  2. 低精度：支持INT8量化，牺牲精度换取性能
  3. 脉动阵列：高效矩阵运算架构
  4. 确定性：简化设计，追求稳定性能

4.2 TPU v1架构详解

4.2.1 TPU v1硬件架构

TPU v1是Google第一款商用TPU芯片，采用28nm工艺制造，主频700MHz，热设计功耗为40瓦。

核心组件：

  1. MXU（矩阵乘法单元）
  2. 256×256脉动阵列
  3. 65536个ALU
  4. 每个时钟周期可执行65536次INT8乘加运算

  5. 理论峰值：92 TOPS

  6. 累加器（Accumulators）

  7. 4MiB容量
  8. 存储MXU的输出结果

  9. 4096个256元素32位累加器

  10. 统一缓冲区（Unified Buffer）

  11. 28MiB SRAM
  12. 作为MXU的输入缓存

  13. 减少对外部内存的访问

  14. DDR3 DRAM

  15. 8GiB容量
  16. 用于存储模型权重
  17. 34 GB/s带宽

4.2.2 脉动阵列原理

脉动阵列（Systolic Array）是TPU的核心创新，也是理解TPU架构的关键。

什么是脉动阵列？

脉动阵列是一种并行计算架构，数据在阵列中"流动"并进行计算，每个处理单元（PE）只与相邻单元直接通信。这种架构的名称来源于希腊语"systole"（收缩），因为数据在阵列中的流动节奏类似于心脏泵血。

脉动阵列的工作原理

传统的CPU/GPU计算方式是：数据从内存加载到寄存器，执行计算，再存储回内存。每次计算都需要访问内存。

脉动阵列则不同：数据从阵列边缘流入，在流动过程中被多个PE复用，最终流出时完成计算。这种设计大大减少了对内存的访问次数。

以一个简化的2×2矩阵乘法为例：

T=1:  a[0,0]×b[0,0] → PE[0,0]
T=2:  a[0,0]×b[0,1], a[0,1]×b[1,0] → PE[0,1], PE[1,0]
T=3:  a[0,0]×b[0,2], a[0,1]×b[1,2], a[0,2]×b[2,0] → ...

数据在PE之间脉动式流动，同时参与多个计算。

TPU MXU的具体实现

TPU v1的MXU是一个256×256的脉动阵列：

  • 每次可处理256×256的矩阵块
  • 输入向量从顶部流入
  • 权重矩阵预加载到阵列中
  • 每个时钟周期完成256次乘加
  • 脉动频率：700MHz
           权重矩阵(256×256)
           ↓↓↓↓
    ← ← ← ← ← ← ← ←
    ← ← ← ← ← ← ← ←
输入向量 → PE PE PE PE →
(256)    ← ← ← ← ← ← ← ←
    ← ← ← ← ← ← ← ←
    ← ← ← ← ← ← ← ←
           ↓↓↓
        输出结果

4.2.3 TPU v1的量化技术

神经网络对计算精度的要求低于传统科学计算。TPU v1充分利用这一特点，引入INT8量化技术。

量化原理

将32位浮点权重和激活值转换为8位整数：

  • FP32: 1位符号 + 8位指数 + 23位尾数
  • INT8: 0-255（无符号）或-128~127（有符号）

量化过程： 

FP32值 → 缩放到INT8范围 → 存储为INT8

量化优势

  1. 内存节省：权重存储空间减少4倍
  2. 带宽提升：相同带宽传输4倍数据量
  3. 计算加速：INT8计算比FP32更快
  4. 能效提高：硬件更简单，更节能

4.2.4 TPU v1指令集

TPU v1使用CISC指令集，包含12条指令：

指令 功能
Read_Host_Memory 从CPU读取数据到Unified Buffer
Read_Weight 从DRAM读取权重到Weight FIFO
Matrix_Multiply 执行矩阵乘法/卷积
Activate 执行ReLU、Sigmoid等激活函数
Write_Host_Memory 将结果写回CPU内存

TPU采用四级流水线处理这些指令，通过指令重叠提高硬件利用率。

4.3 TPU v2与v3

4.3.1 TPU v2架构升级

TPU v2于2017年发布，相比v1有了重大架构升级。

主要改进

  1. HBM高带宽内存
  2. 16 GiB HBM
  3. 700 GB/s带宽

  4. 解决v1的内存带宽瓶颈

  5. BF16精度支持

  6. v1仅支持INT8推理
  7. v2新增BF16训练支持

  8. 平衡精度与性能

  9. 双芯片模块

  10. 4个TPU v2组成一个模块
  11. 180 TFLOPS算力

  12. 高带宽互联

  13. TPU Pod

  14. 256片TPU v2组成
  15. 11.5 PFLOPS算力
  16. 专用网络互联

BF16 vs FP16

BF16（Brain Float）是Google专门为深度学习设计的浮点格式：

格式 符号位 指数位 尾数位 动态范围
FP32 1 8 23 极大
FP16 1 5 10 有限
BF16 1 8 7 与FP32相当

BF16保留与FP32相同的指数位，确保相同的动态范围，只减少尾数位（精度）。这对深度学习训练来说非常合适。

4.3.2 TPU v3

TPU v3于2018年发布，相比v2性能提升显著。

规格对比

规格 TPU v2 TPU v3
芯片面积增加 - 6%
时钟频率提升 700 MHz 940 MHz
MXU数量 1 2
HBM容量 16 GiB 32 GiB
内存带宽 700 GB/s 900 GB/s
TFLOPS (BF16) 45 123
Pod规模 256芯片 1024芯片
Pod算力 12 PFLOPS 126 PFLOPS

2D Torus互联

TPU v3采用2D Torus互联拓扑：

  • 每个TPU与4个邻居直接相连
  • 支持4096芯片的全互联
  • 高带宽，低延迟
  • 适合大规模分布式训练

4.4 TPU v4

4.4.1 TPU v4主要特性

TPU v4于2021年发布，是Google迄今为止最强大的TPU。

架构创新

  1. 7nm工艺
  2. 从16nm升级到7nm

  3. 显著提升能效

  4. 3D Torus互联

  5. 从2D升级到3D Torus
  6. 更高维度的邻居连接

  7. 提升全互联带宽

  8. 更大规模

  9. 支持4096芯片互联
  10. Pod算力达1 exaFLOPS
  11. 满足大模型训练需求

规格参数

参数 TPU v4
工艺 7nm
芯片内存 32 GiB HBM
内存带宽 1.2 TB/s
BF16算力 275 TFLOPS
芯片间互联 3D Torus
Pod规模 4096芯片

4.4.2 3D Torus拓扑

3D Torus是TPU v4的关键创新。相比2D Torus，3D Torus中的每个节点可以与6个邻居直接相连（左右、上下、前后），而不是4个。

3D Torus优势

  1. 更高带宽
  2. 更多等距路径

  3. 负载均衡更好

  4. 更低延迟

  5. 跳数更少

  6. 路径选择更多

  7. 更强容错

  8. 绕过故障节点
  9. 可靠性更高

拓扑示意图

        ┌─────────────────────────────┐
       /│                            /│
      / │        3D Torus           / │
     /  │      每个节点与6个        /  │
    ┌─────────────────────────────┐   │
    │    │   邻居直接相连           │   │
    │────┼────────────────────────│───┤
    │    │                        │   │
    │    └────────────────────────│───┘
    │   /                        │  /
    │  /                         │ /
    │ /                          │/
    └─────────────────────────────┘

4.5 Edge TPU与Pixel Neural Core

4.5.1 Edge TPU

Edge TPU是Google面向边缘设备推出的AI加速器。

规格特点

参数 Edge TPU
算力 4 TOPS
能效 2 TOPS/W
功耗 ~2W
精度 INT8
尺寸 微型片上系统

应用场景

  • 智能手机
  • IoT设备
  • 嵌入式系统
  • 边缘推理

4.5.2 Pixel Neural Core

Pixel Neural Core是Google Pixel手机上专用的人工智能处理器。

技术特点

  • 基于Edge TPU架构
  • 针对手机拍照优化
  • 加速HDR+、人像模式等
  • 实时语义分割

4.6 Tesla DOJO

4.6.1 DOJO概述

Tesla DOJO是Tesla专门为自动驾驶训练设计的超级计算机系统。Tesla在2021年AI Day上首次公布了DOJO的详细设计。

设计目标

  • 训练自动驾驶视觉神经网络
  • 处理海量视频数据
  • 支持超大规模模型
  • 极致性能优化

4.6.2 DOJO硬件架构

D1芯片

DOJO的核心是Tesla自研的D1芯片：

参数 D1
工艺 7nm
晶体管 500亿
算力 362 TFLOPS (BF16)
带宽 10 TB/s
面积 645 mm²

D1特点

  • 无缓存设计
  • 专用编译器
  • 极大芯片尺寸
  • 面向训练优化

训练瓦片（Training Tile）

  • 25个D1芯片组成
  • 9 PFLOPS算力
  • 160 GB HBM
  • 专用互联网络

ExaPOD

  • 120个训练瓦片
  • 1.1 EFLOPS算力
  • 13 TB HBM
  • 是世界上最快的AI训练系统之一

4.6.3 DOJO软件栈

DOJO采用专用软件栈：

  1. Dojo Neural Network
  2. 分布式训练框架
  3. 自动微分

  4. 模型并行

  5. Dojo编译器

  6. 针对D1架构优化
  7. 自动化并行化

  8. 图优化

  9. Dojo Runtime

  10. 硬件抽象层
  11. 任务调度
  12. 通信管理

4.6.4 DOJO vs NVIDIA

特性 DOJO NVIDIA A100
定位 自动驾驶训练 通用AI训练
算力 362 TFLOPS 312 TFLOPS
内存带宽 10 TB/s 2 TB/s
芯片互联 专用 NVLink
编程方式 专用 CUDA

DOJO专为Tesla的视觉自动驾驶方案优化，在特定工作负载下性能优于通用GPU。

4.7 英特尔Gaudi

4.7.1 Gaudi概述

Intel Gaudi是Intel面向AI训练和推理的专用加速器。Intel在2019年收购了AI芯片公司Habana Labs后，将其Gaudi产品线纳入旗下。

产品定位

  • 云端AI训练
  • 高效推理
  • 与Intel Xeon CPU协同

4.7.2 Gaudi架构

Gaudi 2规格

参数 Gaudi 2
工艺 7nm
算力 (FP16) 400 TFLOPS
HBM容量 96 GB
HBM带宽 2.4 TB/s
RDMA带宽 400 Gb/s
TDP 350W

架构特点

  1. 张量处理核心
  2. 专门优化矩阵运算
  3. 支持FP16/BF16/FP32

  4. 高效卷积加速

  5. HBM2e高带宽内存

  6. 96 GB容量

  7. 满足大模型需求

  8. ROCO RDMA

  9. 远程直接内存访问
  10. 高效多芯片互联
  11. 降低集群通信开销

4.7.3 Gaudi软件生态

SynapseAI

Intel的AI软件栈：

  • 优化的深度学习框架
  • 算子库
  • 编译器优化

支持框架

  • TensorFlow
  • PyTorch
  • ONNX

与NVIDIA竞争

Intel Gaudi定位与NVIDIA A100/H100竞争：

特性 Gaudi 2 H100 SXM
FP16算力 400 TFLOPS 1979 TFLOPS
HBM 96 GB 80 GB
互联带宽 400 Gb/s 900 GB/s
生态成熟度 较低 极强

虽然性能存在差距，但Gaudi在性价比和特定场景下有优势。

4.8 其他国外AI芯片

4.8.1 Graphcore IPU

Graphcore IPU是英国芯片公司的产品，采用独特的多指令多数据（MIMD）架构。

IPU特点

  • 大规模并行（1472核）
  • 903 GB/s带宽
  • IPU-POD跨芯片互联
  • Poplar软件栈

4.8.2 Cerebras WSE

Cerebras Wafer Scale Engine是世界上最大的芯片：

  • 850,000个AI核心
  • 46,225 mm²面积
  • 2.6万亿晶体管
  • 晶圆级集成

4.8.3 Groq

Groq是一家专注于推理的芯片公司：

  • 单一芯片架构
  • 确定性执行
  • 极致推理性能
  • TSP（张量流处理器）

本章小结

本章介绍了国外主要AI芯片：

  1. Google TPU：
  2. 从v1到v4的四代演进
  3. 脉动阵列是核心创新
  4. 量化技术实现高效推理

  5. BF16精度支持训练

  6. Tesla DOJO：

  7. 专为自动驾驶设计
  8. 自研D1芯片
  9. 分布式训练优化

  10. 软件栈垂直整合

  11. Intel Gaudi：

  12. 面向云端AI训练
  13. 高带宽互联

  14. 差异化竞争策略

  15. 其他AI芯片：

  16. Graphcore IPU
  17. Cerebras WSE
  18. Groq TSP

这些芯片代表了AI芯片的不同技术路线和应用场景。

思考与练习

  1. TPU分析：分析Google TPU采用的脉动阵列与NVIDIA GPU采用的SIMD/SIMT架构的核心区别，以及各自的优势和局限性。

  2. 量化技术：解释什么是INT8量化？为什么TPU v1采用INT8量化能够提升推理性能？量化对模型精度的影响如何控制？

  3. BF16精度：比较FP16和BF16两种浮点格式的差异。为什么Google选择BF16作为TPU v2的训练精度？

  4. DOJO设计：分析Tesla DOJO超级计算机的设计理念。它与通用GPU集群（如NVIDIA DGX）的主要区别是什么？

  5. 竞争格局：分析当前全球AI芯片市场的竞争格局。除了技术性能，还有哪些因素影响AI芯片的市场竞争力？

  6. 架构创新思考：你认为未来AI芯片架构可能有哪些重大创新方向？需要解决哪些技术挑战？