由于采用了多攝像頭輸入和深度卷積骨干網(wǎng)絡，用于訓練自動駕駛感知模型的 GPU 內(nèi)存占用很大。當前減少內(nèi)存占用的方法往往會導致額外的計算開銷或工作負載的失衡。

本文介紹了 NVIDIA 和智能電動汽車開發(fā)商蔚來的聯(lián)合研究。具體來說，文中探討了張量并行卷積神經(jīng)網(wǎng)絡（CNN）訓練如何有助于減少 GPU 內(nèi)存占用，并展示了蔚來如何提高自動駕駛汽車感知模型的訓練效率和 GPU 利用率。

自動駕駛的感知模型訓練

自動駕駛感知任務采用多攝像頭數(shù)據(jù)作為輸入，卷積神經(jīng)網(wǎng)絡（CNN）作為骨干（backbone）來提取特征。由于 CNN 的前向激活值（activations）都是形狀為(N, C, H, W)的特征圖（feature maps）（其中 N、C、H、W 分別代表圖像數(shù)、通道數(shù)、高度和寬度）。這些激活值需要被保存下來用于反向傳播，因此骨干網(wǎng)絡的訓練通常會占據(jù)顯著的內(nèi)存大小。

例如，有 6 路相機以 RGB 格式輸入分辨率為 720p 的圖像，批大?。╞atchsize）設置為 1，那么骨干網(wǎng)絡的輸入形狀為(6, 3, 720, 1280)。對于如 RegNet 或 ConvNeXt 這樣的骨干網(wǎng)絡而言，激活值的內(nèi)存占用是遠大于模型權(quán)重和優(yōu)化器狀態(tài)的內(nèi)存占用的，并且可能會超出 GPU 的內(nèi)存大小限制。

蔚來汽車自動駕駛團隊在這一領域的研究表明，使用更深的模型和更高的圖像分辨率可以顯著提高感知精度，尤其是對尺寸小和距離遠的目標的識別；同時，蔚來 Aquila 超感系統(tǒng)搭載 11 個 800 萬像素高清攝像頭，每秒可產(chǎn)生 8GB 圖像數(shù)據(jù)。

GPU 內(nèi)存優(yōu)化需求

深度模型和高分辨率輸入對于 GPU 內(nèi)存優(yōu)化提出了更高的要求。當前解決激活值 GPU 內(nèi)存占用過大的技術(shù)有梯度檢查點（gradient checkpointing），即在前向傳播的過程中，只保留部分層的激活值。而對于其他層的激活值，則在反向傳播的時候重新計算。

這樣可以節(jié)省一定的 GPU 內(nèi)存，但會增加計算的開銷，拖慢模型訓練。此外，設置梯度檢查點通常需要開發(fā)者根據(jù)模型結(jié)構(gòu)來選擇和調(diào)試，這給模型訓練過程引入了額外的代價。

蔚來還使用了流水線并行技術(shù)，將神經(jīng)網(wǎng)絡按照 GPU 內(nèi)存開銷進行平均分段，部署到多個 GPU 上進行訓練。此方法雖然將存儲需求平分到多個 GPU 上，但是因為計算不平衡，會導致明顯的 GPU 間負載不均衡現(xiàn)象，一些 GPU 的計算資源無法被充分利用。

基于 PyTorch DTensor 的張量并行 CNN 訓練

綜合考慮以上因素，NVIDIA 和蔚來合作設計并實現(xiàn)了張量并行（Tensor Parallel）卷積神經(jīng)網(wǎng)絡訓練方案，將輸入值和中間激活值切分到多個 GPU 上。而對于模型權(quán)重和優(yōu)化器狀態(tài)，我們采用和數(shù)據(jù)并行訓練相同的策略，將其復制到各個 GPU 上。該方法能夠降低對單個 GPU 的內(nèi)存占用和帶寬壓力。

PyTorch 2.0 中引入的 DTensor 提供了一系列原語（primitives）來表達張量的分布如切片（sharding）和重復（replication），使用戶能夠方便地進行分布式計算而無需顯式調(diào)用通信算子，因為 DTensor 的底層實現(xiàn)已經(jīng)封裝了通信庫，如 NVIDIA 集合通信庫 (NCCL)。

