如何将 AI 模型的训练时间缩短

好者，无法使用 H100 等高端 AI GPU。大多数人都知道在自己的消费级 GPU 上训练模型并等待很长时间才能完成 1 个微不足道的时期是多么困难。

在我的硕士论文期间，我能够将持续时间减少到其初始持续时间的仅 4.65％，换句话说，我实现了 2153％ 的加速。

背景故事

对于我的论文，压缩数据集为 472 GB，我计划用这些数据训练 32 个模型。在设置好一切并测试训练流程后，我认为这在我使用的 2080Ti GPU 上行不通，因为一次训练迭代就花了10 小时 40 分 04 秒……我预计至少需要 50 次迭代才能得到有用的结果。具体来说：
32 个模型 x 10.7 小时 x 50 次迭代 = 713.33 天。

策略一：更多 GPU

在强调了 AI 爱好者在有限的计算能力方面所面临的困难后，这一点似乎有点违反直觉。但让我们先面对现实，我们谈论的是 AI，如果你想在这个领域有所成就，就会涉及一些成本。但 4 个翻新的 RTX 2080 TI（总价 480 欧元）与一个低档 AI GPU（约 6-7000 欧元）之间存在显著差异。

回到故事，我决定在多节点Slurm 集群配置中使用 4 个 RTX 2080 TI 。这意味着我有四台每台配备一个 GPU 的 PC，而不是一台配备 4 个 GPU 的 PC。最初，这是相当多的开销，但最终，值得设置它。一个好处是您可以关闭当前不需要的节点并节省能源。另一个好处是每台 PC 仍然可以单独用于其他事情，例如游戏、photoshop、渲染或其他 GPU 密集型任务。

一切设置和配置完成后，我使用PyTorch 进行数据并行处理，首次提速 363% 。为了让它更容易，我将代码移植到PyTorch-lightning，让训练器处理数据的分布。

让我向你展示一个简化的 PyTorch-lighting 示例：

trainer_config = { 
  “accelerator” : “gpu” , 
  “precision” : “32” , 
  “devices” : 1 ,      #每个节点的 GPU 
  “num_nodes” : 4 ,    #节点数
  “max_epochs” : 50
 } 

mock_data_config_train = { 
  'batch_size' : 32 , 
  'shuffle' : True , 
  'num_workers' : 4 , #每个节点 4x GPU
 } 

mock_data_config_test = { 
  'batch_size' : 32 , 
  'shuffle' : False , 
  'num_workers' : 4 , #每个节点 4x GPU
 } 

trainer = pl.Trainer(**trainer_config) 
train_dl = DataLoader(train_data, **mock_data_config_train) 
valid_dl = DataLoader(valid_data, **mock_data_config_test) 
model = ... # 你的 PyTorch-lighting 包装模型

trainer.fit( 
  model=lightning_model, 
  train_dataloaders=train_dl, 
  val_dataloaders=validation_dl 
)

现在，Slurm 提交脚本在简单的 bash 中：

#!/bin/bash -l 

# SLURM 提交脚本
#SBATCH --nodes=4 
#SBATCH --gres=gpu:1 
#SBATCH --ntasks-per-node=1 
#SBATCH --output=var/logs/run.out 


. /opt/conda/etc/profile.d/conda.sh 
conda activate ddp_env 
srun python /var/train/train.py #我们上面的脚本

策略二——优化数据处理

为了进一步优化训练速度，我研究了数据加载器。因为我正在处理视频，所以我需要处理大量数据，然后模型才能对数据进行训练。这发生在与训练并行的第二个任务中，但由于数据量大且处理复杂，这大大减慢了训练速度（呃）。

让我来告诉你我的意思……在图 1 中，你可以看到缓慢的数据处理对训练的影响。每个黄色块代表一批处理的数据，每个蓝色块代表一个训练步骤。

图 1. 数据处理速度慢时的处理时间。图片由作者提供。

