CN116501502A - 一种基于Pytorch框架的数据并行优化方法 - Google Patents

Abstract

本发明公开了一种基于Pytorch框架的数据并行优化方法，属于计算机技术领域，具体为：加载模型阶段；按照前一轮训练后各节点之间的性能比例，将对应份的数据划分并加载至对应节点中；各节点将加载的数据按照当前轮批大小，对待训练模型进行当前轮训练；第一轮训练中的批大小为总批大小与节点数的比值，第二轮训练中的批大小根据第一轮训练后的性能比例分配总批大小，第三轮及之后的训练中根据前一轮训练所占用的时间计算时间波动率，若大于波动率阀值，根据第一轮训练后的性能比例分配总批大小；否则，采用前一轮训练的批大小。本发明可及时应对集群性能波动，同时避免由于节点性能微小波动而频繁调整批大小，进而带来的额外性能损耗。

Description

一种基于Pytorch框架的数据并行优化方法

技术领域

本发明属于计算机技术领域，具体涉及一种基于Pytorch框架的数据并行优化方法。

背景技术

随着人工智能的快速发展，深度学习模型参数和训练数据越来越大。模型参数规模逐步提升至千亿、万亿，同时用于训练的数据量级也显著提升。虽然对数据进行训练会提高模型的能力，但训练数据、模型参数的增多也会导致训练速度的降低，因此需要对数据并行方向进行研究以减少分布式训练时间。

PyTorch（一个开源的Python机器学习库）是深度学习研究和应用中常使用的训练框架。其中，PyTorch将值转换成张量，张量是具有丰富数据操作集的通用n维数组；模块定义了从输入到输出的转换，其正向传递期间的行为由前向传播函数指定；模块可以包含张量作为参数。例如，线性模块包含权重参数和偏差参数，其前向传播函数通过将输入乘以权重并添加偏差来生成输出。

PyTorch中提供了一个数据并行（Distributed Data Paralle, DDP）框架，以实现在多个进程或节点之间的数据并行训练。在训练过程中，每个节点的模型和优化器都是独立的，DDP通过确保所有模型的副本从完全相同的状态开始，来保证正确性，并在每个反向传递后拥有相同的参数。因此，即使来自不同节点的优化器都是相互独立的，但可通过在每次迭代结束前将不同节点的梯度进行同步，再更新本地模型参数，使得每个节点在每次迭代训练后的模型都能保持一致。

目前业界已有提出对于DDP的优化方法，例如，“DLB: A dynamic load balancestrategy for distributed training of deep neural networks[J]. IEEETransactions on Emerging Topics in Computational Intelligence, 2022,1-11（《深度神经网络分布式训练的动态负载平衡策略》，《IEEE计算智能新兴课题汇刊》）”公开了一种DLB策略，即基于动态负载平衡策略进行深度学习的分布式训练方法，用于解决梯度同步等待问题。具体来说，首先根据各节点在上一次的训练耗时来评估该节点的性能，然后根据节点的当前性能自适应调整批大小和数据集大小，以减少节点之间的等待时间，从而提高集群利用率。

但DLB策略存在以下问题：

（1）在每个epoch结束时才为下一次epoch训练调整批大小，而对于epoch耗时长的训练任务可能不能及时调整。如训练一个epoch耗时1小时，如果在这1小时因为其他用户使用集群资源，造成性能波动，则不能及时的调整批大小。

（2）DLB算法仅根据节点性能划分批大小，并没有考虑该节点显存大小，机器性能好的节点可能被分配到过大训练批大小，导致显存溢出造成程序终止。

（3）在每训练完一次epoch就重新划分批大小，当集群性能只是微小波动时也频繁调整训练数据批大小，会带来额外性能损耗。

发明内容

为了解决上述现有技术中存在的技术问题，本发明提供了一种基于Pytorch框架的数据并行优化方法，通过更细粒度的动态调整批大小和数据集分区，可对集群性能波动做出及时调整，避免显存溢出问题，以及由机器性能微小波动带来的额外性能损耗。

本发明的技术方案如下：

一种基于Pytorch框架的数据并行优化方法，包括以下步骤：

步骤1、加载模型阶段：

将基于Pytorch框架编写的待训练模型加载至J个节点中，使各节点均拥有待训练模型副本，J个节点的初始性能相同；并设置训练参数，包括：初始数据集的等分数量N，总批大小b，波动率阀值T，显存阈值S，以及显存调整步进P；

步骤2、加载数据阶段：

在对待训练模型进行当前轮训练前，按照前一轮训练后J个节点之间的性能比例，将对应份的数据不重复地划分为J个数据集分区，并行加载至对应节点中；其中，各节点的性能高低与被划分的数据量正相关；对第一轮训练，按照J个节点之间的初始性能比例，将对应份的数据不重复地划分为J个数据集分区；

