使用张量并行技术进行自动驾驶感知模型训练

由于采用了多摄像头输入和深度卷积骨干网络，用于训练自动驾驶感知模型的 GPU 内存占用很大。当前减少内存占用的方法往往会导致额外的计算开销或工作负载的失衡。

本文介绍了 NVIDIA 和智能电动汽车开发商蔚来的联合研究。具体来说，文中探讨了张量并行卷积神经网络（CNN）训练如何有助于减少 GPU 内存占用，并展示了蔚来如何提高自动驾驶汽车感知模型的训练效率和 GPU 利用率。

自动驾驶的感知模型训练

自动驾驶感知任务采用多摄像头数据作为输入，卷积神经网络（CNN）作为骨干（backbone）来提取特征。由于 CNN 的前向激活值（activations）都是形状为(N, C, H, W)的特征图（feature maps）（其中 N、C、H、W 分别代表图像数、通道数、高度和宽度）。这些激活值需要被保存下来用于反向传播，因此骨干网络的训练通常会占据显著的内存大小。

例如，有 6 路相机以 RGB 格式输入分辨率为 720p 的图像，批大小（batchsize）设置为 1，那么骨干网络的输入形状为(6, 3, 720, 1280)。对于如 RegNet 或 ConvNeXt 这样的骨干网络而言，激活值的内存占用是远大于模型权重和优化器状态的内存占用的，并且可能会超出 GPU 的内存大小限制。

蔚来汽车自动驾驶团队在这一领域的研究表明，使用更深的模型和更高的图像分辨率可以显著提高感知精度，尤其是对尺寸小和距离远的目标的识别；同时，蔚来 Aquila 超感系统搭载 11 个 800 万像素高清摄像头，每秒可产生 8GB 图像数据。

GPU 内存优化需求

深度模型和高分辨率输入对于 GPU 内存优化提出了更高的要求。当前解决激活值 GPU 内存占用过大的技术有梯度检查点（gradient checkpointing），即在前向传播的过程中，只保留部分层的激活值。而对于其他层的激活值，则在反向传播的时候重新计算。

这样可以节省一定的 GPU 内存，但会增加计算的开销，拖慢模型训练。此外，设置梯度检查点通常需要开发者根据模型结构来选择和调试，这给模型训练过程引入了额外的代价。

蔚来还使用了流水线并行技术，将神经网络按照 GPU 内存开销进行平均分段，部署到多个 GPU 上进行训练。此方法虽然将存储需求平分到多个 GPU 上，但是因为计算不平衡，会导致明显的 GPU 间负载不均衡现象，一些 GPU 的计算资源无法被充分利用。

基于 PyTorch DTensor 的张量并行 CNN 训练

综合考虑以上因素，NVIDIA 和蔚来合作设计并实现了张量并行（Tensor Parallel）卷积神经网络训练方案，将输入值和中间激活值切分到多个 GPU 上。而对于模型权重和优化器状态，我们采用和数据并行训练相同的策略，将其复制到各个 GPU 上。该方法能够降低对单个 GPU 的内存占用和带宽压力。

PyTorch 2.0 中引入的 DTensor 提供了一系列原语（primitives）来表达张量的分布如切片（sharding）和重复（replication），使用户能够方便地进行分布式计算而无需显式调用通信算子，因为 DTensor 的底层实现已经封装了通信库，如 NVIDIA 集合通信库 (NCCL)。

有了 DTensor 的抽象，用户可以方便地搭建各种并行训练策略，如张量并行（Tensor Parallel），分布式数据并行（Distributed Data Parallel）和完全切片数据并行（Fully Sharded Data Parallel）。

实现

以用于视觉任务的 CNN 模型 ConvNeXt-XL 为例，我们将展示 Tensor Parallel 卷积神经网络训练的实现。DTensor 放置方式如下：

• 模型参数：Replicate

• 重复放置在各个 GPU 上，模型包含 3.50 亿个参数，以 FP32 存储时占据 1.4GB GPU 内存。

• 模型输入：Shard(3)

• 切分（N, C, H, W）的 W 维度，将输入分片放到各个 GPU 上。例如，在 4 个 GPU 上对形状为(7, 3, 512, 2048) 的输入执行 Shard(3) 会生成四个切片，形状为 (7, 3, 512, 512)。

• 激活值：Shard(3)

• 切分（N, C, H, W）的 W 维度，将激活值分片放在各个 GPU 上

