📢 转载信息
原文链接:https://machinelearningmastery.com/3-ways-to-speed-up-model-training-without-more-gpus/
原文作者:Shittu Olumide
在本文中,您将学习三种经过验证的方法,通过优化精度、内存和数据流,在不增加任何新GPU的情况下加速模型训练。
我们将涵盖的主题包括:
- 混合精度和内存技术如何安全地提高吞吐量
- 使用梯度累积以更大的“虚拟”批次进行训练
- 使用ZeRO进行分片和卸载,以便在现有硬件上容纳更大的模型
让我们不要再浪费时间了。

3 Ways to Speed Up Model Training Without More GPUs
Image by Editor
引言
训练大型模型可能会非常缓慢,人们的第一反应通常是要求更多的GPU。但额外的硬件并非总是可行的选项。预算和云限制等问题会阻碍这一需求。好消息是,有方法可以在不增加任何GPU的情况下显著加快训练速度。
加速训练不仅仅关乎原始计算能力;它关乎更有效地利用您已有的资源。大量的训练时间浪费在了内存交换、GPU空闲和未优化的数据管道上。通过改进代码和硬件的通信方式,您可以从训练运行中节省数小时甚至数天的时间。
方法 1:混合精度和内存优化
在无需新GPU的情况下加速训练的最简单方法之一是使用混合精度。现代GPU被设计成比标准的32位浮点数(FP32)更快地处理半精度(FP16)或bfloat16运算。通过使用较小的数据类型进行存储和计算,可以减少内存使用和带宽,从而允许更多数据同时装入GPU,这意味着操作完成得更快。
其核心思想很简单:
- 对大多数操作使用较低精度(FP16或BF16)
- 将关键部分(如损失缩放和少数累加操作)保持在全精度(FP32)以维持稳定性
如果操作得当,混合精度通常可以带来1.5 – 2倍的训练速度提升,同时准确性几乎没有下降。它原生支持于PyTorch、TensorFlow和JAX,并且大多数NVIDIA、AMD和Apple GPU现在都为此提供了硬件加速。
这是一个启用自动混合精度的PyTorch示例:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
|
# Mixed Precision Example (PyTorch)
import torch
from torch import nn, optim
from torch.cuda.amp import GradScaler, autocast
model = nn.Linear(512, 10).cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
for inputs, targets in dataloader:
optimizer.zero_grad()
with autocast(): # operations run in lower precision
outputs = model(inputs.cuda())
loss = nn.functional.cross_entropy(outputs, targets.cuda())
scaler.scale(loss).backward() # scaled to prevent underflow
scaler.step(optimizer)
scaler.update()
|
工作原理:
autocast()
自动为每个操作选择FP16或FP32GradScaler()
通过动态调整损失尺度来防止下溢- GPU执行速度更快,因为它在每个操作中移动和计算的字节更少
您也可以通过PyTorch的自动混合精度 (AMP) 或用于旧版设置的Apex 库在全局激活它。对于较新的设备(A100、H100、RTX 40 系列),bfloat16 (BF16) 通常比FP16更稳定。
内存优化与混合精度相辅相成。两个常见的技巧是:
- 梯度检查点 (Gradient checkpointing):仅保存关键的激活值,并在反向传播过程中重新计算其他激活值,用计算量换取内存
- 激活卸载 (Activation offloading):将不常使用的张量暂时移动到CPU内存中
在PyTorch中,可以使用以下方式启用它们:
1
|
from torch.utils.checkpoint import checkpoint
|
或者使用DeepSpeed、Hugging Face Accelerate或bitsandbytes自动配置。
何时使用:
- 如果您的模型内存占用紧张,或者批次大小较小
- 您使用的是较新的GPU(RTX 20 系列或更高版本)
- 您可以容忍训练过程中轻微的数值变化
通常可以预期获得30–100%的训练速度提升和高达50%的内存使用减少,具体取决于模型大小和硬件。
方法 2:梯度累积和有效批次大小技巧
有时,加速训练的最大障碍不是计算能力,而是GPU内存。您可能希望使用大批次来提高梯度稳定性,但GPU在达到该大小之前就会耗尽内存。
梯度累积巧妙地解决了这个问题。不是一次性处理一个巨大的批次,而是将其分成更小的微批次。对每个微批次执行前向和后向传播,累积梯度,并且仅在经过几次迭代后才更新模型权重。这使得您可以使用相同的硬件来模拟大批次训练。
以下是PyTorch中的实现方式:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
|
# Gradient Accumulation Example (PyTorch)
import torch
from torch import nn
from torch.cuda.amp import GradScaler, autocast
# Assumes `model`, `optimizer`, and `dataloader` are defined elsewhere
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()
accum_steps = 4 # accumulate gradients over 4 mini-batches
for i, (inputs, targets) in enumerate(dataloader):
with autocast(): # works nicely with mixed precision
outputs = model(inputs.cuda())
loss = criterion(outputs, targets.cuda()) / accum_steps # normalize
scaler.scale(loss).backward()
if (i + 1) % accum_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
|
工作原理:
- 损失除以累积步数,以保持梯度平衡
- 梯度在步与步之间存储在内存中,而不是被清除
- 经过
accum_steps
个微批次后,优化器执行一次更新
这个简单的更改允许您使用高达四到八倍大的虚拟批次,在不超出GPU内存的情况下提高稳定性和潜在的收敛速度。
重要性:
- 更大的有效批次减少了梯度更新中的噪声,改善了复杂模型的收敛性
- 您可以将其与混合精度结合以获得额外的优势
- 当内存而非计算能力是限制因素时,它尤其有效
何时使用:
- 当您遇到大批次的“内存不足”错误时
- 当您想要在不更换硬件的情况下获得大批次的优势时
- 当您的数据加载器或增强流水线能够跟上每次更新所需的多个微步骤时
方法 3:智能卸载和分片训练 (ZeRO)
随着模型规模的扩大,GPU内存成为主要的瓶颈,远早于计算能力。您可能拥有训练模型的原始能力,但没有足够的内存同时容纳其所有参数、梯度和优化器状态。这时就需要智能卸载和分片训练。
其思路是智能地分割和分配内存使用,而不是在每个GPU上复制所有内容。DeepSpeed和Hugging Face Accelerate等框架通过ZeRO(Zero Redundancy Optimizer)等技术来实现这一点。
ZeRO 的工作原理
通常,多GPU设置中每个GPU都保存了完整的模型参数、梯度和优化器状态的副本。对于大型模型来说,这非常浪费。ZeRO通过在设备之间分片这些状态来打破这种重复:
- ZeRO 阶段 1:分片优化器状态
- ZeRO 阶段 2:分片优化器状态和梯度
- ZeRO 阶段 3:分片所有内容,包括模型参数
现在,每个GPU只保留了总内存占用的一小部分,但它们仍然协同工作以完成完整的更新。这使得那些远超单个GPU内存容量的模型也能高效地进行训练。
简单示例 (DeepSpeed)
下面是一个启用ZeRO优化的基本DeepSpeed配置片段:
1
2
3
4
5
6
7
8
9
|
{
"train_batch_size": 64,
"fp16": { "enabled": true },
"zero_optimization": {
"stage": 2,
"offload_optimizer": { "device": "cpu", "pin_memory": true },
"offload_param": { "device": "cpu" }
}
}
|
然后在您的脚本中:
1
2
|
import deepspeed
model, optimizer, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config='ds_config.json')
|
它做了什么:
- 为更快的计算启用了混合精度(fp16)
- 激活了ZeRO阶段2,在设备之间分片优化器状态和梯度
- 当GPU内存紧张时,将未使用的张量卸载到CPU内存中
何时使用:
- 您正在训练一个大型模型(数亿或数十亿参数)
- 即使使用混合精度,您也会耗尽GPU内存
- 您正在使用多个GPU或分布式节点
附加提示
上述三种主要方法——混合精度、梯度累积和ZeRO卸载——可以在不增加硬件的情况下提供大部分性能提升。但还有一些较小、通常被忽视的优化措施可以带来明显的差异,尤其是在与主要方法结合使用时。
让我们看看其中一些在几乎所有训练设置中都适用的优化。
1. 优化数据管道
GPU利用率通常会下降,因为在下一个批次准备好处理之前,模型就完成了计算。解决方法是并行化和预取数据。
在PyTorch中,您可以通过调整DataLoader来提高数据吞吐量:
1
|
train_loader = DataLoader(dataset, batch_size=64, num_workers=8, pin_memory=True, prefetch_factor=4)
|
num_workers
使用多个CPU线程进行加载pin_memory=True
加快主机到GPU的传输速度prefetch_factor
确保在GPU请求之前准备好批次
如果您正在处理大型数据集,请将它们存储在针对顺序读取优化的格式中,例如WebDataset、TFRecord或Parquet,而不是纯粹的图像或文本文件。
2. 优化前分析
在应用高级技术之前,请找出您的训练循环实际花费时间的地方。框架提供了内置的分析器:
您通常会发现最大的瓶颈不是GPU,而是像数据增强、日志记录或缓慢的损失计算这样的操作。修复这些问题可以立即提高速度,而无需任何算法更改。
3. 使用提前停止和课程学习
并非所有样本在整个训练过程中贡献相同。提前停止 (Early stopping) 在性能趋于平稳后阻止不必要的Epoch。课程学习 (Curriculum learning) 从较简单的示例开始训练,然后引入更难的示例,帮助模型更快地收敛。
1
2
3
4
|
if validation_loss > best_loss:
patience_counter += 1
if patience_counter >= patience_limit:
break # early stop
|
这个小模式可以在大型数据集上节省数小时的训练时间,同时对准确性的影响极小。
4. 定期监控内存和利用率
了解模型实际使用了多少内存有助于平衡批次大小、累积和卸载。在PyTorch中,您可以使用以下方式记录GPU内存统计信息:
1
|
print(f"Max memory used: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")
|
像nvidia-smi、GPUtil或Weights & Biases系统指标这样的监控工具可以帮助尽早发现利用率低的GPU。
5. 智能组合技术
最大的收获来自于堆叠这些策略:
- 混合精度 + 梯度累积 = 更快更稳定的训练
- ZeRO卸载 + 数据管道优化 = 在没有内存错误的情况下训练更大的模型
- 提前停止 + 分析 = 减少浪费的Epoch
何时使用每种方法
为了让您更容易决定哪种方法适合您的设置,这里有一个比较我们介绍的三种主要技术的摘要表,包括它们预期的好处、最适合的场景和权衡。
方法 | 最适合 | 如何提供帮助 | 典型速度提升 | 内存影响 | 复杂性 | 关键工具/文档 |
---|---|---|---|---|---|---|
混合精度和内存优化 | 任何内存占用紧张的模型 | 使用较低精度 (FP16/BF16) 和更轻的张量来减少计算和传输开销 | 1.5 – 2× 训练速度更快 | 30–50% 内存减少 | 低 | PyTorch AMP, NVIDIA Apex |
梯度累积和有效批次大小 | 受GPU内存限制但需要大批次大小的模型 | 通过在较小批次之间累积梯度来模拟大批次训练 | 提高收敛稳定性;通过更少的重启带来间接的速度提升 | 中等额外内存(临时梯度) | 低 – 中 | DeepSpeed 文档, PyTorch 论坛 |
智能卸载和分片训练 (ZeRO) | 无法在GPU内存中容纳的超大模型 | 跨设备或CPU分片优化器状态、梯度和参数 | 10–30% 吞吐量提升;可训练 2–4× 更大的模型 | 释放了大部分GPU内存 | 中 – 高 | DeepSpeed ZeRO, Hugging Face Accelerate |
🚀 想要体验更好更全面的AI调用?
欢迎使用青云聚合API,约为官网价格的十分之一,支持300+全球最新模型,以及全球各种生图生视频模型,无需翻墙高速稳定,文档丰富,小白也可以简单操作。
评论区