HuggingFace Transformer库如何在训练时降低GPU显存

---

0. 引言

Hugging Face Transformers 是一个非常流行的库，用于处理自然语言处理任务。在使用该库进行模型训练时，可能会遇到显存（GPU内存）不足的问题。本文介绍可以尝试的方法来降低显存占用，

1. 减少批次大小 (Batch Size)

批次大小是影响显存占用的主要因素之一。减小批次大小可以显著减少显存消耗，但可能会延长训练时间。

TrainingArguments(
	per_device_train_batch_size=2/4/...  # 每个设备上的训练批次大小
)

如果你有多个GPU，则实际的每轮训练批次大小将是每个设备上的批次大小乘以GPU的数量。

2. 使用 Gradient Accumulation

当你减小了批次大小但仍然希望保持较大的有效批次大小时，可以使用梯度累积（Gradient Accumulation）。这意味着在更新权重之前，会累积多个小批次的梯度。

TrainingArguments(
	gradient_accumulation_steps=2/4/...  # 每个设备上的训练批次大小
)

比如我们设置了 per_device_train_batch_size 为4，表示每个 GPU 设备上的批次大小为4。同时，我们设置了 gradient_accumulation_steps 为2，这意味着在更新权重之前，会累积两个小批次的梯度。因此，尽管每个设备上的实际批次大小为4，但通过累积两个批次的梯度，我们可以得到一个相当于每轮训练8个样本的效果。

3. 混合精度训练 (Mixed Precision Training)

使用 FP16 或 BF16 进行混合精度训练可以在保证训练效果的同时大幅减少显存占用。Hugging Face 提供了与 PyTorch 和 TensorFlow 的混合精度训练支持。

4. 序列长度 (Sequence Length)

减小输入序列的长度也可以降低显存消耗。可以通过截断或摘要等方法来实现。

5. 使用 DeepSpeed

DeepSpeed 是一个深度学习优化库，特别适合大规模模型训练。它提供了许多功能，如 ZeRO、ZeRO-Offload 和 Pipeline Parallelism，这些都可以帮助减少显存需求。

ZeRO 的原理请参考：《大模型训练ZeRO内存优化原理详解》

6. 减少CUDA内存分配块的大小

设置 PyTorch 中 CUDA 分配器的行为，告诉 PyTorch 在分配 GPU 内存时采用一种特定的策略，以减少碎片化并提高内存管理效率。

os.environ["PYTOCH_CUDA_ALLOC_CONF"]="max_split_size_mb:128"

max_split_size_mb:128 这意味着当请求的内存小于 128MB 时，PyTorch 将尝试将内存分配成最大不超过128MB 的块。

默认情况下，PyTorch 会尽可能地分配一个连续的大块内存来满足请求，但这可能导致碎片化问题。通过限制分配的最大块大小，可以使得较小的内存请求更容易找到可用的空间，从而减少碎片化的发生。

但太小的 max_split_size_mb 可能会引发频繁分配内存，会降低工作效率。

---

欢迎关注本人，我是喜欢搞事的程序猿； 一起进步，一起学习；

欢迎关注知乎/CSDN：SmallerFL

也欢迎关注我的wx公众号（精选高质量文章）：一个比特定乾坤