【深度学习】模型训练时减少GPU显存占用

最新推荐文章于 2025-05-04 10:32:33 发布

原创

最新推荐文章于 2025-05-04 10:32:33 发布 · 1.8k 阅读

14 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能

训练过程中，显存的主要占用来自：

激活值（Activations）：前向传播过程中存储的中间计算结果。
梯度存储：用于反向传播。
权重和偏置参数。

一、梯度检查点（Gradient Checkpointing）

在深度学习中，反向传播需要保留前向传播过程中生成的中间激活值（activations）来计算梯度。

Gradient Checkpointing 提供了一种折衷方案，允许用时间换空间。

普通训练
- 在前向传播中，每一层的激活值会保存在内存中，供反向传播使用。
- 这样可以减少重复计算，但占用了大量显存。
Gradient Checkpointing
- 在前向传播时，只保存某些关键层（称为“检查点”）的激活值，而不是所有层的激活值。
- 在反向传播时，未保存的激活值会被重新计算。
- 这种方法减少了显存需求，但增加了计算量。

在 Hugging Face 的 Trainer 中，可以通过 TrainingArguments 启用：

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=16,
    gradient_checkpointing=True,  # 启用梯度检查点
    num_train_epochs=3,
    learning_rate=5e-5
)