步骤3、各节点将加载的数据按照当前轮的批大小，对待训练模型进行当前轮训练，包括前向传播阶段、损失计算阶段、反向传播阶段和模型参数更新阶段，获得各节点在当前轮训练过程中所占用的时间；

其中，在第一轮的训练过程中，以总批大小b与节点数J的比值为当前轮的批大小；

在第二轮的训练过程中，先计算各节点在第一轮训练后的性能，按照J个节点之间的性能比例，基于总批大小b重新分配当前轮训练所需的各节点的批大小；

在第三轮及之后的训练过程中，先根据各节点在前一轮训练过程中所占用的时间，计算前一轮的时间波动率，若大于波动率阀值T，则计算各节点在前一轮训练后的性能，再按照J个节点之间的性能比例，基于总批大小b重新分配当前轮训练所需的各节点的批大小，各节点的性能高低与被分配的批大小正相关；否则，将前一轮训练采用的批大小作为当前轮训练所需的批大小；

并且，在第二轮及之后的训练过程中，还包括对各节点所需显存大小的调节过程，具体为：根据Pytorch框架运行所需显存大小以及待训练模型在训练过程中所占用显存大小，计算各节点所需的显存大小，对显存大小大于显存阈值S的节点，将前一轮训练后与前两轮训练后之间的性能增加量减少P%，并将所减少的P%按照性能比例分配至其他节点，得到各节点在前一轮训练后的新性能，根据各节点之间的新性能比例，再次基于总批大小b重新分配当前轮训练所需的各节点的批大小；

重新进行对各节点所需显存大小的调节过程，直至各节点的显存大小不超过显存阈值S；

步骤4、重复执行步骤2~3，直至加载第N份数据，并完成第N轮训练。

进一步地，计算第i,i=2,3,…, N-1轮的时间波动率time_volatility_i的公式为：

；

其中，tⁱ _j,j=1,2,…, J为第j个节点在第i轮训练过程中所占用的时间；t^i-1 _j,j=1,2,…, J为第j个节点在第i-1轮训练过程中所占用的时间。

进一步地，计算第j个节点在第i,i=2,3,…, N-1轮训练后的性能pⁱ _j, j=1,2,…,J的公式为：

；

其中，bⁱ _j,j=1,2,…, J为第j个节点在第i轮训练所采用的批大小。

进一步地，基于总批大小b重新分配第i+1,i=2,3,…, N-1轮训练所需的第j个节点的批大小bⁱ⁺¹ _j的计算公式为：

。

进一步地，步骤2中若按照前一轮训练后J个节点之间的性能比例划分的J个数据集分区的数据量非整数，则对各数据集分区的数据量向下取整，将剩余未被划分的数据丢弃。

进一步地，步骤3还包括对bⁱ⁺¹ _j进行整数化处理的过程，具体为：

先对bⁱ⁺¹ _j向上取整，得到所有向上取整后的bⁱ⁺¹ _j总和，与总批大小b相差d，将各节点的性能pⁱ _j升序排列；对前d个节点，在对应向上取整后的批大小基础上减1，作为整数化后的第i+1轮训练所需的批大小；对其他节点，将对应向上取整后的批大小作为整数化后的第i+1轮训练所需的批大小。

进一步地，步骤1中波动率阀值T的确定过程具体为：

在训练前，先对待训练模型运行预设的脚本，并在脚本执行期间以固定步进调整待确定波动率阀值，得到各待确定波动率阀值所对应的加速比，取最大加速比所对应的待确定波动率阀值作为波动率阀值T。

进一步地，步骤1中显存阈值S根据硬件显存大小确定，P的取值为20。

进一步地，J、N和b均为大于等于1的正整数。

与现有技术相比，本发明的有益效果如下：

1、本发明提出了一种基于Pytorch框架的数据并行优化方法，根据每轮训练后的时间波动情况，判断是否需要动态调整各节点下一轮训练所需的批大小，可更及时地应对集群性能波动，同时避免由于节点性能微小波动而频繁调整批大小，进而带来的额外性能损耗；

2、本发明在重新分配数据集分区和训练批大小后不直接进行训练，而是先判断训练需显存是否溢出，再进行训练，避免性能好的节点可能被分配到过大训练批大小，而导致显存溢出造成程序终止的问题。

附图说明

图1为本发明实施例1提出的基于Pytorch框架的数据并行优化方法的流程图。

具体实施方式

下面结合附图和实施例，详述本发明的技术方案。

实施例1

本实施例提出了一种基于Pytorch框架的图像数据并行优化方法，流程如图1所示，包括以下步骤：

