【diffusers极速入门（六）】缓存梯度和自动放缩学习率以及代码详解

系列文章目录

• 【diffusers 极速入门（一）】pipeline 实际调用的是什么？ call 方法!
• 【diffusers 极速入门（二）】如何得到扩散去噪的中间结果？Pipeline callbacks 管道回调函数
• 【diffusers极速入门（三）】生成的图像尺寸与 UNet 和 VAE 之间的关系
• 【diffusers极速入门（四）】EMA 操作是什么？
• 【diffusers极速入门（五）】扩散模型中的 Scheduler（noise_scheduler）的作用是什么？

---

---

在深度学习训练过程中，缓存梯度和自动放缩学习率是两个非常重要的技术，它们能够帮助提高训练效率和稳定性。以下是对这两个技术的详细解释以及对应的代码解释。

本文解释的代码来源于 /path/to/diffusers/examples/dreambooth/train_dreambooth_sd3.py

一、缓存梯度是什么？

标准训练过程： 前向传播和反向传播各进行一次，没有额外的计算。
缓存梯度技术： 前向传播仅进行一次，但反向传播过程中由于需要重新计算未存储的激活值，导致反向传播时间增加。

• 缓存梯度（Gradient Checkpointing）是一种技术，通过在前向传播过程中存储部分中间激活值，减少内存占用量。
• 这样可以在反向传播时重新计算这些激活值，以节省显存。
• 这种方法通常会导致反向传播速度变慢，但可以训练更深和更大的模型。

为什么只是反向传播变慢？

1. 训练过程中，标准的前向和反向传播过程：
   • 前向传播（Forward Pass）： 计算模型的输出，同时存储所有中间激活值（即每一层的输出）。
   • 反向传播（Backward Pass）： 使用存储的中间激活值计算梯度，并更新模型参数。
2. 缓存梯度技术的工作原理
   • 其主要目标是减少内存占用。
   • 它的工作原理如下：
     • 前向： 仅存储部分关键的中间激活值（检查点），其余的中间激活值不存储。
     • 反向： 当需要用到未存储的中间激活值时，重新计算这些激活值。
3. 为什么反向传播变慢？
   因为需要重新计算激活值。在反向传播过程中，当需要计算梯度时，必须重新计算之前未存储的中间激活值。这意味着：每次反向传播到一个未存储的层时，需要从最近的检查点重新进行前向传播，计算出该层的激活值。这种重新计算会导致额外的计算开销。
4. 具体示例如下
   假设我们有一个简单的三层网络，使用缓存梯度技术并在第二层处设置检查点：

前向传播

前向传播

前向传播

反向传播:计算第3层梯度

反向传播: 计算第2层梯度

反向传播: 计算第1层梯度

输入数据

第1层激活值-已存储

第2层:激活值-检查点已存储

第3层: 计算但不存储激活值

• 如果我们前向只在第2层设置检查点，计算第3层梯度时无需额外计算，但计算第1层梯度时需要重新计算第2层的激活值。
• 这意味着每次反向传播到第1层时，都需要重新进行一次前向传播到第2层，增加了计算量。

相关代码

if args.gradient_checkpointing:
    transformer.enable_gradient_checkpointing()
    if args.train_text_encoder:
        text_encoder_one.gradient_checkpointing_enable()
        text_encoder_two.gradient_checkpointing_enable()
        text_encoder_three.gradient_checkpointing_enable()
...
parser.add_argument(
    "--gradient_checkpointing",
    action="store_true",
    help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
)

• 检查是否启用缓存梯度： if args.gradient_checkpointing: 检查命令行参数中是否启用了缓存梯度。
• 启用缓存梯度：
  • transformer.enable_gradient_checkpointing()：为 transformer 模型启用缓存梯度。
  • text_encoder_one.gradient_checkpointing_enable() 等：为多个文本编码器启用缓存梯度。
• 命令行参数： --gradient_checkpointing 是一个布尔类型参数，用于指示是否启用缓存梯度。启用后，可以通过 action=“store_true” 将其设置为 true。

二、自动放缩学习率

• 自动放缩学习率是一种根据训练配置自动调整初始学习率的方法。
• 当使用多个 GPU 或者增加批量大小时，需要相应地调整学习率以保持训练的稳定性和效果。
• 常见的做法是将学习率（learning_rate）乘以一个因子，该因子与 GPU 数量（num_processes）、梯度累积步数（gradient_accumulation_steps）和批量大小（train_batch_size）相关。
• 具体如下方代码所示

if args.scale_lr:
    args.learning_rate = (
        args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
    )
...
parser.add_argument(
    "--scale_lr",
    action="store_true",
    default=False,
    help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
)
parser.add_argument(
    "--learning_rate",
    type=float,
    default=1e-4,
    help="Initial learning rate (after the potential warmup period) to use.",
)
parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
parser.add_argument(
        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
    )

• 检查是否启用自动放缩学习率： if args.scale_lr: 检查命令行参数中是否启用了自动放缩学习率。
• 计算新学习率： 根据梯度累积步数、批量大小和 GPU 数量（由 accelerator.num_processes 提供）调整学习率。
  • 通过 args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes 计算得到新的学习率。
  • 命令行参数：
    • –scale_lr 是一个布尔类型参数，用于指示是否启用自动放缩学习率，默认值为 False。
    • –learning_rate 用于指定初始学习率，默认值为 1e-4。

总结

• 缓存梯度（Gradient Checkpointing）： 通过在前向传播过程中存储部分中间激活值，减少显存占用，代价是反向传播速度变慢。适用于需要训练非常深或大的模型的情况。
• 自动放缩学习率（Learning Rate Scaling）： 根据训练配置（如 GPU 数量（num_processes）、梯度累积步数和批量大小）自动调整学习率，以适应不同的训练环境和配置，确保训练的稳定性和效果。