有了 DTensor 的抽象，用戶可以方便地搭建各種并行訓練策略，如張量并行（Tensor Parallel），分布式數(shù)據(jù)并行（Distributed Data Parallel）和完全切片數(shù)據(jù)并行（Fully Sharded Data Parallel）。

實現(xiàn)

以用于視覺任務的 CNN 模型 ConvNeXt-XL 為例，我們將展示 Tensor Parallel 卷積神經(jīng)網(wǎng)絡訓練的實現(xiàn)。DTensor 放置方式如下：

模型參數(shù)：Replicate

重復放置在各個 GPU 上，模型包含 3.50 億個參數(shù)，以 FP32 存儲時占據(jù) 1.4GB GPU 內(nèi)存。

模型輸入：Shard(3)

切分（N, C, H, W）的 W 維度，將輸入分片放到各個 GPU 上。例如，在 4 個 GPU 上對形狀為(7, 3, 512, 2048) 的輸入執(zhí)行 Shard(3) 會生成四個切片，形狀為 (7, 3, 512, 512)。

激活值：Shard(3)

切分（N, C, H, W）的 W 維度，將激活值分片放在各個 GPU 上

模型參數(shù)的梯度：Replicate

重復放置在各個 GPU 上。

優(yōu)化器狀態(tài)：Replicate

重復放置在各個 GPU 上。

上述配置可以通過 DTensor 提供的 API 來實現(xiàn)，且用戶只需指明模型參數(shù)和模型輸入的放置方式，其他張量的放置方式會自動生成。

而要達成張量并行的訓練，我們需要給卷積算子 aten.convolution 和 aten.convolution_backward 注冊傳播規(guī)則，這將根據(jù)輸入 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，在使用其數(shù)值時會自動執(zhí)行規(guī)約操作，默認規(guī)約算子為求和。

接下來，我們便要給出張量并行的卷積算子前向和反向的實現(xiàn)。由于將激活值切分到了多個 GPU 上，1 個 GPU 上的本地卷積可能需要相鄰 GPU 上激活值的邊緣數(shù)據(jù)，這就需要 GPU 之間進行通信。在 ConvNeXt-XL 模型中，其降采樣層的卷積不存在該問題，而 Block 中的逐深度卷積則需要處理該問題。

如果無需交換數(shù)據(jù)，用戶可以直接調(diào)用卷積的前向和反向算子，傳入本地張量即可。如果需要交換本地激活值張量邊緣數(shù)據(jù)，則使用如圖 1 和圖 2 所示的卷積前向算法和反向算法，省略了圖中的 N 和 C 維度，并假設卷積核大小為 5x5，padding 為 2，stride 為 1。

圖 1 張量并行卷積前向算法示意圖

如圖 1 所示，當卷積核大小為 5x5，padding 為 2，stride 為 1 時，每個 GPU 上的本地 input 都需要取用相鄰 GPU 的寬度為 2 的輸入邊緣，并將收到的邊緣數(shù)據(jù)拼接到自身上。換句話說，需要 GPU 間的通信來確保張量并行卷積的正確性。這種數(shù)據(jù)交換，可以通過調(diào)用 PyTorch 封裝的 NCCL 發(fā)送接受通信算子來實現(xiàn)。

值得一提的是，在多個 GPU 上存在激活切片時，卷積算子的有些 padding 是不需要的。因此本地卷積前向傳播完成后，需要切除 output 中由不需要的 padding 引入的無效像素，如圖 1 中的藍色條所示。

圖 2 顯示了張量并行卷積的反向傳播。首先，在梯度輸出上應用 zero padding，這與前向傳播過程中的輸出切除操作相對應。對本地輸入同樣要進行數(shù)據(jù)交換、拼接和 padding 操作。

之后，通過調(diào)用每個 GPU 上的卷積反向算子，即可獲得權(quán)重梯度、偏置梯度和梯度輸入。

圖 2 張量并行卷積反向傳播工作流程