步骤1、加载模型阶段：

服务节点将基于Pytorch框架编写的待训练模型加载至J个节点中，使各节点均拥有待训练模型副本，J个节点的初始性能相同；并设置训练参数，包括：初始图像数据集的等分数量N，总批大小b，波动率阀值T，显存阈值S，显存调整步进P；

其中，本实施例针对的待训练模型为图像识别模型；

步骤2、令i=1，此时J个节点的性能相同，第i轮训练时第j个节点的批大小bⁱ _j =b/J, j=1,2,…, J；

步骤3、加载数据阶段：

按照J个节点之间的性能比例，将第i份数据不重复地划分为J个数据集分区，并行加载至对应节点中；其中，各节点的性能高低与被划分的数据量正相关；若按性能比例划分的J个数据集分区的数据量非整数，则对各数据集分区的数据量向下取整，将剩余未被划分的数据丢弃；

步骤4、第j, j=1,2,…, J个节点将加载的数据按照批大小bⁱ _j，对待训练模型进行第i轮训练，包括前向传播阶段、损失计算阶段、反向传播阶段和模型参数更新阶段，获得第j个节点在第i轮训练过程中所占用的时间tⁱ _j, j=1,2,…, J；

其中，在反向传播阶段，待训练模型会生成与各节点对应的局部梯度，发送至服务节点，服务节点通过集合通信库的All-Reduce操作进行聚合，计算得到平均梯度；在模型参数更新阶段，服务节点再将所得平均梯度更新至各节点；

步骤5、判断i是否大于1，若是，则计算第i轮的时间波动率time_volatility_i：

（1）；

转至步骤6；否则，转至步骤7；

步骤6、判断时间波动率time_volatility_i是否大于波动率阀值T，若是，则转至步骤7；否则，第i+1轮训练保持第i轮训练所需的各节点的批大小，转至步骤8；

步骤7、计算第j个节点在第i轮训练后的性能pⁱ _j, j=1,2,…, J：

（2）；

按照第i轮训练后J个节点之间的性能比例，基于总批大小b重新分配第i+1轮训练所需的各节点的批大小bⁱ⁺¹ _j，各节点的性能高低与被分配的批大小正相关，具体为：

计算bⁱ⁺¹ _j：

（3）；

对bⁱ⁺¹ _j向上取整，得到所有向上取整后的bⁱ⁺¹ _j总和，与总批大小b相差d，将各节点的性能pⁱ _j升序排列；对前d个节点，在对应向上取整后的批大小基础上减1，作为整数化后的第i+1轮训练所需的批大小；对其他节点，将对应向上取整后的批大小作为整数化后的第i+1轮训练所需的批大小；

转至步骤8；

步骤8、对各节点在第i轮训练后所需显存大小进行调节，具体为：基于Torch-summary，根据Pytorch框架运行所需显存大小以及待训练模型在训练过程中所占用显存大小，待训练模型在训练过程中所占用显存大小包括待训练模型自身结构参数和待训练模型计算过程中所产生的中间变量（与各节点被分配的数据集分区大小有关），计算各节点所需的显存大小，对显存大小大于显存阈值S的节点，将第i轮训练后与第i-1轮训练后之间的性能增加量减少20%，并将所减少的20%按照性能比例分配至其他节点，得到各节点在第i轮训练后的新性能，根据各节点之间的新性能比例，再次基于总批大小b重新分配第i+1轮训练所需的各节点的批大小bⁱ⁺¹ _j；其中i=1时，第i-1轮训练后的性能为初始性能；

重新进行对各节点在第i轮训练后所需显存大小的调节过程，直至各节点的显存大小不超过显存阈值S，转至步骤9；

步骤9、判断i是否等于N，若是，则对待训练模型的分布式训练完成；否则，令i=i+1，转回步骤3。

示例地，对J=3，b=128的情形，3个节点在第i轮训练过程中所占用的时间为{1, 2,3}，在第i轮训练过程中对应分配的批大小为{32, 32, 64}；根据式（2），计算得到在第i轮训练后对应的性能{32, 16, 64/3}；根据式（3），按照3个节点之间的比例32：16：64/3，计算得到3个节点在第i+1轮训练所需的批大小为{59.07, 29.5, 39.3}，向上取整后为{60,30, 40}，求得差值d为2，将各节点的性能升序排列，将前2个节点向上取整后的批大小减1，得到最终3个节点在第i+1轮训练所需的批大小为{60, 29, 39}。

