特斯拉在AI Day 2022先介绍了AI编译器,后面又介绍了Dojo的硬件软件,软件部分和AI编译器有部分重叠,本文介绍还是延用AI Day的思路,分为三部分:AI编译和推理,Dojo硬件,Dojo软件。
特斯拉车道检测网络
特斯拉车道网络使用稀疏点预测和计算来识别和预测车道线的位置和形状,极大地减少了计算量,使得网络可以在高帧率下运行,提供实时的车道信息。稀疏点预测和计算是特斯拉FSD推理系统的重要组成部分,通过高效的计算和存储策略,优化了稀疏点的预测和处理过程。
1 稀疏点预测
热图生成:网络首先预测最可能的空间位置点的热图(heat map),这一步骤使用卷积神经网络或其他适合的架构生成每个可能位置的概率分布。
ArgMax操作:对生成的热图应用ArgMax操作,找到概率最高的空间位置索引,这一步骤确定了最可能的稀疏点的位置。
One-hot编码:将找到的空间位置索引转换为one-hot编码。one-hot编码是一种稀疏表示,只在一个位置上有1,其余位置都是0。
2 稀疏计算
嵌入向量选择:使用one-hot编码从嵌入表(embedding table)中选择对应的嵌入向量。嵌入表是在训练期间学习的,包含了各个位置的特征表示;特斯拉在SRAM中构建了一个查找表,设计了嵌入向量的维度,使得可以通过矩阵乘法来实现高效查找。
嵌入缓存:选中的嵌入向量存储在一个token缓存中,这一缓存机制避免了每次迭代都重新计算嵌入,从而提高了计算效率和速度。
矩阵乘法实现查找:通过矩阵乘法实现嵌入查找和计算,这种方式利用了硬件加速器(如特斯拉FSD芯片)对矩阵运算的优化,提高了计算效率。
3 稀疏点处理
特征融合:稀疏点的特征通过自回归和迭代的方式逐步融合,每一步处理都结合前一步的稀疏点特征,逐步构建全局的空间表示。
迭代优化:每次迭代中,网络生成新一轮的稀疏点预测,并将这些点的特征与之前的结果结合,进行进一步处理,这一过程反复进行,直到得到最终的稀疏点集合。
特斯拉的AI编译器
特斯拉的AI编译器是为其自动驾驶系统优化深度学习模型推理而设计的,主要目标是将复杂的神经网络图优化为高效的推理代码,以在FSD计算硬件上运行。
1 AI编译器功能
特斯拉的AI编译器旨在处理大规模神经网络图,将其分割、优化并生成适合在其自研硬件上运行的高效代码。这个编译器处理的图包含了数十万的节点和连接,涵盖了自动驾驶系统中的各种任务,包括感知、预测和规划。
2 编译器工作流程
(可结合之前介绍的AI编译器背景内容:[编译器]-2 AI编译器介绍)
1)图分割与划分(Graph Partitioning and Splitting):编译器首先接收一个完整的神经网络图,该图可能包含多个任务和子任务,使用图分割技术,将大图分割成多个独立的子图,每个子图代表一个相对独立的计算单元,可以在硬件上独立运行。
2)子图优化(Subgraph Optimization):对每个子图进行独立优化,利用特斯拉硬件的特性进行定制化调整。具体优化包括操作融合、常量折叠和子图内的并行化等。
3)生成高效代码(Code Generation):编译器将优化后的子图转换为适合FSD硬件运行的低级代码,生成的代码能够充分利用硬件的特性,如硬件加速的矩阵运算和特定的AI计算单元。
4)链接与调度(Linking and Scheduling):将各个独立的子图代码链接起来,形成一个完整的推理程序,通过优化的调度算法,确保子图之间的数据传输和计算顺序最优化,以实现最低的推理延迟和最高的计算效率。
3 常规优化策略
1)操作融合(Operation Fusion):将多个相邻的算子融合为一个算子,减少中间数据存储和传输的开销。例如将卷积操作和激活函数融合在一起,直接计算激活后的结果。
2)常量折叠(Constant Folding):在编译时预计算所有可能的常量表达式,减少运行时的计算量。例如将所有不依赖于输入数据的固定计算提前计算并存储。
3)内存优化(Memory Optimization):优化内存分配和数据布局,以减少内存访问延迟和带宽消耗。利用内存复用技术,使得同一块内存区域可以在不同时间段存储不同的数据,减少总的内存需求。
4)并行化(Parallelization):充分利用FSD硬件的多核架构和并行计算能力,将计算任务拆分并行执行,利用专用AI加速单元的并行计算特性,提高整体计算速度。
5)稀疏计算(Sparse Computation):针对稀疏神经网络图,利用稀疏矩阵乘法和稀疏激活函数等技术减少无效计算,通过只计算非零值的操作,减少总体计算量。
6)流水线和批处理(Pipelining and Batching):利用流水线技术将不同计算任务的执行重叠起来,提高计算资源的利用率。通过批处理技术,将多个输入数据同时处理,减少单次推理的平均计算时间。
DOJO硬件
1 Dojo Tile
Dojo Tile是Dojo加速器的核心计算单元,每个Tile包含25个D1芯片(并行处理大规模的矩阵乘法运算),这些芯片专门设计用于高效的矩阵运算和深度学习任务。Tile的设计考虑了高带宽、低延迟的特点,使其能够在大规模计算任务中保持高效。
2 System Tray
System Tray是Dojo加速器中一个关键的组件,高度集成的系统托盘负责机械和热管理,并集成了高功率传输功能;每个托盘支持六个Tile,总计算能力达到54 PetaFLOPS,配备640 GB高带宽内存。具体特点如下:
1)层压母线(Laminated Bus Bar):这是一种高功率传输组件,能够在极高的功率密度下进行机械和热管理。它的高度为75毫米,能够支持六个Tile,总重量为135千克。
2)数据接口:System Tray提供高带宽的DRAM接口,用于存储训练数据,并通过TTP(Tesla Transfer Protocol)客户协议与训练Tile进行通信,确保了数据能够高效地传输到每个Tile。
3 Exapod
Exapod是Dojo系统的更大级别的集成单元,旨在将多个Dojo加速器组合在一起,总计算能力达到一ExaFLOP,确保大规模机器学习任务的高效处理。配备冗余电源,可以将三相480伏交流电转换为52伏直流电,确保系统在高功率需求下的稳定运行。
DOJO软件
特斯拉全自动驾驶(FSD)软件系统依赖于高度优化的软硬件协同工作,主要组件包括JIT神经网络编译器、LLVM后端、数据摄取流水线等。
JIT NN编译器和LLVM后端:JIT NN编译器即时编译神经网络模型,将高层次的神经网络代码转换为低层次的机器指令。LLVM后端利用LLVM编译器基础设施生成高效的机器代码,使得编译后的代码能够在Dojo硬件上高效运行。
摄取流水线(Ingest Pipeline):负责将数据高效地传输到硬件中,确保硬件始终有足够的数据进行处理,避免因数据饥饿而导致的性能下降。
1 性能优化
性能的提升依赖于利用率(utilization)和加速器占用率(occupancy)。
硬件提供了高峰值性能,编译器的任务是从硬件中提取最大化的利用率,而摄取流水线的任务是确保数据吞吐量足够高。特斯拉的Dojo系统设计用于高利用率处理大规模模型,不仅加速计算密集型部分,还加速延迟敏感部分(如Batch Normalization)和带宽敏感部分(如梯度汇总和参数聚合)。
2 编译器优化策略
1)提取并行性
模型并行:在多个加速器上并行执行不同部分的模型。
数据并行:在多个加速器上并行处理不同批次的数据。
张量并行:在多个加速器上并行处理单个张量的不同部分。
2)SRAM中的张量存储:张量通常在SRAM中分片存储,并在执行层时即时复制。
3)重叠计算和数据传输:在计算的同时进行数据传输,最大限度地利用计算资源。
3 性能测试结果
一句话,特斯拉的硬件性能远超行业其它GPU。