• 模型参数的梯度：Replicate

• 重复放置在各个 GPU 上。

• 优化器状态：Replicate

• 重复放置在各个 GPU 上。

上述配置可以通过 DTensor 提供的 API 来实现，且用户只需指明模型参数和模型输入的放置方式，其他张量的放置方式会自动生成。

而要达成张量并行的训练，我们需要给卷积算子 aten.convolution 和 aten.convolution_backward 注册传播规则，这将根据输入 DTensor 的放置方式来确定输出 DTensor 的放置方式：

• aten.convolution

• Input 放置方式为 Shard(3)，weight 和 bias 放置方式为 Replicate，output 放置方式为 Shard(3)

• aten.convolution_backward

• grad_output 放置方式为 Shard(3)，weight和 bias 放置方式为 Replicate，grad_input 放置方式为 Shard(3)，grad_weight 和 grad_bias 方式方式为 _Partial

放置方式为 _Partial 的 DTensor，在使用其数值时会自动执行规约操作，默认规约算子为求和。

接下来，我们便要给出张量并行的卷积算子前向和反向的实现。由于将激活值切分到了多个 GPU 上，1 个 GPU 上的本地卷积可能需要相邻 GPU 上激活值的边缘数据，这就需要 GPU 之间进行通信。在 ConvNeXt-XL 模型中，其降采样层的卷积不存在该问题，而 Block 中的逐深度卷积则需要处理该问题。

如果无需交换数据，用户可以直接调用卷积的前向和反向算子，传入本地张量即可。如果需要交换本地激活值张量边缘数据，则使用如图 1 和图 2 所示的卷积前向算法和反向算法，省略了图中的 N 和 C 维度，并假设卷积核大小为 5x5，padding 为 2，stride 为 1。

图 1 张量并行卷积前向算法示意图

如图 1 所示，当卷积核大小为 5x5，padding 为 2，stride 为 1 时，每个 GPU 上的本地 input 都需要取用相邻 GPU 的宽度为 2 的输入边缘，并将收到的边缘数据拼接到自身上。换句话说，需要 GPU 间的通信来确保张量并行卷积的正确性。这种数据交换，可以通过调用 PyTorch 封装的 NCCL 发送接受通信算子来实现。

值得一提的是，在多个 GPU 上存在激活切片时，卷积算子的有些 padding 是不需要的。因此本地卷积前向传播完成后，需要切除 output 中由不需要的 padding 引入的无效像素，如图 1 中的蓝色条所示。

图 2 显示了张量并行卷积的反向传播。首先，在梯度输出上应用 zero padding，这与前向传播过程中的输出切除操作相对应。对本地输入同样要进行数据交换、拼接和 padding 操作。

之后，通过调用每个 GPU 上的卷积反向算子，即可获得权重梯度、偏置梯度和梯度输入。

图 2 张量并行卷积反向传播工作流程

权重梯度和偏置梯度的 DTensor 放置方式是 _Partial，因此使用时会自动对它们的值进行多 GPU 规约操作。梯度输入的 DTensor 放置方式是 Shard(3)。

最后，本地梯度输入的边缘像素会被发送到邻近 GPU 并在相应位置累积，如图 2 中的橙色条所示。

除了卷积层之外，ConvNeXt-XL 还有一些层需要处理以支持张量并行训练。例如我们需要为 DropPath 层使用的 aten.bernoulli 算子传播规则。该算子应被置于随机数生成追踪器的分布式区域内，以保证各个 GPU 上的一致性。

所有代码已经并入了 PyTorch GitHub repo 的主分支，用户使用时直接调用 DTensor 的上层 API 便可实现张量并行的卷积神经网络训练。

使用张量并行训练 ConvNeXt 的基准效果

我们在 NVIDIA DGX AI 平台上进行了基准测试，研究 ConvNeXt-XL 训练的速度和 GPU 内存占用。梯度检查点技术和 DTensor 是兼容的，并且结合两项技术，GPU 的内存占用能够更显著地降低。

测试的基线是在 1 个 NVIDIA GPU 上使用 PyTorch 原生 Tensor，输入大小为(7, 3, 512, 1024)时的结果：不使用梯度检查点时 GPU 内存占用为 43.28 GiB，一次训练迭代时间为 723 ms；使用梯度检查点时 GPU 内存占用为 11.89 GiB，一次训练迭代时间为 934 ms。

