由于采用了多摄像头输入和深度卷积骨干网络，用于训练自动驾驶感知模型的 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 的算力和互连功能，使感知模型训练更加普及。