如何在 PyTorch 中释放 GPU memory

Question

我有一个句子列表，我正在尝试使用以下代码使用多个模型计算困惑度：

from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch
import numpy as np
model_name = 'cointegrated/rubert-tiny'
model = AutoModelForMaskedLM.from_pretrained(model_name).cuda()
tokenizer = AutoTokenizer.from_pretrained(model_name)

def score(model, tokenizer, sentence):
    tensor_input = tokenizer.encode(sentence, return_tensors='pt')
    repeat_input = tensor_input.repeat(tensor_input.size(-1)-2, 1)
    mask = torch.ones(tensor_input.size(-1) - 1).diag(1)[:-2]
    masked_input = repeat_input.masked_fill(mask == 1, tokenizer.mask_token_id)
    labels = repeat_input.masked_fill( masked_input != tokenizer.mask_token_id, -100)
    with torch.inference_mode():
        loss = model(masked_input.cuda(), labels=labels.cuda()).loss
    return np.exp(loss.item())


print(score(sentence='London is the capital of Great Britain.', model=model, tokenizer=tokenizer)) 
# 4.541251105675365

大多数模型运行良好，但有些句子似乎会引发错误：

RuntimeError: CUDA out of memory. Tried to allocate 10.34 GiB (GPU 0; 23.69 GiB total capacity; 10.97 GiB already allocated; 6.94 GiB free; 14.69 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

这是有道理的，因为有些很长。 所以我所做的是添加类似try, except RuntimeError, pass的东西。

这似乎一直有效，直到大约 210 个句子，然后它才输出错误：

CUDA error: an illegal memory access was encountered CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

我发现这个有很多讨论和想法，有些是关于潜在的故障 GPU？ 但我知道我的 GPU 可以正常工作，因为这个确切的代码适用于其他型号。 这里也讨论了批量大小，这就是为什么我认为它可能与释放 memory 有关。

我尝试运行torch.cuda.empty_cache()来释放 memory，就像这里每隔一段时间后一样，但它没有用（抛出同样的错误）。

更新：我过滤了长度超过 550 的句子，这似乎消除了CUDA error: an illegal memory access was encountered CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1. CUDA error: an illegal memory access was encountered CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1. 错误。

Answer 1

您需要在gc.collect()之前应用torch.cuda.empty_cache()我还将 model 拉到 cpu，然后删除该 model 及其检查点。 试试适合你的方法：

import gc

model.cpu()
del model, checkpoint
gc.collect()
torch.cuda.empty_cache()

Answer 2

实际上，PyTorch 会重新分配 memory 并在需要时释放缓存。 否则，您需要删除对张量的所有引用，这会降低您的代码速度。 torch.cuda.empty_cache()不会阻止缓存，但它将允许释放可重用的 memory。 另一方面，您可以重新启动整个 kernel。

    model.to(torch.device("cuda"))
    del model
    gc.collect()
    torch.cuda.empty_cache()

Answer 3

我没有确切的答案，但我可以分享一些我在类似情况下采用的故障排除技术......希望它可能会有所帮助。

1. 首先，不幸的是， CUDA error有时含糊不清，因此您应该考虑在 CPU 上运行代码以查看是否真的发生了其他事情（请参阅此处）

2. 如果问题是关于 memory，这里有两个我使用的自定义工具：

from torch import cuda


def get_less_used_gpu(gpus=None, debug=False):
    """Inspect cached/reserved and allocated memory on specified gpus and return the id of the less used device"""
    if gpus is None:
        warn = 'Falling back to default: all gpus'
        gpus = range(cuda.device_count())
    elif isinstance(gpus, str):
        gpus = [int(el) for el in gpus.split(',')]

    # check gpus arg VS available gpus
    sys_gpus = list(range(cuda.device_count()))
    if len(gpus) > len(sys_gpus):
        gpus = sys_gpus
        warn = f'WARNING: Specified {len(gpus)} gpus, but only {cuda.device_count()} available. Falling back to default: all gpus.\nIDs:\t{list(gpus)}'
    elif set(gpus).difference(sys_gpus):
        # take correctly specified and add as much bad specifications as unused system gpus
        available_gpus = set(gpus).intersection(sys_gpus)
        unavailable_gpus = set(gpus).difference(sys_gpus)
        unused_gpus = set(sys_gpus).difference(gpus)
        gpus = list(available_gpus) + list(unused_gpus)[:len(unavailable_gpus)]
        warn = f'GPU ids {unavailable_gpus} not available. Falling back to {len(gpus)} device(s).\nIDs:\t{list(gpus)}'

    cur_allocated_mem = {}
    cur_cached_mem = {}
    max_allocated_mem = {}
    max_cached_mem = {}
    for i in gpus:
        cur_allocated_mem[i] = cuda.memory_allocated(i)
        cur_cached_mem[i] = cuda.memory_reserved(i)
        max_allocated_mem[i] = cuda.max_memory_allocated(i)
        max_cached_mem[i] = cuda.max_memory_reserved(i)
    min_allocated = min(cur_allocated_mem, key=cur_allocated_mem.get)
    if debug:
        print(warn)
        print('Current allocated memory:', {f'cuda:{k}': v for k, v in cur_allocated_mem.items()})
        print('Current reserved memory:', {f'cuda:{k}': v for k, v in cur_cached_mem.items()})
        print('Maximum allocated memory:', {f'cuda:{k}': v for k, v in max_allocated_mem.items()})
        print('Maximum reserved memory:', {f'cuda:{k}': v for k, v in max_cached_mem.items()})
        print('Suggested GPU:', min_allocated)
    return min_allocated


def free_memory(to_delete: list, debug=False):
    import gc
    import inspect
    calling_namespace = inspect.currentframe().f_back
    if debug:
        print('Before:')
        get_less_used_gpu(debug=True)

    for _var in to_delete:
        calling_namespace.f_locals.pop(_var, None)
        gc.collect()
        cuda.empty_cache()
    if debug:
        print('After:')
        get_less_used_gpu(debug=True)

2.1 free_memory允许您结合gc.collect和cuda.empty_cache从命名空间中删除一些所需的对象并释放它们的 memory（您可以将变量名称列表作为to_delete参数传递）。 这很有用，因为您可能有未使用的对象占用 memory。 例如，假设您遍历 3 个模型，那么当您进行第二次迭代时，第一个模型可能仍然需要一些 gpu memory（我不知道为什么，但我在笔记本电脑中遇到过这种情况，也是我能找到的唯一解决方案是重新启动笔记本电脑或显式释放内存）。 但是，我不得不说这并不总是实用的，因为您需要知道哪些变量持有 GPU memory...而且情况并非总是如此，尤其是当您有很多与Z20F35E630DAF44DBFA4C3F68内部相关的渐变时信息）。 您还可以尝试的一件事是使用with torch.no_grad():而不是with torch.inference_mode(): ； 它们应该是等效的，但可能值得一试......

2.2 如果您有一个 multigpu 环境，您可以考虑交替切换到使用较少的 GPU，这要归功于其他get_less_used_gpu

3. 此外，您可以尝试跟踪 GPU 的使用情况，以查看错误何时发生并从那里进行调试。 如果您在 linux 平台上，我可以建议的最好/最简单的方法是使用nvtop

希望有些东西是有用的:)