PAI-TurboX： 面向自动驾驶的训练推理加速框架

技术背景

自动驾驶领域对高效且实时的数据处理能力提出了极高要求。如何优化深度学习模型的训练与推理流程，提升计算资源利用率，缩短模型迭代周期，并在保证精度的同时，实现低延迟、高吞吐量的部署，始终是行业面临的核心挑战。针对这些难题，PAI-TurboX 为自动驾驶场景中的复杂数据预处理、离线大规模模型训练和实时智能驾驶推理，提供了全方位的加速解决方案。这些方案不仅显著提升了感知、规划、控制等多模块系统的训练与推理效率，还有效加速端到端和世界模型的开发进程，助力自动驾驶技术不断迈向更高水平。

自动驾驶系统结构示例图（Autonomous Driving Small-Scale Cars: A Survey of Recent Development，https://arxiv.org/pdf/2404.06229）

自动驾驶面临的两大难题

• 训练期间，大规模训练样本数据加载和复杂的预处理，导致数据加载成为了训练系统的短板，使得 GPU 总是等待数据，无法满负荷运行，从而浪费 GPU 资源。

• 自动驾驶模型，需要融合图像、点云等多个模态数据，并且需要建模时序信息，导致无论是感知模型、规控模型或者是端到端模型结构都异常复杂，算子数量巨大，存在着大量的 CPU 和 GPU 交替切换，导致系统执行效率低下和优化难度大。

核心技术亮点

PAI-TurboX 从系统、数据、模型三个方面针对自动驾驶面临的两大难题进行优化。

在系统侧，TurboX 通过优化 CPU 亲和性、优化垃圾回收机制、动态编译、流水线并行和操作系统优化等方案，提升模型训练推理效率。

• 系统级性能调优，深度集成 CPU Affinity 亲和性优化方案，精准控制进程与 CPU 核绑定，减少上下文切换开销；优化 Python 垃圾回收机制，重构垃圾回收管理策略，降低 GC 频率与延迟抖动；以及 HugePage 内存加速技术，结合 HugePage 提升内存访问效率，显著降低 I/O 延迟，全面提升训练与推理任务的稳定性与吞吐能力。

• 支持模型动态编译技术，实时优化计算图结构，按需高效执行代码，显著提升模型的训练和推理速度与硬件利用率，实现性能与灵活性的双重突破。

• 流水线并行执行，在训练和推理过程中，高性能的 GPU 训练推理计算前后，往往伴随着密集的 CPU 算子计算，将 CPU 计算和 GPU 计算流水并行调度执行，能够大幅减少 GPU 等待时间，提升整个计算 Job 的吞吐。

• 操作系统（OS）优化：针对智驾场景下系统运行的特点，使用阿里云操作系统团队提供的高效动态内存管理、系统运行参数自适应调优、系统资源防争抢等优化能力，保障智驾相关模型训练和推理的运行效率。

• 高效动态内存管理：通过自研的内核内存特性，降低系统动态内存管理所占用的 CPU 计算资源，解决大规模数据处理所引发的系统内存碎片化，从而提升模型训练、推理整体吞吐和稳定性。

• 系统参数自动调优：提供自动系统参数调优工具 keentune，可以自动匹配出最适合模型训练的系统优化参数并完成自动设置，在用户无感的情况下自动提升模型训练效率。

• 系统资源防争抢：细粒度的系统资源监控与隔离能力，自动为高优先级的模型训练任务预留计算、存储、带宽等资源，缓解训练数据的在线处理对训练吞吐负面影响。

• 在数据侧，PAI-TurboX 提出了高性能的 DataLoader 引擎、优化了数据预处理流程和实现了智能训练样本分组，有效提升了数据处理效率。

• 高性能 DataLoader 引擎：通过循环加载数据，避免训练时等待数据，利用锁页内存，加快数据从 CPU 到 GPU 的传输效率，实现训练的快速启动。

• 深度重构预处理流程：涵盖智能调整大小、精准裁剪等核心环节，构建高效统一的数据增强与加载流水线，显著降低预处理延迟，提升数据吞吐能力。

• 智能训练样本分组：采用智能分组采样引擎，精准聚合同类样本，突破传统随机采样的局限，并减少分组的多次拷贝，实现按类别或特征对样本进行精准聚合，实现高性能训练。

