LangChain调用本地HuggingFace模型时张量设备不匹配RuntimeError修复

问题概述 #

使用LangChain构建Agent，采用本地部署的gpt-oss-20b模型，通过transformers pipeline封装为LangChain的HuggingFacePipeline。模型已通过device_map="auto"加载至GPU，但调用AgentExecutor.invoke()时触发设备不匹配错误，核心报错：

RuntimeError: Expected all tensors to be on the same device, but got index is on cpu, different from other tensors on cuda:0

复现代码 #

import os
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.prompts import PromptTemplate
from langchain_community.llms import HuggingFacePipeline
from transformers import pipeline
import torch

# 1. Set up the Language Model
model_path = "../gpt-oss-20b-local" 

try:
    # Create a transformers pipeline for text generation
    local_pipe = pipeline(
        "text-generation",
        model=model_path,
        dtype="auto",
        device_map="auto", # Should handle placing the model on GPU
        max_new_tokens=256,
    )

    # Wrap the pipeline for LangChain
    llm = HuggingFacePipeline(
        pipeline=local_pipe,
        model_kwargs={"temperature": 0.5},
    )
    print("Local LLM Loaded successfully.")

except Exception as e:
    print(f"Error loading local model: {e}")
    exit()

# 2. Define Tools
search = DuckDuckGoSearchRun()
tools = [search]

# 3. Create Prompt
template = """
Answer the following questions as best you can. You have access to the following tools:
{tools}
Use the following format:
Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question
Begin!
Question: {input}
Thought:{agent_scratchpad}
"""
prompt = PromptTemplate.from_template(template)

# 4. Create Agent and Executor
agent = create_react_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

print("Agent Executor created. Ready to receive input.")
print("=" * 50)

# 6. Run the Agent
question = "Who is the current prime minister of the United Kingdom and what is their political party?"
response = agent_executor.invoke({"input": question}) # Error occurs here

print("-" * 50)
print(f"Final Response: {response['output']}")

完整报错日志 #

Fetching 40 files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:00<00:00, 3852.05it/s]
Fetching 40 files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:00<00:00, 2968.16it/s]
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  1.95it/s]
Device set to use cuda:0
/meysam/test-oss/agent.py:28: LangChainDeprecationWarning: The class `HuggingFacePipeline` was deprecated in LangChain 0.0.37 and will be removed in 1.0. An updated version of the class exists in the :class:`~langchain-huggingface package and should be used instead. To use it run `pip install -U :class:`~langchain-huggingface` and import as `from :class:`~langchain_huggingface import HuggingFacePipeline``.
  llm = HuggingFacePipeline(
Local LLM Loaded successfully.
Tools defined.
Agent created.
Agent Executor created. Ready to receive input.
==================================================


> Entering new AgentExecutor chain...
/meysam/envs/new_env/lib/python3.12/site-packages/transformers/generation/utils.py:2412: UserWarning: You are calling .generate() with the `input_ids` being on a device type different than your model's device. `input_ids` is on cpu, whereas the model is on cuda. You may experience unexpected behaviors or slower generation. Please make sure that you have put `input_ids` to the correct device by calling for example input_ids = input_ids.to('cuda') before running `.generate()`.
  warnings.warn(
Exception in thread Thread-3 (generate):
Traceback (most recent call last):
  File "/usr/lib/python3.12/threading.py", line 1075, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.12/threading.py", line 1012, in run
    self._target(*self._args, **self._kwargs)
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/transformers/generation/utils.py", line 2539, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/transformers/generation/utils.py", line 2867, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/transformers/utils/generic.py", line 940, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 663, in forward
    outputs: MoeModelOutputWithPast = self.model(
                                      ^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/transformers/utils/generic.py", line 1064, in wrapper
    outputs = func(self, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 474, in forward
    inputs_embeds = self.embed_tokens(input_ids)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/modules/sparse.py", line 192, in forward
    return F.embedding(
           ^^^^^^^^^^^^
  File "/meysam/envs/new_env/lib/python3.12/site-packages/torch/nn/functional.py", line 2546, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but got index is on cpu, different from other tensors on cuda:0 (when checking argument in method wrapper_CUDA__index_select)

修复方案 #

方案1：显式指定Pipeline的GPU设备 #

虽然设置了device_map="auto"，但显式指定device参数可以确保pipeline强制将输入张量移到GPU：

local_pipe = pipeline(
    "text-generation",
    model=model_path,
    dtype="auto",
    device_map="auto",
    device=0,  # 指定cuda:0设备，也可以用torch.cuda.current_device()
    max_new_tokens=256,
)

方案2：使用LangChain官方推荐的新版HuggingFacePipeline #

旧版langchain_community.llms.HuggingFacePipeline已废弃，新版langchain_huggingface.HuggingFacePipeline对设备处理更完善：

1. 先安装依赖：

pip install langchain-huggingface

2. 修改导入语句：

# 替换原来的from langchain_community.llms import HuggingFacePipeline
from langchain_huggingface import HuggingFacePipeline

3. 保持其他代码不变，重新运行。

方案3：自定义LLM类强制输入移至GPU #

如果以上方案无效，可自定义LLM类，手动控制输入张量的设备：

from langchain_core.llms import LLM
from typing import Optional, List, Any

class CustomLocalLLM(LLM):
    pipeline: pipeline

    @property
    def _llm_type(self) -> str:
        return "custom"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> str:
        # 将输入张量移到模型所在设备
        inputs = self.pipeline.tokenizer(prompt, return_tensors="pt").to(self.pipeline.model.device)
        # 生成文本
        outputs = self.pipeline.model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.5,
            stop=stop,
            **kwargs
        )
        # 解码结果
        return self.pipeline.tokenizer.decode(outputs[0], skip_special_tokens=True)

# 替换原来的llm初始化
llm = CustomLocalLLM(pipeline=local_pipe)

验证修复 #

重新运行脚本，若AgentExecutor能正常执行并返回结果，说明设备不匹配问题已解决。

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