Как исправить ошибку GPU/CPU в LangChain и Transformers с помощью HuggingFacePipeline

Пошаговое решение ошибки несоответствия устройств (GPU/CPU) при генерации в LangChain+Transformers: фикс device в HuggingFacePipeline, пример кода, ReAct агент.

При сборке агента LangChain в стиле ReAct поверх локально развернутой модели HuggingFace через transformers легко столкнуться с несоответствием устройств: сама модель находится на GPU, а входы остаются на CPU. В итоге генерация падает с жесткой ошибкой, даже если device_map установлен в auto и модель без возражений загружается на cuda:0.

Воспроизводимый пример и контекст

В настройке используется transformers.pipeline для генерации текста, обернутый в HuggingFacePipeline из LangChain, плюс простой инструмент DuckDuckGoSearchRun. Модель инициализируется на GPU с device_map="auto", но AgentExecutor падает в момент вызова generate.

RuntimeError: Ожидалось, что все тензоры будут на одном устройстве, но индекс находится на cpu, в отличие от остальных тензоров на cuda:0

Вы вызываете .generate() с input_ids на устройстве, отличном от устройства вашей модели. input_ids на cpu, тогда как модель — на cuda.

Минимальный скрипт

Ниже — воспроизводимый пример и рабочая правка. Логика не меняется; отличаются только имена локальных переменных.

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
# Путь к локальной директории с чекпоинтом модели
ckpt_dir = "../gpt-oss-20b-local"
try:
    # Сборка пайплайна генерации на transformers
    gen_pipe = pipeline(
        "text-generation",
        model=ckpt_dir,
        dtype="auto",
        device_map="auto",
        max_new_tokens=256,
    )
    # Обертка для LangChain и явная фиксация устройства
    core_llm = HuggingFacePipeline(
        pipeline=gen_pipe,
        model_kwargs={"temperature": 0.5, "device": 0},
    )
    print("Local model is ready.")
except Exception as exc:
    print(f"Load failure: {exc}")
    exit()
# Одна утилита для веб‑поиска
web_search = DuckDuckGoSearchRun()
skillset = [web_search]
# Промпт в стиле ReAct
prompt_str = """
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}
"""
schema_prompt = PromptTemplate.from_template(prompt_str)
# Сборка агента и исполнителя
think_agent = create_react_agent(core_llm, skillset, schema_prompt)
runner = AgentExecutor(agent=think_agent, tools=skillset, verbose=True)
print("Agent initialized.")
print("=" * 50)
# Запуск запроса
user_q = "Who is the current prime minister of the United Kingdom and what is their political party?"
outcome = runner.invoke({"input": user_q})
print("-" * 50)
print(f"Final Response: {outcome['output']}")

Что на самом деле происходит

Трассировка стека прямо указывает на проблему: веса модели на cuda:0, а входные индексы, которые используются при поиске в эмбеддингах, приходят с CPU. Во время генерации transformers выводит предупреждение, что input_ids находятся на cpu, тогда как модель — на cuda, и в итоге прямой проход падает с ошибкой “Expected all tensors to be on the same device.” Проще говоря, обертка не переносит входы на то же устройство, где находится модель.

Исправление

Зафиксируйте устройство в обертке HuggingFacePipeline. Передача параметра device в model_kwargs гарантирует, что входы будут отправляться на тот GPU, где размещена модель. Установка device в 0 устраняет несоответствие:

core_llm = HuggingFacePipeline(
    pipeline=gen_pipe,
    model_kwargs={"temperature": 0.5, "device": 0},
)

Это выравнивает input_ids с моделью на cuda:0 и разблокирует генерацию внутри AgentExecutor.

Почему это важно

Смешанное исполнение на разных устройствах приводит к жестким сбоям в базовых операторах вроде embedding и index_select, а также может ухудшать производительность, если незаметно происходит откат на CPU. Чтобы обеспечить стабильный инференс при связке пайплайнов transformers с AgentExecutor из LangChain, важно держать модель и ее входы на одном устройстве.

В журналах также есть уведомление о деприкации: HuggingFacePipeline из LangChain будет удален в одном из будущих релизов; обновленный класс доступен в пакете langchain-huggingface с путем импорта из langchain_huggingface. Стоит учесть этот маршрут миграции, чтобы сэкономить время при обновлениях.

Выводы

Если агент LangChain выдает ошибку несоответствия устройств во время генерации, даже при загрузке модели с device_map="auto", зафиксируйте устройство в обертке, передав device: 0. Это гарантирует, что входные тензоры стабильно попадут на cuda:0 и устранит ошибку во время выполнения. Параллельно обратите внимание на уведомление о деприкации HuggingFacePipeline и при необходимости планируйте переход на импорт из langchain_huggingface.