使用GPU运行Haystack PromptNode时出现CUDA内存不足错误求助

以下是针对你遇到问题的可行解决建议：

• 模型量化压缩：通过4位/8位量化加载模型，大幅降低显存占用。需要先安装bitsandbytes库，之后修改代码：

prompt_node = PromptNode(model_name_or_path='google/flan-t5-xl',
                         default_prompt_template=lfqa_prompt,
                         use_gpu=True,
                         max_length=300,
                         model_kwargs={"load_in_4bit": True})

也可以将load_in_4bit替换为load_in_8bit，根据显存情况选择。

• 自动设备映射：让PyTorch自动分配模型层到GPU和CPU，平衡显存负载：

prompt_node = PromptNode(model_name_or_path='google/flan-t5-xl',
                         default_prompt_template=lfqa_prompt,
                         use_gpu=True,
                         max_length=300,
                         model_kwargs={"device_map": "auto"})

• 优化显存分配策略：按照报错提示设置环境变量，减少显存碎片：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

可根据实际显存情况调整max_split_size_mb的数值。

• 降低输入批次大小：如果是批量处理数据，调小每次输入的样本数量，避免一次性占用过多显存。

• 切换小尺寸模型：若以上方法仍无法解决，可替换为google/flan-t5-large或google/flan-t5-base，这类模型显存需求更低，能在14.85GiB的GPU上正常运行。

内容的提问来源于stack exchange，提问作者sherin_a27