權(quán)重梯度和偏置梯度的 DTensor 放置方式是 _Partial，因此使用時會自動對它們的值進行多 GPU 規(guī)約操作。梯度輸入的 DTensor 放置方式是 Shard(3)。

最后，本地梯度輸入的邊緣像素會被發(fā)送到鄰近 GPU 并在相應位置累積，如圖 2 中的橙色條所示。

除了卷積層之外，ConvNeXt-XL 還有一些層需要處理以支持張量并行訓練。例如我們需要為 DropPath 層使用的 aten.bernoulli 算子傳播規(guī)則。該算子應被置于隨機數(shù)生成追蹤器的分布式區(qū)域內(nèi)，以保證各個 GPU 上的一致性。

所有代碼已經(jīng)并入了 PyTorch GitHub repo 的主分支，用戶使用時直接調(diào)用 DTensor 的上層 API 便可實現(xiàn)張量并行的卷積神經(jīng)網(wǎng)絡訓練。

使用張量并行訓練 ConvNeXt 的基準效果

我們在 NVIDIA DGX AI 平臺上進行了基準測試，研究 ConvNeXt-XL 訓練的速度和 GPU 內(nèi)存占用。梯度檢查點技術(shù)和 DTensor 是兼容的，并且結(jié)合兩項技術(shù)，GPU 的內(nèi)存占用能夠更顯著地降低。

測試的基線是在 1 個 NVIDIA GPU 上使用 PyTorch 原生 Tensor，輸入大小為(7, 3, 512, 1024)時的結(jié)果：不使用梯度檢查點時 GPU 內(nèi)存占用為 43.28 GiB，一次訓練迭代時間為 723 ms；使用梯度檢查點時 GPU 內(nèi)存占用為 11.89 GiB，一次訓練迭代時間為 934 ms。

全部測試結(jié)果如圖 3 和圖 4 所示：全局輸入形狀為 (7，3，512，W)，其中 W 從 1024 到 8192 不等。實線為未使用梯度檢查點時的結(jié)果，虛線為使用梯度檢查點時的結(jié)果。

圖 3 各種測試條件下的 GPU 內(nèi)存占用

圖 4 各種測試條件下一次訓練迭代耗時

如圖 3 所示，使用 DTensor 切分激活值可以有效降低 ConvNeXt-XL 訓練的 GPU 內(nèi)存占用，并且同時使用 DTensor 和梯度檢查點，ConvNeXt-XL 訓練的 GPU 內(nèi)存占用可以降到很低的水平。如圖 4 所示，張量并行方法有很好的弱擴展性；在問題規(guī)模足夠大時，也有不錯的強擴展性。下面是不使用梯度檢查點時的擴展性：

全局輸入(7, 3, 512, 2048)給 2 個 GPU 時，一次迭代時間為 937 ms

全局輸入(7, 3, 512, 4096)給 4 個 GPU 時，一次迭代時間為 952 ms

全局輸入(7, 3, 512, 4096)給 8 個 GPU 時，一次迭代時間為 647 ms

結(jié)論

蔚來自動駕駛開發(fā)平臺（NADP）是蔚來專門用于研發(fā)核心自動駕駛服務的平臺。該平臺可提供高性能計算和全鏈工具，用來處理每天成千上萬的日常推理和訓練任務，以確保主動安全和駕駛輔助功能的持續(xù)演進。使用 DTensor 實現(xiàn)的張量并行 CNN 訓練能夠有效提高 NADP 上的訓練效率。

該關鍵性的方案使得 NADP 能夠進行萬卡規(guī)模的并行計算，它提高了對 GPU 的利用率，降低了訓練模型的成本，支持了更靈活的模型結(jié)構(gòu)?；鶞蕼y試顯示，在蔚來自動駕駛場景下，該方法表現(xiàn)良好，有效解決了視覺大模型的訓練難題。

基于 PyTorch DTensor 的張量并行 CNN 訓練可顯著減少內(nèi)存占用并保持良好的可擴展性。我們預計該方法將充分利用多個 GPU 的算力和互連功能，使感知模型訓練更加普及。

審核編輯：劉清