全部测试结果如图 3 和图 4 所示：全局输入形状为 (7，3，512，W)，其中 W 从 1024 到 8192 不等。实线为未使用梯度检查点时的结果，虚线为使用梯度检查点时的结果。

图 3 各种测试条件下的 GPU 内存占用

图 4 各种测试条件下一次训练迭代耗时

如图 3 所示，使用 DTensor 切分激活值可以有效降低 ConvNeXt-XL 训练的 GPU 内存占用，并且同时使用 DTensor 和梯度检查点，ConvNeXt-XL 训练的 GPU 内存占用可以降到很低的水平。如图 4 所示，张量并行方法有很好的弱扩展性；在问题规模足够大时，也有不错的强扩展性。下面是不使用梯度检查点时的扩展性：

• 全局输入(7, 3, 512, 2048)给 2 个 GPU 时，一次迭代时间为 937 ms
• 全局输入(7, 3, 512, 4096)给 4 个 GPU 时，一次迭代时间为 952 ms
• 全局输入(7, 3, 512, 4096)给 8 个 GPU 时，一次迭代时间为 647 ms

结论

蔚来自动驾驶开发平台（NADP）是蔚来专门用于研发核心自动驾驶服务的平台。该平台可提供高性能计算和全链工具，用来处理每天成千上万的日常推理和训练任务，以确保主动安全和驾驶辅助功能的持续演进。使用 DTensor 实现的张量并行 CNN 训练能够有效提高 NADP 上的训练效率。

该关键性的方案使得 NADP 能够进行万卡规模的并行计算，它提高了对 GPU 的利用率，降低了训练模型的成本，支持了更灵活的模型结构。基准测试显示，在蔚来自动驾驶场景下，该方法表现良好，有效解决了视觉大模型的训练难题。

基于 PyTorch DTensor 的张量并行 CNN 训练可显著减少内存占用并保持良好的可扩展性。我们预计该方法将充分利用多个 GPU 的算力和互连功能，使感知模型训练更加普及。

更多详情，请访问 PyTorch GitHub repo。

https://github.com/pytorch/pytorch/blob/main/torch/distributed/_tensor/examples/convnext_example.py

关于作者

曹添, 蔚来

曹添（Tian Cao）现任蔚来深度学习平台负责人。该平台提供适用于深度学习的高性能计算、高性能分布式存储和CICD。平台每天处理数十万推理和训练任务以及数千万数据处理任务。在加入蔚来之前，曹添曾在Momenta工作，主要负责提高深度学习的生产效率。

白宇利，蔚来

白宇利（Yuli Bai）目前担任蔚来汽车云端工程负责人，管理AI研发所需的数据、工具及超大规模基础设施，调度超过200EOPS的端云混合计算能力。此前，他作为Momenta的基础设施平台负责人，带领团队构建高性能AI平台解决方案。白宇利毕业于韩国延世大学计算机科学专业，曾在韩国知名教育娱乐公司担任全栈工程师。

刘洁, NVIDIA

Kingsley Liu现任NVIDIA AI开发者技术工程师，从事自动驾驶模型训练和推理优化方面的工作。在此之前，他曾从事NVIDIA Merlin HugeCTR推荐系统的开发工作。Kingsley Liu拥有清华大学信息与通信工程硕士学位，研究方向为无线通信与深度学习的交叉领域。

王泽寰, NVIDIA

王泽寰，现任NVIDIA Merlin HugeCTR和RnnEngine作者兼开发经理，同时也是NVIDIA亚太区开发者技术解决方案团队经理。自2014年以来主要从事深度学习GPU加速方面的工作。此前曾从事智能视频分析和HPC粒子模拟方面的工作。

陈文恺，NVIDIA

陈文恺，现任NVIDIA解决方案架构师，负责推进基于NVIDIA SDK的自动驾驶汽车和智能座舱工作流。他拥有北京邮电大学电子信息专业硕士学位，研究方向为计算机视觉。

李博，NVIDIA

李博，现任NVIDIA汽车行业开发者关系经理，负责拓展汽车行业的开发者生态系统和推进关键的AI工作负载。

郭若乾，NVIDIA

郭若乾现任NVIDIA AI开发者技术工程师。浙江大学硕士学位，于2020年加入NVIDIA，主要从事计算机视觉领域的算法加速以及大语言模型推理优化库TensorRT-LLM的开发工作。