id="0__1">0.

前言

在将预训练的机器学习模型投入生产环境之前，模型训练是不可或缺的关键环节。

随着深度学习的发展，大模型往往具有数百万乃至数十亿参数。

使用反向传播来调整这些参数需要大量的内存和计算资源。

即便如此，模型训练仍然可能需要数天甚至数月时间才能完成。

在本节中，我们将探讨如何通过跨机器和机器内多进程的分布式训练来加速模型训练过程。

我们将系统学习

PyTorch

API——torch.distributed、torch.multiprocessing

torch.utils.data.distributed.DistributedSampler，使用这些

API

通过本节学习，将能够充分释放硬件设备的训练潜力。

对于超大规模模型训练而言，本节所探讨的工具不仅至关重要，在某些情况下甚至是不可或缺的。

进行分布式训练

在本节中，我们将模型训练过程从常规训练转换为分布式训练，探讨

PyTorch

提供的分布式训练工具，这些工具能显著提升训练速度并优化硬件使用效率。

以常规方式训练模型

(1)

首先导入所需库：

punctuation">.functional

punctuation">.utils

class="token

punctuation">.device

class="token

punctuation">)

接下来，定义卷积神经网络

class-name">ConvNet

class="token

punctuation">.Module

class="token

punctuation">(ConvNet

class="token

punctuation">.__init__

class="token

punctuation">.Conv2d

class="token

punctuation">.Conv2d

class="token

punctuation">.Dropout

class="token

punctuation">.Dropout

class="token

punctuation">.Linear

class="token

punctuation">.Linear

class="token

punctuation">.max_pool2d

class="token

punctuation">.flatten

class="token

punctuation">.log_softmax

class="token

punctuation">.utils

class="token

punctuation">.DataLoader

class="token

punctuation">.MNIST

class="token

operator">=transforms

class="token

punctuation">.Compose

class="token

punctuation">.ToTensor

class="token

punctuation">.Normalize

class="token

punctuation">.Adadelta

class="token

punctuation">(model

class="token

punctuation">.parameters

class="token

punctuation">.train

class="token

punctuation">)

PyTorch

(ConvNet)，并定义了优化器。

punctuation">.epochs

class="token

punctuation">(train_dataloader

class="token

punctuation">(device

class="token

punctuation">(device

class="token

punctuation">.nll_loss

class="token

punctuation">(pred_prob

class="token

punctuation">.zero_grad

class="token

punctuation">.backward

class="token

punctuation">(train_dataloader

class="token

punctuation">(train_dataloader

class="token

punctuation">)

epoch

数。

在循环内，通过批数据的方式遍历整个训练数据集，本节中批大小为

128

个训练数据点的批次，使用模型进行前向传播，以计算预测概率。

然后，我们将预测结果结合真实标签计算批次损失，并通过反向传播利用该损失梯度来调整模型参数。

将所有组件整合在

punctuation">.ArgumentParser

class="token

punctuation">.add_argument

class="token

punctuation">.parse_args

class="token

punctuation">}

class="token

punctuation">)

Python

数。

我们还对训练过程进行了计时，以便可以将它与分布式训练过程进行比较。

main()

punctuation">)

Python

epoch，因为当前重点不在于模型精度，而在于模型训练耗时。

可以看到输出结果如下所示：

src="https://i-blog.csdnimg.cn/direct/ef6dc7ca957d45e6931be750a12882bb.png#pic_center">

训练

1

API，即使存在跨进程或跨机器重复传递数据的额外开销，模型训练速度也能显著提升。

punctuation">.functional

punctuation">.utils

class="token

punctuation">.multiprocessing

punctuation">.distributed

torch.multiprocessing

用于在单台机器上生成多个

核心数生成对应数量的进程)，而torch.distributed则实现不同机器间的通信协作，使它们能共同完成模型训练。

执行时，我们需要在每台参与训练的机器上显式启动训练脚本。

PyTorch

Gloo

)会自动处理机器间的通信协调。

在每台机器内部，多进程机制会进一步将训练任务并行分配到各个进程。

模型架构定义部分保持不变：

class-name">ConvNet

class="token

punctuation">.Module

class="token

punctuation">(ConvNet

class="token

punctuation">.__init__

class="token

punctuation">.Conv2d

class="token

punctuation">.Conv2d

class="token

punctuation">.Dropout2d

class="token

punctuation">.Dropout2d

