系列文章目录
- 【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)是一种技术,通过在前向传播过程中存储部分中间激活值,减少内存占用量。
- 这样可以在反向传播时重新计算这些激活值,以节省显存。
- 这种方法通常会导致反向传播速度变慢,但可以训练更深和更大的模型。
为什么只是反向传播变慢?
- 训练过程中,标准的前向和反向传播过程:
- 前向传播(Forward Pass): 计算模型的输出,同时存储所有中间激活值(即每一层的输出)。
- 反向传播(Backward Pass): 使用存储的中间激活值计算梯度,并更新模型参数。
- 缓存梯度技术的工作原理
- 其主要目标是减少内存占用。
- 它的工作原理如下:
- 前向: 仅存储部分关键的中间激活值(检查点),其余的中间激活值不存储。
- 反向: 当需要用到未存储的中间激活值时,重新计算这些激活值。
- 为什么反向传播变慢?
因为需要重新计算激活值。在反向传播过程中,当需要计算梯度时,必须重新计算之前未存储的中间激活值。这意味着:每次反向传播到一个未存储的层时,需要从最近的检查点重新进行前向传播,计算出该层的激活值。这种重新计算会导致额外的计算开销。 - 具体示例如下
假设我们有一个简单的三层网络,使用缓存梯度技术并在第二层处设置检查点:
- 如果我们前向只在第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)、梯度累积步数和批量大小)自动调整学习率,以适应不同的训练环境和配置,确保训练的稳定性和效果。