使用 Intel® oneAPI 工具包优化分布式 AI 训练

深度学习工作负载的增长速度非常快。基于深度学习的算法处理海量数据，以识别图像分类、物体检测、时间序列预测等模式。随着数据可用性的增加，深度学习模型的复杂性也随之增加。像ResNet*和Visual Geometry Group (VGG)这样的模型拥有数百万个参数，执行的浮点运算次数高达数十亿。最近的模型，如Generative Pretrained Transformer 3 (GPT-3)和Bidirectional Encoder Representations from Transformers (BERT)，拥有数十亿到数万亿个参数。因此，训练深度学习模型是一个计算密集且耗时的过程。为了缩短解决方案的实现时间，除了优化单核和多核性能外，还可以考虑扩展到多个节点（例如，分布式训练）。这项任务可以通过拆分模型（模型并行）、拆分数据（数据并行）或这两种方案的组合（混合并行）来实现。

本文重点介绍在计算节点集群上对基于深度学习的算法进行调优和扩展。我们使用半监督生成对抗网络（S-GAN）进行图像分类来说明这一点。我们演示了在OpenMP*、TensorFlow*和Intel® MPI库中的各种调优选项。在单个节点上，我们的优化实现了2倍的速度提升。在八节点集群上进行扩展可实现16倍的总速度提升。在单个节点上，Intel MPI库的图像吞吐量（每秒图像数）比Open MPI*库高出27%；在八节点集群上，高出18%。许多这些性能提升是在没有代码修改的情况下实现的，这表明这些优化可以有效地应用于其他应用程序。  

S-GAN

监督学习需要大量的标记数据。标记和标注必须由人工专家手动完成，因此非常耗时且昂贵。半监督学习是一种使用标记数据和未标记数据来训练模型的技术。通常，标记数据点的数量远少于未标记数据点的数量。半监督学习利用数据中的模式和趋势进行分类。半监督学习是一种使用标记数据和未标记数据来训练模型的技术。通常，标记数据点的数量远少于未标记数据点的数量。半监督学习利用数据中的模式和趋势进行分类。

S-GAN通过生成模型生成数据点来应对对大量训练数据的需求。生成对抗网络（GAN）是一种通过图像判别器模型使用大型未标记数据集训练图像生成器模型的架构。GAN包含两个模型：生成模型和判别模型。它们在一个零和博弈中一起训练。生成器的任务是生成与数据集中存在的相似的数据。判别器的任务是在生成的数据中识别真实数据。S-GAN通过向分类任务添加监督判别器（分类器）来扩展GAN架构（图1）。这样就得到了一个在未见过的数据上泛化良好的分类器。

软件和硬件

实验中使用了三个Intel® oneAPI工具包（v2021.1）。

1. Intel® oneAPI Base Toolkit
2. Intel® oneAPI HPC工具包
3. Intel® oneAPI AI分析工具包

Intel® Distribution for Python*（v3.6），它使用Intel® oneAPI数学内核库和Intel® oneAPI数据分析库，用于加速核心Python数值包。S-GAN模型是使用Intel® TensorFlow*优化（v1.15）实现的。Horovod*（v0.20.2）用于分布式训练。Horovod依赖MPI进行节点间通信，因此比较了两种MPI库的性能：Intel MPI库和Open MPI（v4.0.5）。Intel® VTune™ Profiler用于分析性能。所有测试都在Intel的Endeavor集群上进行，使用Intel® Xeon® Platinum 8260L处理器，通过100 Gbps的Intel® Omni-Path Fabric连接（图2）。

调优方法

在优化应用程序之前测量基线性能是一个好习惯。性能分析工具有助于识别潜在的优化区域，例如线程、向量化、I/O、多节点通信等。我们在Intel VTune Profiler中使用了Application Performance Snapshot。Application Performance Snapshot分析显示，应用程序产生了过多的线程。有些线程来自OpenMP，而有些线程来自TensorFlow调用的Eigen库。OpenMP和Eigen线程的数量超过了每个节点的逻辑核心数，导致资源超额分配，这通常会损害性能。

选择最佳线程数

第一步是找到最佳线程数。我们通过设置OMP_NUM_THREADS环境变量，在单个计算节点上使用低分辨率图像（256 x 256像素）并限制训练轮数（两轮）以节省时间，测试了不同数量的OpenMP线程。我们发现50个线程提供了最佳性能，每秒处理19.67张图像（图3）。

接下来，我们使用TensorFlow线程API来控制线程数。inter_op_parallelism_threads的值指定了独立非阻塞操作使用的线程数，而intra_op_parallelism_threads的值指定了矩阵乘法和归约等单个操作使用的线程数。图4显示了这两个参数不同值下的性能。深绿色块表示良好性能，深红色块表示性能差。白色块表示运行失败。