class="token

punctuation">.Linear

class="token

punctuation">.Linear

class="token

punctuation">.max_pool2d

class="token

punctuation">.flatten

class="token

punctuation">.log_softmax

class="token

punctuation">(cpu_num

class="token

punctuation">.machine_id

punctuation">.num_processes

punctuation">.init_process_group

class="token

punctuation">.world_size

class="token

punctuation">.manual_seed

class="token

punctuation">.device

class="token

punctuation">)

rank

值——这本质上是该进程在整个分布式系统中的顺序标识符。

举例来说，若使用

2

class="katex-mathml">rank=n×4+krank=n\times

4+k

0.0315em">k

init_process_group，该方法为每个启动的进程配置以下关键参数：

Gloo)

参与分布式训练的进程总量(由

指定)，亦称

world_size

当前启动进程的全局

rank

值

init_process_group

方法会阻塞所有进程，直到跨机器的全部进程都完成初始化才会继续执行。

PyTorch

• Gloo
• NCCL
• MPI

简而言之，对于

CPU

punctuation">.MNIST

class="token

operator">=transforms

class="token

punctuation">.Compose

class="token

punctuation">.ToTensor

class="token

punctuation">.Normalize

class="token

punctuation">.utils

class="token

punctuation">.distributed

class="token

punctuation">.DistributedSampler

class="token

punctuation">.world_size

class="token

punctuation">.utils

class="token

punctuation">.DataLoader

class="token

operator">=train_dataset

class="token

punctuation">.batch_size

class="token

operator">=train_sampler

class="token

punctuation">.Adadelta

class="token

punctuation">(model

class="token

punctuation">.parameters

class="token

punctuation">.parallel

class="token

punctuation">.DistributedDataParallel

class="token

punctuation">(model

class="token

punctuation">.train

class="token

punctuation">)

与单机训练相比，分布式训练的关键改进体现在数据加载与模型封装两个层面。

我们将

MNIST

DistributedSampler

shuffle

False，因为数据分配已由采样器控制。

nn.parallel.DistributedDataParallel

DistributedDataParallel

• 分布式环境中的每个派生进程都会获得独立的模型副本
• 每个进程的模型都维护自己的优化器，并与全局迭代保持同步的局部优化步骤
• 在每次分布式训练迭代时，各进程独立计算损失值及梯度，随后跨进程对这些梯度求取平均值
• 平均后的梯度将通过全局反向传播机制同步到所有模型副本，用于调整参数
• 由于全局反向传播步骤的存在，所有模型参数在每次迭代时都保持一致，从而实现自动同步

DistributedDataParallel

Python

进程运行在独立的解释器上，有效规避了在单一解释器下多线程实例化多个模型可能引发的全局解释器锁

