PyTorch加载LLaMA-2微调模型时CUDA与CPU张量设备不匹配问题

我尝试用transformers库加载微调后的LLaMA-2模型做文本生成，指定设备为cuda:3，也把停止标记张量移到了这个设备，但调用文本生成管道时触发RuntimeError，提示同时存在cuda:3和CPU的张量，想知道漏了哪些没移到CUDA的张量。

原代码 #

device = 'cuda:3'

tokenizer = transformers.AutoTokenizer.from_pretrained('llama-2-7b-chat-fine-tuned.bin')

stop_list = ['\nHuman:', '\n```\n']
stop_token_ids = [tokenizer(x)['input_ids'] for x in stop_list]
stop_token_ids

stop_token_ids = [torch.LongTensor(x).to(device) for x in stop_token_ids]
stop_token_ids

class StopOnTokens(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        for stop_ids in stop_token_ids:
            if torch.eq(input_ids[0][-len(stop_ids.to(device)):], stop_ids.to(device)).all():
                return True
        return False

stopping_criteria = StoppingCriteriaList([StopOnTokens()])

generate_text = transformers.pipeline(model='llama-2-7b-chat-fine-tuned.bin',
                                      tokenizer=tokenizer,
                                      return_full_text=True,
                                      task='text-generation',
                                      stopping_criteria=stopping_criteria,
                                      temperature=0.1,
                                      max_new_tokens=512,
                                      repetition_penalty=1.1)


res = generate_text("How are you doing?")
print(res[0]["generated_text"])

错误信息 #

Traceback (most recent call last):
  File "test_fine_tuned_llamas_2.py", line 75, in <module>
    res = generate_text("How are you doing?")
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/pipelines/text_generation.py", line 200, in __call__
    return super().__call__(text_inputs, **kwargs)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/pipelines/base.py", line 1122, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/pipelines/base.py", line 1129, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/pipelines/base.py", line 1028, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/pipelines/text_generation.py", line 261, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/generation/utils.py", line 1538, in generate
    return self.greedy_search(
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/generation/utils.py", line 2423, in greedy_search
    if stopping_criteria(input_ids, scores):
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/generation/stopping_criteria.py", line 127, in __call__
    return any(criteria(input_ids, scores) for criteria in self)
  File "/sharedvolume/felipe/project1/test_project1/lib/python3.8/site-packages/transformers/generation/stopping_criteria.py", line 127, in <genexpr>
    return any(criteria(input_ids, scores) for criteria in self)
  File "test_fine_tuned_llamas_2.py", line 59, in __call__
    if torch.eq(input_ids[0][-len(stop_ids.to(device)):], stop_ids.to(device)).all():
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cpu!

问题原因及修复方案 #

核心问题 #

1. 模型未指定设备：创建transformers pipeline时没有指定device，默认会把模型加载到cuda:0（或CPU），导致模型、输入张量的设备与你提前移到cuda:3的stop_token_ids不一致。
2. 冗余的设备迁移操作：StopOnTokens类的__call__方法中重复调用stop_ids.to(device)，虽然stop_ids已在cuda:3，但更关键的是input_ids在模型所在设备，两者对比时设备不匹配。

修改后的代码 #

device = 'cuda:3'
device_idx = 3  # pipeline的device参数接受整数索引

tokenizer = transformers.AutoTokenizer.from_pretrained('llama-2-7b-chat-fine-tuned.bin')

stop_list = ['\nHuman:', '\n```\n']
stop_token_ids = [tokenizer(x)['input_ids'] for x in stop_list]
stop_token_ids = [torch.LongTensor(x).to(device) for x in stop_token_ids]

class StopOnTokens(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        for stop_ids in stop_token_ids:
            # 先检查序列长度，避免索引越界
            if len(input_ids[0]) >= len(stop_ids):
                # 直接对比，无需重复迁移设备
                if torch.eq(input_ids[0][-len(stop_ids):], stop_ids).all():
                    return True
        return False

stopping_criteria = StoppingCriteriaList([StopOnTokens()])

generate_text = transformers.pipeline(
    model='llama-2-7b-chat-fine-tuned.bin',
    tokenizer=tokenizer,
    return_full_text=True,
    task='text-generation',
    stopping_criteria=stopping_criteria,
    temperature=0.1,
    max_new_tokens=512,
    repetition_penalty=1.1,
    device=device_idx  # 关键：指定模型加载到cuda:3
)

res = generate_text("How are you doing?")
print(res[0]["generated_text"])

关键修改点 #

• 给pipeline添加device=device_idx参数，确保模型、输入张量都加载到cuda:3，与stop_token_ids设备统一。
• 移除StopOnTokens中重复的stop_ids.to(device)调用，因为stop_token_ids已提前迁移到目标设备。
• 添加序列长度检查，避免生成的文本长度短于停止标记时出现索引越界错误。

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