对于每个节点上的单个MPI rank，发现inter_op_parallelism_threads和intra_op_parallelism_threads的最佳值分别为0和45，对应每秒12.34张图像。这低于使用OMP_NUM_THREADS（图3）达到的每秒19.67张图像的值，因此我们使用此环境变量来控制线程数，而不是使用TensorFlow线程API。

选择每个节点上的最佳MPI Rank数

该应用程序使用MPI和OpenMP混合并行，因此找到每个节点上OpenMP线程和MPI Rank的最佳组合非常重要。我们重复了图3的测试（例如，1 ≤ OMP_NUM_THREADS ≤ 96），但改变了每个节点上的MPI Rank数（21 ≤ ppn ≤24，其中ppn是每个节点的Rank数，如图5所示）。最佳性能（每秒26.3张图像）是在每个节点上使用8个MPI Rank和9个OpenMP线程时实现的。

解决内存泄漏问题

到目前为止描述的实验使用了低分辨率图像。对于更高分辨率的图像（1360 x 1360像素），每个工作进程的内存占用量显著增加，因此每个节点运行8个MPI Rank会导致内存不足错误。结果发现单个Rank的内存占用量超过129 GB。由于每个节点上只有约200 GB的DRAM可用，因此只能启动一个Rank，这对于双插槽机器来说不是最优选择（因为存在非统一内存访问（NUMA）问题）。通过Python实用程序Memory Profiler运行我们的应用程序，发现存在内存泄漏。在模型的前向传播过程中，约有56 GB的内存未被释放（图6，上方）。事实证明，这是TensorFlow中TensorFlow issue #33009和Keras*issue #13118的Keras Model.predict方法中一个已知的bug。根据TensorFlow社区的建议，将Model.predict替换为Model.predict_on_batch，这降低了每个进程的总内存消耗（图6，下方）。内存泄漏修复后，对于高分辨率图像，我们仍然只能在每个节点上启动两个Rank。我们没有进一步优化内存消耗，尽管这可能是可能的。

设置其他OpenMP*和Intel® MPI库环境变量

另外三个环境变量（KMP_BLOCKTIME, KMP_AFFINITY和I_MPI_PIN_DOMAIN）的最佳设置也被进行了探索。根据之前的性能建议，我们进行了表1所示的六个实验。请注意，设置#5中的环境变量设置也用于其他五个设置。设置#3提供了最佳的训练性能（图7）。

表1. 其他OpenMP和Intel MPI库环境变量的实验设置

选择MPI库

接下来，我们收集了之前调优实验的最佳设置，并比较了Intel MPI库与Open MPI在1360 x 1360像素的高分辨率图像下的性能，每个节点两个MPI Rank，33个OpenMP线程，共100个epoch。使用了以下Intel MPI库命令行：

$ OMP_NUM_THREADS=33 KMP_BLOCKTIME=1 I_MPI_PIN_DOMAIN=auto:compact \
    mpirun -n 2/4/8/16 -ppn 2 -f $HOSTFILE \
    python SGAN_pob.py --image_size 1360 --num_epochs 100 --save_checkpoints False

以下大致等效的运行时设置用于Open MPI*：

$ OMP_NUM_THREADS=33 KMP_BLOCKTIME=1 \
    mpirun - n 2/4/8/16 -hostfile $HOSTFILE - bind-to socket -map-by ppr:1:socket \
    --mca btl ^openib --mca pml ob1 --report-bindings - x LD_LIBRARY_PATH \
    python SGAN_pob.py --image_size 1360 --num_epochs 100 --save_checkpoints FALSE

在单节点和多节点场景下，Intel MPI库的性能均优于Open MPI（图8）。

尽管基于每秒图像数的并行扩展看起来是线性的，但基于时间的扩展是非线性的，这表明应用程序效率不高。进一步分析表明，评估的频率没有根据总节点数正确扩展。随着节点数的增加，评估的频率也随之增加。此外，这些评估是由MPI主Rank执行的，因此所有其他Rank都会等待直到主Rank完成评估。我们通过每15次训练迭代执行一次评估步骤来修复这个问题，而与节点数无关（图9和图10）。此修改未影响模型准确性。

结论

在本文中，我们提出了一种增量优化方法，以实现更好的S-GAN训练性能。通过运行时和源代码级别的优化，解决了内存问题、扩展问题和单节点性能效率低下的问题。使用Horovod在多节点集群上实现了S-GAN模型的分布式训练。比较了两种MPI库（Intel MPI库和Open MPI）。

图11总结了我们优化工作带来的性能提升。单节点优化实现了2倍的加速比，而多节点优化带来了进一步的8倍加速比，使总解决方案时间缩短了16倍，且模型准确性没有损失。