Langchain构建RAG应用多GPU部署遇HuggingFace Pipeline显存不足问题

你当前的问题出在硬指定了单GPU设备，没利用多GPU做模型并行，导致单卡负载过高。Langchain的HuggingFacePipeline.from_model_id确实不直接支持device_map参数，但可以通过手动创建Transformers原生Pipeline的方式绕开限制，同时充分利用多GPU资源。

以下是具体解决方法：

1. 手动创建支持模型并行的Pipeline，再传入Langchain #

直接用HuggingFace Transformers的pipeline函数创建支持device_map="auto"的文本生成管道，再用Langchain的HuggingFacePipeline包装，模型会自动分配到多个GPU上：

from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline, HuggingFaceEmbeddings
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"

# 手动加载模型和tokenizer，启用模型并行
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"  # 自动分配模型层到多个GPU
)

# 创建Transformers原生Pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 包装成Langchain的LLM
llm = ChatHuggingFace(llm=pipe)

2. 优化Embedding模型的GPU分配 #

你的Embedding当前指定了cuda:1，可以同样用device_map让Embedding模型合理分配，或者如果模型较小，指定到空闲GPU，同时关闭不必要的多进程（多进程可能额外占用内存）：

embedding = HuggingFaceEmbeddings(
    model_name=MODEL_NAME,
    model_kwargs={
        "torch_dtype": torch.float16,
        "device_map": "auto"  # 自动分配Embedding模型到GPU
    },
    multi_process=False  # 多进程可能增加内存开销，多GPU环境下建议关闭
)

3. 额外内存优化建议 #

• 启用量化加载：如果模型依然占内存过高，用BitsAndBytes做4bit/8bit量化，进一步降低显存占用：

  model = AutoModelForCausalLM.from_pretrained(
      MODEL_NAME,
      torch_dtype=torch.float16,
      device_map="auto",
      load_in_4bit=True,  # 4bit量化
      bnb_4bit_quant_type="nf4",
      bnb_4bit_compute_dtype=torch.float16
  )
  

• 调整生成参数：减少max_new_tokens数值，或设置合适的batch_size，避免一次性处理过多请求导致显存溢出。

内容的提问来源于stack exchange，提问作者Cihan Yalçın