多GPU环境下PyTorch长上下文推理仍遇CUDA内存不足问题

我在运行PyTorch模型时遇到CUDA内存不足错误，系统配备8块NVIDIA GPU，模型已通过device_map='auto'分布式部署在所有GPU上，但输入是约20k tokens的长上下文，批量大小为1的推理任务中GPU 0仍报内存不足。

模型加载代码 #

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("MODEL_TYPE")
model = AutoModelForCausalLM.from_pretrained("MODEL_TYPE", output_attentions=True, device_map = 'auto', torch_dtype=torch.float16, low_cpu_mem_usage=True)

报错信息 #

OutOfMemoryError: CUDA out of memory. Tried to allocate 20.11 GiB. GPU 0 has a total capacty of 22.17 GiB of which 16.06 GiB is free. Including non-PyTorch memory, this process has 6.10 GiB memory in use. Of the allocated memory 5.57 GiB is allocated by PyTorch, and 308.62 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

推理代码 #

inputs = tokenizer(prompt, return_tensors="pt")
torch.cuda.empty_cache()
model.generation_config.temperature = temp
model.eval()
with torch.no_grad():
    output = model.generate(inputs.input_ids, max_length=25000, output_attentions=False,output_scores=False, return_dict_in_generate=True)
    print("temp:",model.generation_config.temperature)
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

response = tokenizer.batch_decode(output[0], skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]

尝试过将输入强制转移到其他GPU（inputs = tokenizer(prompt, return_tensors="pt").to("cuda:1")）但无效，请问如何高效利用可用GPU处理长上下文输入以避免内存不足错误？

解决方法 #

1. 启用梯度检查点 #

梯度检查点通过复用计算资源减少中间张量的内存占用，适合长上下文场景。加载模型时添加配置：

model = AutoModelForCausalLM.from_pretrained(
    "MODEL_TYPE",
    output_attentions=True,
    device_map='auto',
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    gradient_checkpointing=True
)
# 推理前启用兼容模式
model.gradient_checkpointing_enable()

2. 优化模型设备分配策略 #

device_map='auto'可能导致GPU 0负载集中，改用balanced模式让模型层均匀分布在多GPU上：

model = AutoModelForCausalLM.from_pretrained(
    "MODEL_TYPE",
    output_attentions=True,
    device_map='balanced',  # 替换auto为balanced
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

也可以自定义device_map字典，手动将大模型层分配到显存充足的GPU，避免单卡过载。

3. 启用FlashAttention优化注意力计算 #

长上下文下注意力机制是内存占用核心，FlashAttention能大幅降低内存消耗并提升速度。如果模型支持（如Llama、GPT系列），加载时启用：

model = AutoModelForCausalLM.from_pretrained(
    "MODEL_TYPE",
    output_attentions=True,
    device_map='auto',
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    use_flash_attention_2=True
)

需确保安装最新版transformers和flash-attn依赖。

4. 优化生成阶段的KV缓存 #

长上下文生成时，KV缓存会持续占用内存。改用max_new_tokens替代max_length，仅计算新增token的长度；同时启用动态缓存：

output = model.generate(
    inputs.input_ids,
    max_new_tokens=5000,  # 输入已20k，只需设置新增token数
    output_attentions=False,
    output_scores=False,
    return_dict_in_generate=True,
    cache_implementation="dynamic"
)

5. 调整内存分配参数减少碎片 #

设置PyTorch内存分配参数，缓解内存碎片问题。可在代码开头添加：

import torch
torch.backends.cuda.max_split_size_mb = 512

或在启动脚本前设置环境变量：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512

6. 关闭不必要的模型输出 #

如果推理不需要注意力输出，加载模型时直接关闭output_attentions=True，避免模型保留冗余层占用内存：

model = AutoModelForCausalLM.from_pretrained(
    "MODEL_TYPE",
    output_attentions=False,  # 关闭不必要的输出
    device_map='auto',
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

内容的提问来源于stack exchange，提问作者Flying-Meta