如何在训练 Huggingface Transformers (Pegasus) 模型时防止 VRAM 中途丢失？

Question

我正在通过 Huggingface 转换器采用预训练的 pegasus 模型（具体来说，

google/pegasus-cnn_dailymail

，我正在通过 Pytorch 使用 Huggingface 转换器），我想根据自己的数据对其进行微调。然而，这是一个相当大的数据集，我遇到了训练中途 VRAM 耗尽的问题，由于数据集的大小，甚至可能在训练开始后几天，这使得试错方法非常低效。

我想知道如何提前确保它不会耗尽内存。我认为模型的内存使用量在某种程度上与输入的大小成正比，所以我已经将

truncation=True

、

padding=True

、

max_length=1024

传递给我的分词器，如果我的理解是正确的，应该使每行相同大小的分词器的所有输出。考虑到批量大小也是一个常数，我认为使用的 VRAM 量应该是稳定的。所以我应该能够将数据集分割成可管理的部分，只看第一次运行的 ram/vram 使用情况，并推断它会从头到尾顺利运行。

然而，事实似乎恰恰相反。我一直在观察任何时候使用的 VRAM 数量，它可能变化很大，从一次约 12GB 到突然需要超过 24GB 并崩溃（因为我没有超过 24GB）。

那么，我如何确保在整个训练过程中使用的 vram 数量保持在合理范围内，并避免在我已经进入训练过程几天后由于缺少 vram 而导致崩溃？

Answer 1

padding=True

实际上并没有填充到

max_length

，而是填充到您传递给分词器的列表中最长的样本。要填充到

max_length

，您需要设置

padding='max_length'

.

Answer 2

有时 pytorch 不会收集所有垃圾。如果您尝试开始清理缓存而不是等到垃圾收集完毕，就会发生这种情况。所以你可以试试这个代码：

import torch
import time
import gc
from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo

def clear_gpu_memory():
    torch.cuda.empty_cache()
    gc.collect()

def wait_until_enough_gpu_memory(min_memory_available, max_retries=10, sleep_time=5):
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(torch.cuda.current_device())

    for _ in range(max_retries):
        info = nvmlDeviceGetMemoryInfo(handle)
        if info.free >= min_memory_available:
            break
        print(f"Waiting for {min_memory_available} bytes of free GPU memory. Retrying in {sleep_time} seconds...")
        time.sleep(sleep_time)
    else:
        raise RuntimeError(f"Failed to acquire {min_memory_available} bytes of free GPU memory after {max_retries} retries.")

# Usage example
min_memory_available = 2 * 1024 * 1024 * 1024  # 2GB
clear_gpu_memory()
wait_until_enough_gpu_memory(min_memory_available)