如您所见，如果批次的准备时间比前一个训练步骤更长，则模型必须等到下一个批次准备就绪。

理想情况下，您希望数据处理速度至少与训练步骤一样快，以确保训练过程顺利进行。这可以在图 2 中看到。

图 2. 数据处理速度足够快时的处理时间。图片由作者提供。

要确定数据处理是否是瓶颈，请比较训练数据加载器和测试数据加载器在一个训练步骤中花费的时间。
如果它们的持续时间大致相同，则数据加载器不是管道中的瓶颈。

如果您的数据加载器是一个瓶颈，那么第一步是弄清楚哪些组件可以通过额外的一次预处理步骤来完成。

让我们具体一点

在我训练的模型的数据加载器中，我使用预先训练的 MCTNN 人脸识别模型以 30% 的余量裁剪人脸，然后从每张图片中提取人脸标志。然后使用这些人脸标志计算凸包和感兴趣区域，以对某些面部特征应用专门的数据增强。完成这些增强后，生成的图像在第二个增强管道中进行处理，以另外应用更通用的Albumentation 变换。从这个描述中你就可以得出，数据处理功能非常复杂，计算量很大。

为了减少数据工作进程的持续时间，我预先计算了所有可能的情况，并将原始值以可用于应用增强的格式存储在新的parquet 数据集中。在一个节点上完成此预计算大约需要 10 个小时。您可以在图 3 中看到新设计。

图 3. 改进后的数据加载功能。图片由作者提供

与最初的一个时期的持续时间相比，新设计的速度提高了 1329%。

策略三——ZeRO 优化器 (DeepSpeed)

进一步减少一个 epoch 所需时间的下一步是使用微软开发的DeepSpeed 库。它还与 PyTorch-lightning 紧密集成。该库附带多个预定义配置：

• DeepSpeed ZeRO 第 1 阶段 — 分片优化器状态
• DeepSpeed ZeRO 第 2 阶段 — 分片优化器状态和梯度
• DeepSpeed ZeRO 第 3 阶段 — 分片优化器状态、梯度、参数和可选激活
• DeepSpeed ZeRO Stage 2 和 Stage 3 还配备了卸载配置，可进一步优化内存

现在你可能会问自己“内存优化？这对我有什么帮助？”。较低的内存消耗意味着你可以在训练期间使用更大的批处理大小，因此你需要更少地将数据传输到 GPU。但真正的加速发生在切换到 fp16 之后，这是 DeepSpeed 库强烈推荐的。使用 PyTorch-lightning 框架将这些更改应用于代码非常简单，只需要一行新代码。

trainer_config = { 
  "accelerator": "gpu", 
  "precision": "16", #更改了精度
  "devices": 1, #每个节点的 GPU 
  "num_nodes": 4, #节点数
  "strategy": "deepspeed_stage_2", #添加了 DeepSpeed Stage 
  "max_epochs": 50 
} 

mock_data_config_train = { 
  'batch_size': 32, 
  'shuffle': True, 
  'num_workers': 4, #每个节点 4x GPU 
} 

mock_data_config_test = { 
  'batch_size': 32, 
  'shuffle': False, 
  'num_workers': 4, #每个节点 4x GPU 
} 

trainer = pl.Trainer(**trainer_config) 
train_dl = DataLoader(train_data, **mock_data_config_train) 
valid_dl = DataLoader(valid_data, **mock_data_config_test) 
model = ... # 你的 PyTorch-lighting 包装模型

trainer.fit( 
  model=lightning_model, 
  train_dataloaders=train_dl, 
  val_dataloaders=validation_dl 
)

这两项变化都使速度提升至最终值的 2153%。

结论

经过所有的优化，结果可以总结如下：

表 1. 所有结果的摘要

使用上述所有策略，我能够在一个月内完成我的硕士论文项目，或者准确地说，大约 33 天，而不是 713 天。