• 
  

• 在模型侧，PAI-TurboX 从算子、Module、优化器、优化策略、量化、CPU 和 GPU 设备的自动切换等多个角度显著提升模型训练与推理的速度。

• 高性能算子优化：针对常见的自动驾驶感知任务中的核心算子进行了深度优化，实现了高效体素化 cuda kernel，提升 point cloud 体素化的效率；通过 shared memory 优化 nms kernel 和 deformable kernel 的执行效率；通过 CUDA 重写并融合重复逻辑，去除冗余操作来优化稀疏卷积算子、矩阵求逆运算、坐标变换小矩阵乘法等常见的自动驾驶算子，显著减少访存延迟，大幅提升训练吞吐；

• 高性能 Module 优化：在多卡分布式训练中，批归一化层 BatchNorm 需要做同步处理 SyncBN，去除了冗余的 GPU->CPU 同步，提升了自动驾驶场景 BN 同步阶段的执行效率。

• Workload 设备重映射: 将模型训练过程中适合 CPU 计算的部分迁移到 dataloader 部分预先计算，将计算密集型的数据处理迁移到 gpu 上计算，以降低 cpu 处理的压力、提升 gpu 的执行效率。

• 为提升模型推理效率：采用 INT8 量化、BF16/TF32 半精度计算及自动混合精度（AMP）策略，结合硬件特性对算子进行细粒度精度优化，在保障精度的前提下显著降低计算开销与内存带宽需求，实现异构平台下的高性能推理部署。

• 优化器融合：通过融合优化器技术将动量更新、权重衰减与梯度下降操作合并为单一内核执行，降低了 GPU 内存访问压力并提升了整体训练性能。

• 训练策略优化：使用 CudaGraph 来优化多个 kernel 启动，比如将 EMA 的更新逻辑整合进 CudaGraph 中，减少 Python 层对模型参数的逐层访问与更新。

实测性能速度提升

PAI-TurboX 支持大量自动驾驶模型优化，如 BEVFusion、BEVFormer、GameFormer、MTR、MapTR、MapQR、Open-Clip、Sparse4D、SparseDrive 和 OpenVLA 等模型。

BEVFusion 流程图（BEVFusion: A Simple and Robust LiDAR-Camera Fusion Framework）

下面的图像展示了在 PAI-TurboX 在自动驾驶各类模型上都有着充分的速度提升。BEVFusion 模型主要是融合激光雷达点云和图像特征，产生一个多模态的 BEV 特征，以支持各种 3D 感知任务，应用 PAI-TurboX 可以将训练时间缩短 58.5%。MapTR 一个结构化的端到端框架，用于高效的在线矢量高清地图构建，应用 PAI-TurboX 可以将训练时间缩短 53.0%。Sparse4D 是一个高性能高效率的长时序纯稀疏融合感知算法，应用 PAI-TurboX 可以将训练时间缩短 49.8%。

SparseDrive 是一个通过稀疏场景表示实现了端到端自动驾驶的模型，应用 PAI-TurboX 可以在这个算法的两个训练阶段获得明显的训练速度提升，相同训练步数下两个阶段的时间可以分别缩减 51.5%和 48.5%。

在 PAI 平台中使用 PAI-TurboX

1、在 PAI-Notebook Gallery 广场找到 PAI-Turbox 加速例子的 Notebook，可以使用搜索框找到对应模型卡片或通过链接直达该模型：

https://pai.console.aliyun.com/#/dsw-gallery/preview/deepLearning/cv/turbox

2、点击“在 DSW 中打开”按钮，创建对应所需的实例。

3、填入实例所需的必填项后，完成实例的创建。

4、创建完成后，选择对应的实例进入 Notebook。

5、进入 Notebook，按步骤进行点击，即可完成 PAI-TurboX 框架的拉起。

未来规划

除了进一步优化当前的一些加速方案外，PAI-TurboX 将进一步探索新的加速策略。随着单机 GPU 算力变得越来越强大，在自动驾驶模型的部分训练和推理场景中，GPU 处理数据的速度比 CPU 加载和预处理数据的速度更快，为了充分利用 GPU 的算力资源，提出 RemoteDataloader 解决方案来将数据加载和预处理与 GPU 处理机器分离，通过其它节点预先加载和预处理数据。