教程向：如何提高多卡训练速度（附github代码+实验结果）

教程向：如何提高多卡训练速度（附github代码+实验结果）

摘要

数据并行（DDP）的多卡训练能够显著减少模型的训练时长。但是由于显卡带宽，同步等问题，多卡训练往往会带来一定的性能损耗。较少人关注到选择不同位置的GPU会带来不同的并行效率，本文主要是探讨实验室常见的机型配置，包括NVLink，PCIE-bridge等硬件对训练效率产生的影响，并探讨了多节点训练的训练效率与训练损耗。

前置知识

本文默认读者已经有深度学习基础、并且了解pytorch的DDP原理。后文重点关注多卡训练时硬件侧会影响训练效率的部分。如果需要了解基础知识，可查阅文末附的优秀博客。

为什么要多卡训练

在以前大多数情况下多卡训练都是为了加速。比如用TitianRTX 24G显卡在ImageNet数据集上跑Resnet50从零训练，大概要跑4-5天。但是如果有8张卡就可以1天多跑完。

多卡训练的原理简单回顾

简单说明DDP的工作原理就是。假设我们想在八张显卡上训练模型。我们就将模型复制八份到每一张显卡上，然后数据切分成八份在8张卡上，然后让每张卡都进行前向传递和反向传播，每个batch执行完后八张卡要进行一次all reduce来同步权重（不然八张卡上的权重会变的不一样）。过程类似下图。

但执行效率的损失坏就坏在八张卡要同步上，为了保证八张卡的权重一次，每次batch都需要同步一次，会引入两个问题：

1. 每次同步卡间传输的参数量超大
2. 如果有卡提前计算完了，要等其他卡也计算完。相当于每个batch的速度取决于最慢的卡。

第一点是后文主要讨论的内容，第二点的话需要注意就是假设你要在GPU0,1,2,3上开启实验，你发现GPU2已经有人在训练了，就不要用了，因为单块卡就会直接拖慢剩余3块卡的速度。

除了模型并行外，还有没有别的方法呢。也有，比如参数服务器parameters server原理这个就是异步的计算。但一般实验室也就是撑死了八块卡，其实这个效率还算能接受。

影响带宽的硬件因素

那么，因为DDP的多卡训练，卡间通信是拖慢速度的主要因素。如何提高卡间通信速度呢。首先要先理解多GPU服务器的结构。

首先一个常见的误区就是很多人会以为插在同一个电脑（服务器）中的八块卡间传输效率一样。实际上如果是一般的单CPU4卡台式机，确实差别不会很明显。因为所有的显卡都在一个PCIe桥或者直连CPU。

但是如果是八卡服务器，一般是带两个CPU的。八张卡会分为两组：0，1，2，3GPU在CPU0，4，5，6，7在CPU1。这里以常用的超微4124八卡服务器为例，官方文档 可以看到连卡的PCIe被分为了两组分别连在两个CPU上。如果跨组调用GPU将带来较大的延迟。

那么第二个就是是否使用NVLink了。NVLink可以说是老黄黑科技，在显卡间搭建了直连通道，数据不需要走PCIe就能直接在卡间传输。并且带宽、延迟、速度都显著超过走PCIe。（似乎多卡训练时需要指定使用nccl作为后端才能使用，这个待测试。）

这里可以使用nvidia-smi的命令查看卡间是否有nvlink加速（当然如果不是PCIe卡，而是NVSwitch就不用关注这个了，直接随便选着用就行），使用如下命令查看显卡间连接拓扑图：

nvidia-smi topo -m

结果如下，可以看到0-3号GPU是在同一CPU上的，4-7GPU是在另一个CPU上。0-1，2-3，4-5，6-7GPU之间使用了NVLink连接，NVLink版本为4。使用NVLink能够极大提升卡间通信速度，尤其是在张量并行的情况下。

那么如果需要跨节点进行超过八卡的训练，比如两节点16卡训练，那么就要关注网络带宽了，一般推荐节点间采用40G光互联的方式。如果没条件使用10G网线互联也是可以的。一块10G双口网卡二手也就100来块钱，线的成本忽略不计。主要是10G交换机会贵一点（8口便宜的1-2000），如果是两个节点的GPU服务器可以考虑直接两台服务器直连，成本能低不少，不止两个节点的实验室也不差钱买台40G光交换机了。

实验

这里将重点比较下，训练代码完全一致，只是调用在不同位置的GPU，产生的训练性能差异。

实验环境

这里采用huggingface accelerate来完成多卡、多节点的启动。使用SwanLab进行训练跟踪。

训练任务选用的是使用Resnet18进行CIFAR10数据集分类。后续补一个使用transformers的实验（主要是折腾了半天transformers和accelerate的联动，没搞清楚他俩间什么关系）。

对于Accelerate不了解的同学可以参考我的另一篇博客入门向Accelerate教程

相关代码开源在github项目上，欢迎Star👏。实验结果可参考SwanLab项目

硬件环境

• 服务器：2xsupermicro4124gs-tnr + 192core-512G-RAM
• 显卡类型：8x RTX3090 4xNVLink
• 服务器间连接：10G电口，过交换机

软件环境+代码

• Ubuntu24.04 + Python3.11 + CUDA12.1
• 实验代码：https://github.com/ShaohonChen/speed_with_ddp.git

完整训练代码

建议参考github项目获得完整代码。实验结果也开源在SwanLab项目上

为了方便有人没法翻墙打不开还是放一下代码，直接复制此处的代码需要使用accelerate config设置下实验的运行硬件信息（比如训练卡数、是否多节点等），github中有配置好的config文件可以直接使用：

import torch.utils
import torch.utils.data
import torch.utils.data.dataloader

# from tutils import open_dev_mode
import swanlab
from swanlab.integration.accelerate import SwanLabTracker

# swanlab.login(open_dev_mode())

import torch
from torchvision.models import resnet18, ResNet18_Weights
import torchvision

from accelerate import Accelerator
from accelerate.logging import get_logger
import time
import fire


def main(exp="1gpu"):
    # hyperparameters
    config = {
   
        "num_epoch": 5,
        "batch_num": 64,
        "learning_rate": 1e-3,
        "report_step_num": 20,
    }