在SageMaker部署LLM文本生成端点时遇设备不匹配错误求助

核心问题分析 #

你遇到的错误源于输入张量未真正转移到CUDA设备，以及模型加载与设备转移的细节遗漏。以下是针对性修复方案：

1. 修复输入张量的设备转移 #

input_ids.to(device)不会原地修改张量，必须将返回值重新赋值给原变量，否则原张量仍停留在CPU上：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
nlp_model = LlamaForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float16)
# 方案1：生成张量时直接转移设备
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
nlp_model.to(device)

# 方案2：分步赋值确保转移生效
input_ids = tokenizer.encode(prompt, return_tensors="pt")
input_ids = input_ids.to(device)

2. 模型加载时直接指定设备（更可靠） #

使用device_map参数在加载阶段直接将模型分配到CUDA设备，避免后续to(device)操作可能的遗漏：

from transformers import LlamaForCausalLM, AutoTokenizer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
nlp_model = LlamaForCausalLM.from_pretrained(
    model_dir,
    torch_dtype=torch.float16,
    device_map="auto" if torch.cuda.is_available() else None
)
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)

3. 确保所有输入张量设备一致 #

如果使用tokenizer.__call__生成多维度输入（如attention_mask），需将所有张量统一转移到目标设备：

inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = nlp_model.generate(**inputs)

4. SageMaker环境验证 #

• 确认端点使用GPU实例类型（如ml.g5.xlarge及以上），CPU实例会导致torch.cuda.is_available()返回False。
• 优先使用Hugging Face官方SageMaker容器，确保镜像包含匹配的CUDA版本与PyTorch环境。

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