尽管上面对本发明说明性的具体实施方式进行了描述，以便于本技术领域的技术人员理解本发明，但应该清楚，本发明不限于具体实施方式的范围，对本技术领域的普通技术人员来讲，只要各种变化在所附的权利要求限定和确定的本发明的精神和范围内，这些变化是显而易见的，一切利用本发明构思的发明创造均在保护之列。

Claims (9)

1.一种基于Pytorch框架的数据并行优化方法，其特征在于，包括以下步骤：

步骤1、加载模型阶段：

将基于Pytorch框架编写的待训练模型加载至J个节点中，J个节点的初始性能相同；并设置训练参数，包括：初始数据集的等分数量N，总批大小b，波动率阀值T，显存阈值S，以及显存调整步进P；

步骤2、加载数据阶段：

在对待训练模型进行当前轮训练前，按照前一轮训练后J个节点之间的性能比例，将对应份的数据划分为J个数据集分区，并行加载至对应节点中；其中，各节点的性能高低与被划分的数据量正相关；对第一轮训练，按照J个节点之间的初始性能比例，将对应份的数据划分为J个数据集分区；

步骤3、各节点将加载的数据按照当前轮的批大小，对待训练模型进行当前轮训练，包括前向传播阶段、损失计算阶段、反向传播阶段和模型参数更新阶段，获得各节点在当前轮训练过程中所占用的时间；

其中，在第一轮的训练过程中，以总批大小b与节点数J的比值为当前轮的批大小；

在第二轮的训练过程中，先计算各节点在第一轮训练后的性能，按照J个节点之间的性能比例，基于总批大小b重新分配当前轮训练所需的各节点的批大小；

在第三轮及之后的训练过程中，先根据各节点在前一轮训练过程中所占用的时间，计算前一轮的时间波动率，若大于波动率阀值T，则计算各节点在前一轮训练后的性能，再按照J个节点之间的性能比例，基于总批大小b重新分配当前轮训练所需的各节点的批大小，各节点的性能高低与被分配的批大小正相关；否则，将前一轮训练采用的批大小作为当前轮训练所需的批大小；

并且，在第二轮及之后的训练过程中，还包括对各节点所需显存大小的调节过程，具体为：根据Pytorch框架运行所需显存大小以及待训练模型在训练过程中所占用显存大小，计算各节点所需的显存大小，对显存大小大于显存阈值S的节点，将前一轮训练后与前两轮训练后之间的性能增加量减少P%，并将所减少的P%按照性能比例分配至其他节点，得到各节点在前一轮训练后的新性能，根据各节点之间的新性能比例，再次基于总批大小b重新分配当前轮训练所需的各节点的批大小；

重新进行对各节点所需显存大小的调节过程，直至各节点的显存大小不超过显存阈值S；

步骤4、重复执行步骤2~3，直至加载第N份数据，并完成第N轮训练。

2.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，步骤3还包括对bⁱ⁺¹ _j进行整数化处理的过程，具体为：

先对bⁱ⁺¹ _j向上取整，得到所有向上取整后的bⁱ⁺¹ _j总和，与总批大小b相差d，将各节点的性能pⁱ _j升序排列；对前d个节点，在对应向上取整后的批大小基础上减1，作为整数化后的第i+1轮训练所需的批大小；对其他节点，将对应向上取整后的批大小作为整数化后的第i+1轮训练所需的批大小。

3.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，步骤1中波动率阀值T的确定过程具体为：

在训练前，先对待训练模型运行预设的脚本，并在脚本执行期间以固定步进调整待确定波动率阀值，得到各待确定波动率阀值所对应的加速比，取最大加速比所对应的待确定波动率阀值作为波动率阀值T。

4.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，计算第i,i=2,3,…, N-1轮的时间波动率time_volatility_i的公式为：

；

其中，tⁱ _j,j=1,2,…, J为第j个节点在第i轮训练过程中所占用的时间；t^i-1 _j,j=1,2,…, J为第j个节点在第i-1轮训练过程中所占用的时间。

5.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，计算第j个节点在第i,i=2,3,…, N-1轮训练后的性能pⁱ _j, j=1,2,…, J的公式为：

；

其中，bⁱ _j,j=1,2,…, J为第j个节点在第i轮训练所采用的批大小。

6.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，基于总批大小b重新分配第i+1,i=2,3,…, N-1轮训练所需的第j个节点的批大小bⁱ⁺¹ _j的计算公式为：

。

7.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，步骤2中若按照前一轮训练后J个节点之间的性能比例划分的J个数据集分区的数据量非整数，则对各数据集分区的数据量向下取整，将剩余未被划分的数据丢弃。

8.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，步骤1中显存阈值S根据硬件显存大小确定，P的取值为20。

9.根据权利要求1所述基于Pytorch框架的数据并行优化方法，其特征在于，J、N和b均为大于等于1的正整数。