多GPU微调Llama 3.1遇RuntimeError：张量设备不一致（cuda:3与cuda:0）

核心原因分析 #

这个错误的本质是部分张量被手动固定到了特定GPU（比如cuda:0），但Trainer在多GPU模式下会自动将模型/数据分配到不同设备，导致张量跨设备存在。Transformers的Trainer确实支持自动多GPU训练，但前提是代码中没有手动干预设备分配的操作。

具体排查与修复步骤 #

1. 检查数据集预处理代码 #

确保预处理时没有手动将张量移到特定GPU，预处理后的数据集应保持在CPU上，Trainer会自动在训练时将数据分发到对应GPU。

# 错误示例：手动指定GPU
def preprocess_function(examples):
    inputs = tokenizer(examples["text"], truncation=True, padding="max_length")
    inputs["labels"] = inputs["input_ids"].copy().to('cuda:0')  # 禁止手动指定设备
    return inputs

# 正确示例：保持张量在CPU
def preprocess_function(examples):
    inputs = tokenizer(examples["text"], truncation=True, padding="max_length")
    inputs["labels"] = inputs["input_ids"].copy()
    return inputs

2. 检查模型初始化代码 #

不要手动将模型移到特定GPU，让Trainer自动处理设备分配。使用from_pretrained时，避免指定单GPU的device_map。

# 错误示例：强制模型在单GPU
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    device_map="cuda:0"  # 这会和多GPU训练冲突
)

# 正确示例：自动分配多GPU
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    device_map="auto"  # 让框架自动分配到多GPU，或直接不指定device_map
)

3. 检查Trainer配置参数 #

确保Trainer配置中没有手动指定设备，同时使用正确的分布式启动方式。Trainer默认会启用DDP处理多GPU，无需手动初始化，但必须用支持分布式的命令启动脚本。

# 正确的Trainer配置示例
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=4,
    num_train_epochs=3,
    logging_dir="./logs",
    fp16=True,  # 混合精度可正常启用，无需手动指定设备
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

4. 启动命令必须正确 #

多GPU训练不能直接用python启动脚本，需用分布式启动命令：

# 用torchrun指定GPU数量（比如4卡）
torchrun --nproc_per_node=4 your_finetune_script.py
# 或用accelerate启动
accelerate launch your_finetune_script.py

5. 排查自定义逻辑中的张量操作 #

如果有自定义回调、损失函数或数据增强代码，确保所有张量操作都没有手动指定GPU。比如自定义损失函数时，直接使用Trainer传递的张量即可，不要手动执行.to('cuda:x')操作。

验证步骤 #

训练前可以打印设备信息确认：

# 查看模型参数所在设备（多GPU下会显示分布式设备）
print("Model parameter device:", next(model.parameters()).device)
# 查看数据集样本设备（应显示CPU）
print("Dataset sample device:", train_dataset[0]["input_ids"].device)

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