【Hugging Face 开源库】Accelerate 库 —— 大模型训练框架

一杯水果茶！

已于 2025-03-27 17:03:42 修改

阅读量1.3k

点赞数 25

CC 4.0 BY-SA版权

分类专栏：视觉与网络文章标签： Accelerate Hugging face Adma

于 2025-03-26 14:10:09 首次发布

本文链接：https://blog.csdn.net/xiaoyuting999/article/details/146522291

视觉与网络专栏收录该内容

24 篇文章

订阅专栏

大模型训练框架 Accelerate 库
训练对象
- Adam 优化器
- 调度器

大模型训练框架 Accelerate 库

Hugging Face 的 Accelerate 是一个用于 简化和加速深度学习模型训练 的库，它支持在多种硬件配置上进行分布式训练，包括 CPU、GPU、TPU 等。

Accelerate 库的设计目标是 简化多 GPU 和 TPU 的训练过程，同时支持混合精度训练。混合精度训练是指 同时使用单精度（float32）和半精度（float16）进行训练，以提高训练速度和降低内存消耗。

Accelerate 使用步骤

Accelerate 只需添加四行代码，即可使 相同的 PyTorch 代码在任何分布式配置下运行！
让大规模训练和推理变得简单、高效且适应性强。

+ from accelerate import Accelerator
+ accelerator = Accelerator()  #　导入并创建一个 Accelerator 实例

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
      inputs = inputs.to(device)
      targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()

首先 导入并创建一个 Accelerator 实例。
Accelerator 还 知道将 PyTorch 对象移动到哪个设备，因此建议让加速器为我们处理此操作。
```
- device = "cuda"
+ device = accelerator.device
  model.to(device)
```
接下来，为分布式训练 准备 PyTorch 对象（模型、优化器、调度器等）
prepare() 方法负责 将模型放置在适当的容器中（如单个 GPU 或多 GPU） 以适应训练设置，将优化器和调度器调整为使用 Accelerate 的 AcceleratedOptimizer 和 AcceleratedScheduler，并创建一个新的数据加载器，该加载器可以在进程之间进行分片。
- NOTE： Accelerate 只 prepare() 继承自相应 PyTorch 类的对象，例如 torch.optim.Optimizer。
- 可以 prepare() 自定义的 模型对象，如 diffusion_model（需要满足 PyTorch Module 的要求）；Accelerate 将对其进行包装，以确保模型在后续调用中能够正确地使用 Accelerate 提供的功能。
最后，在训练循环中移除对输入和目标的 to(device) 调用，因为 Accelerate 的 DataLoader 类会自动将它们放置在正确的设备上。还应该用 Accelerate 的 backward() 方法替换常规的 backward()，该方法 缩放梯度并使用适当的 backward() 方法（例如，DeepSpeed 或 Megatron）。
```
-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)
```

将一切整合起来，新 Accelerate 训练循环现在应该看起来像这样！

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

Accelerate 提供额外功能——如梯度累积、梯度裁剪、混合精度训练等——可以将这些功能添加到脚本中以改善训练运行。

梯度累积（Gradient accumulation）

梯度累积 通过 在更新权重之前累积多个批次的梯度 来 在更大的批次大小上进行训练。这可以用来克服内存限制。

要在 Accelerate 中启用此功能，请在 Accelerator 类中指定 gradient_accumulation_steps 参数，并将 accumulate() 上下文管理器添加到脚本中。

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

梯度裁剪（Gradient clipping）

梯度裁剪是一种 防止“梯度爆炸” 的技术，Accelerate 提供了：

clipgrad_value() 用于将梯度限制在最小和最大值之间
clipgrad_norm() 用于将梯度规范化到特定值

混合精度（Mixed precision）

混合精度通过使用较低精度的数据类型如 fp16（半精度）来计算梯度，从而 加速训练。

为了在 Accelerate 中获得最佳性能，损失应该在模型内部计算（如 Transformers 模型），因为 模型外部的计算是以全精度进行的。

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

如果在创建 Accelerator 时就会将 mixed_precision 参数设置为 'no'，则 autocast() 仅起到一个透明包装的作用，不会进行混合精度计算。
NOTE： 只需要在前向计算中用 autocast() 包裹那些需要混合精度处理的操作，而 accelerator.backward() 会根据 mixed_precision 的设置自动处理梯度缩放。

Accelerate 会根据 mixed_precision 的设置自动管理 AMP 相关的梯度缩放（scaling），也就是说在调用 accelerator.backward(loss) 时，如果混合精度被启用（例如设置为 'fp16'），梯度缩放已经自动集成进去了。
autocast() 的主要作用是在 前向计算过程中控制操作的数据类型，从而获得混合精度的加速效果。所以，如果没有额外需要低精度执行的操作，backward() 内部的自动缩放已经足够，不需要额外手动处理混合精度。

训练对象

哪些是需要更新的训练参数？

模型参数：例如自定义的 diffusion_model 中的所有权重和偏置等，这些参数在每一步梯度下降时都会更新。
优化器状态：例如 Adam 优化器中维护的动量项、学习率调整等，都会随着训练更新。
EMA 模型的状态：虽然 EMA 的更新是一种滑动平均，但它的状态会在训练过程中定期根据模型参数进行更新。

哪些是不需要更新的配置参数（超参数）？

训练超参数：这些参数 在训练过程中是不变的，只用来控制训练流程和优化器行为。
- train_batch_size：每个训练批次的样本数，决定每步梯度更新时使用的样本量
- gradient_accumulate_every：梯度累积步数，用于模拟更大的批次训练，减少显存消耗
- train_lr：学习率，控制参数更新的步长
- train_num_steps：总的训练步数，决定整个训练过程的迭代次数
- ema_update_every：每隔多少步更新一次 EMA 模型
- ema_decay：EMA 衰减系数，用于平滑模型权重
- adam_betas：Adam 优化器的动量参数
- save_and_sample_every：每隔多少步保存模型和进行采样
- num_samples：用于采样时生成的样本总数，需满足整数平方根条件（便于组织采样输出）
- max_grad_norm：梯度剪裁阈值，用于防止梯度爆炸
调度器设置：scheduler，学习率调度器，用于动态调整学习率，它的配置也是在训练开始前设定的。
数据加载和日志记录配置：例如数据集、DataLoader 的设置以及结果保存的路径等，都属于配置项，不会参与梯度更新。