梯度检查点

前向传播通常会缓存所有中间激活值，以便反向传播时复用。然而，激活值的规模会随着批大小和序列长度的增加而增大。梯度检查点仅保存部分激活值，并丢弃其余部分。这迫使反向传播在需要时即时重新计算部分激活值。

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

训练通常会更慢，因为反向传播需要重新计算被丢弃的激活值，但检查点机制可以减少激活值的内存占用。

设置 gradient_checkpointing=True 以启用。

结合使用梯度累积以进一步降低内存占用。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

部分检查点

完整的梯度检查点会重新计算所有可检查点的层。如果你的运行环境有一定的内存余量，可以减少检查点的层数，以部分内存节省换取速度。

将 every_n_layers 传递给 gradient_checkpointing_enable() 以选择检查点间隔。

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

当 every_n_layers=2 时，第一层以及其后每隔一层的层都会被检查点化。被检查点化的层在前向传播过程中会丢弃其激活值，并在反向传播过程中重新计算它们，而其他层则会在内存中保留其激活值。

model.gradient_checkpointing_enable(every_n_layers=2)

默认值 every_n_layers=1 会对每一层进行检查点。较大的值会对第一层进行检查点，然后每隔 n 层再进行一次，其余层的激活值将保留在内存中。例如，every_n_layers=2 会对第 1、3、5 层等进行检查点。只有继承自 GradientCheckpointingLayer 的模块才会被计入。其他支持梯度检查点的模块将保持启用状态。

要在 Trainer 中使用部分梯度检查点，请在 gradient_checkpointing_kwargs 中设置 every_n_layers。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

卸载已保存的激活值

梯度检查点会在 GPU 上为每个被检查点的层保留一个激活值。在长序列长度下，这会消耗大量显存（大约为 layers x sequence x hidden x bytes_per_element）。将 offload 设置为 true，可将这些激活值保存在固定的主机内存中，代价是在前向传播期间进行一次设备到主机的拷贝，并在反向传播期间进行一次主机到设备的拷贝。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

这两份副本都在计算流上运行，因此这是以较慢的步骤速度换取内存。当运行因其他原因无法容纳时再使用它，而不是为了加速。

后续步骤

阅读GPU 内存使用文档，以了解训练期间 GPU 内存使用量的驱动因素。
请参阅混合精度训练指南，了解如何使用较低精度的数据类型来进一步减少内存占用并加快训练速度。
请参阅 Kernels 指南，了解如何使用自定义融合内核加速训练。

Update on GitHub

Transformers

梯度检查点

部分检查点

卸载已保存的激活值

后续步骤