(Global

限制问题。

这进一步提升了性能表现，特别是对于那些需要大量Python专属运算的模型而言。

punctuation">.epochs

class="token

punctuation">(train_dataloader

class="token

punctuation">(device

class="token

punctuation">(device

class="token

punctuation">.nll_loss

class="token

punctuation">(pred_prob

class="token

punctuation">.zero_grad

class="token

punctuation">.backward

class="token

punctuation">(train_dataloader

class="token

punctuation">(train_dataloader

class="token

punctuation">)

最后，训练循环几乎和单机训练一样。

唯一的区别在于我们限制只有排名为0的进程才能获取日志信息。

这是因为排名为

0

的进程作为参考来跟踪模型训练性能。

如果不加以限制，每个模型训练迭代都会产生与进程数量相同的日志行数。

将所有组件整合在

punctuation">.ArgumentParser

class="token

punctuation">.add_argument

class="token

string">'--num-machines'

class="token

punctuation">.add_argument

class="token

string">'--num-processes'

class="token

punctuation">.add_argument

class="token

string">'--machine-id'

class="token

punctuation">.add_argument

class="token

punctuation">.add_argument

class="token

string">'--batch-size'

class="token

punctuation">.parse_args

class="token

punctuation">.world_size

punctuation">.num_processes

punctuation">.num_machines

os

punctuation">.environ

class="token

string">'MASTER_ADDR'

class="token

punctuation">.environ

class="token

string">'MASTER_PORT'

class="token

punctuation">.spawn

class="token

punctuation">(train

class="token

punctuation">.num_processes

class="token

punctuation">}

class="token

punctuation">)

• num_machines：机器总数量
• num_processes：每台机器上要启动的进程数量
• machine_id：当前机器的序号

  Python

  脚本需要在每台机器上单独启动

• batch_size：每个批次的数据样数量。

  参数作用如下：

  • 所有进程将各自计算梯度，这些梯度将在每次迭代中被平均，从而获得整体梯度
  • 完整的训练数据集被分割为

    个独立的数据集

因此，在每次迭代时，完整的批次数据需要被分割成

现在与

相关联，所以我们将其作为输入参数提供，目的是为了简化训练接口。

参数定义后，计算

• MASTER_ADDR：运行

  rank

  地址

• MASTER_PORT：运行

  rank

  进程的主机上可用的端口号

rank

0

机器负责建立所有后端通信连接，因此整个系统必须能随时定位到该主机，这就是为什么需要提供其

IP

个进程，而非仅运行单个训练进程。

分布式参数会传递给每个派生进程，确保模型训练过程中各进程与机器之间能自主协调。

分布式训练：

优化建议：

如果你的项目访问量较大，建议增加缓存机制。

我们团队在优化后，接口响应时间从800ms降到了50ms，

效果非常明显。

具体的缓存策略可以根据业务场景调整。

punctuation">)

启动分布式训练脚本。

首先使用分布式脚本进行类非分布式运行，将机器数量和进程数量都设置为

1：

number">128

需要注意的是，由于本次训练只使用单个进程，batch_size

与之前非分布式训练时保持一致(仍为

src="https://i-blog.csdnimg.cn/direct/0399a43aea07439d8419bdfe703f3b26.png#pic_center">

若将此结果与上一节非分布式训练的输出对比，可发现训练时间基本相当(约

30

src="https://i-blog.csdnimg.cn/direct/ddabcb373c3a44b8a4d71cea0c4b6175.png#pic_center">

可以看到，训练时间从

30

秒。

训练损失的变化趋势没有受到影响，这表明分布式训练可以加速训练过程，同时保持模型的准确性。

接下来，使用

src="https://i-blog.csdnimg.cn/direct/55914888f36d4ffea3c686c87f8f2235.png#pic_center">

可以看到，训练时间进一步减少，从

20

秒。

训练损失的变化趋势仍然与之前的训练相似。

通过分布式训练，我们已经将训练时间从

30

src="https://i-blog.csdnimg.cn/direct/142ec629bded486bb476584b72cb687e.png#pic_center">

与预期相反，训练时间不仅没有进一步缩短，反而从

15

秒。

由于代码在本地机器执行，系统还存在其他进程(如浏览器)会与部分分布式训练进程争夺资源。

如果分布式训练模型是在远程机器上进行的，同时这些机器的唯一任务就是进行模型训练，在这样的机器上，建议使用与

CPU

最后需要指出的是，由于在本节中我们只使用了一台机器，因此我们只需要启动一个

Python

脚本来开始训练。

然而，如果是在多台机器上进行训练，那么除了修改

外，还需要在每台机器上启动一个

上实施分布式训练深度学习模型的实践探讨，这种方法能带来显著的加速效果。

仅需添加少量代码，就能将常规

PyTorch

模型训练脚本升级为分布式训练模式。

虽然上述实验基于简单的卷积网络，但由于我们完全无需修改模型架构代码，因此这套方案可直接扩展到更复杂的模型训练场景。

接下来，我们将简要讨论如何应用类似的代码更改，实现

GPU

punctuation">.device

class="token

punctuation">.is_available

class="token

punctuation">)

CUDA

通过并行化处理神经网络常规运算(如矩阵乘法和加法)所能提供的显著加速优势。

本节我们将探讨如何通过

GPU

punctuation">(gpu_num

class="token

punctuation">.machine_id

punctuation">.num_gpu_processes

punctuation">.init_process_group

class="token

punctuation">.world_size

class="token

punctuation">.set_device

class="token

punctuation">(gpu_num

class="token

punctuation">(gpu_num

class="token

punctuation">.NLLLoss

class="token

punctuation">(gpu_num

class="token

punctuation">)

在使用

punctuation">.MNIST

class="token

operator">=transforms

class="token

punctuation">.Compose

class="token

punctuation">.ToTensor

class="token

punctuation">.Normalize

class="token

punctuation">.utils

class="token

punctuation">.distributed

class="token

punctuation">.DistributedSampler

class="token

punctuation">.world_size

class="token

punctuation">.utils

class="token

punctuation">.DataLoader

class="token

operator">=train_dataset

class="token

punctuation">.batch_size

class="token

operator">=train_sampler

class="token

punctuation">.Adadelta

class="token

punctuation">(model

class="token

punctuation">.parameters

class="token

punctuation">.parallel

class="token

punctuation">.DistributedDataParallel

class="token

punctuation">(model

class="token

punctuation">[gpu_num

class="token

punctuation">.train

class="token

punctuation">)

DistributedDataParallel

API

包含一个关键参数——device_ids，用于指定调用该

API

True，该参数能显著加速训练过程中从主机(此处指加载数据集的

CPU

的数据传输。

pin_memory

(pin)

内存中，即把数据样本分配到固定的页锁定内存区域。

训练时，这些内存区域的数据会被高效地拷贝到对应

参数配合使用：

punctuation">.epochs

class="token

punctuation">(train_dataloader

class="token

punctuation">(non_blocking

class="token

punctuation">(non_blocking

class="token

punctuation">(pred_prob

class="token

punctuation">.zero_grad

class="token

punctuation">.backward

class="token

punctuation">(train_dataloader

class="token

punctuation">(train_dataloader

class="token

punctuation">)

non_blocking，使得以下两者之间的操作得以重叠：

• 从

数据(真实标签)的传输

GPU

模型训练计算(或

punctuation">.ArgumentParser

class="token

punctuation">.add_argument

class="token

string">'--num-machines'

class="token

punctuation">.add_argument

class="token

string">'--num-gpu-processes'

class="token

punctuation">.add_argument

class="token

string">'--machine-id'

class="token

punctuation">.add_argument

class="token

punctuation">.add_argument

class="token

string">'--batch-size'

class="token

punctuation">.parse_args

class="token

punctuation">.world_size

punctuation">.num_gpu_processes

punctuation">.num_machines

os

punctuation">.environ

class="token

string">'MASTER_ADDR'

class="token

punctuation">.environ

class="token

string">'MASTER_PORT'

class="token

punctuation">.spawn

class="token

punctuation">(train

class="token

punctuation">.num_gpu_processes

class="token

punctuation">}

class="token

punctuation">)

用

torch.cuda.device_count()

GPU

number">128

至此，我们已完成关于使用

分布式模型训练的简要探讨。

上述代码同样适用于其他深度学习模型，当前深度学习模型大多采用

GPU

分布式训练方案。

此外，Horovod、DeepSpeed

PyTorch

id="_417">小结

在本节中，我们探讨了机器学习中一个重要的实践方面——如何优化模型训练过程，介绍了使用

PyTorch

href="https://blog.csdn.net/LOVEmy134611/article/details/149906996"

rel="noopener

nofollow">PyTorch实战（1）——深度学习（Deep

Learning）
PyTorch实战（2）——使用PyTorch构建神经网络
PyTorch实战（3）——PyTorch

vs.

TensorFlow详解
PyTorch实战（4）——卷积神经网络（Convolutional

Neural

Network，CNN）
PyTorch实战（5）——深度卷积神经网络
PyTorch实战（6）——模型微调详解
PyTorch实战（7）——循环神经网络
PyTorch实战（8）——图像描述生成
PyTorch实战（9）——从零开始实现Transformer
PyTorch实战（10）——从零开始实现GPT模型
PyTorch实战（11）——随机连接神经网络（RandWireNN）
PyTorch实战（12）——图神经网络（Graph

Neural

Network，GNN）
PyTorch实战（13）——图卷积网络（Graph

Convolutional

Network，GCN）
PyTorch实战（14）——图注意力网络（Graph

Attention

Network，GAT）
PyTorch实战（15）——基于Transformer的文本生成技术
PyTorch实战（16）——基于LSTM实现音乐生成
PyTorch实战（17）——神经风格迁移
PyTorch实战（18）——自编码器（Autoencoder，AE）
PyTorch实战（19）——变分自编码器（Variational

Autoencoder，VAE）
PyTorch实战（20）——生成对抗网络（Generative

Adversarial

Network，GAN）
PyTorch实战（21）——扩散模型（Diffusion

Model）
PyTorch实战（22）——MuseGAN详解与实现
PyTorch实战（23）——基于Transformer生成音乐
PyTorch实战（24）——深度强化学习
PyTorch实战（25）——使用PyTorch构建DQN模型