使用Hugging Face Trainer做LoRA微调时遇RuntimeError问题咨询

在对小型LLM进行LoRA微调时，遇到错误：RuntimeError: cannot pin 'torch.cuda.FloatTensor' only dense CPU tensors can be pinned。有Discord用户指出问题源于手动将输入移至GPU（使用to(model.device)），称Hugging Face Trainer期望数据在CPU上，会内部处理GPU传输。但我在Trainer官方文档中未找到相关说明，请问该说法是否正确？若错误，该如何解决此问题？

可复现代码（Colab环境） #

先执行安装命令：

%pip install --upgrade torch transformers peft

再运行以下代码：

import torch
from torch.utils.data import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import TrainingArguments
from transformers import Trainer
from peft import LoraConfig, get_peft_model

model_name = "croissantllm/CroissantLLMBase"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

texts = [
    "The first sentence for fine-tuning. </s>",
    "The second sentence for fine-tuning. </s>"
]

inputs = [tokenizer(text, return_tensors="pt").to(model.device) for text in texts]

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"],
)

model = get_peft_model(model, lora_config)

class CustomDataset(Dataset):
    def __init__(self, input_list):
        self.input_list = input_list

    def __len__(self):
        return len(self.input_list)

    def __getitem__(self, idx):
        input_ids = self.input_list[idx]['input_ids'].squeeze()
        labels = input_ids.clone()
        return {"input_ids": input_ids, "labels": labels}

train_dataset = CustomDataset(inputs)

training_args = TrainingArguments(
    output_dir="./lora_croissantllm",
    per_device_train_batch_size=1,
    num_train_epochs=1,
    save_steps=10,
    save_total_limit=2,
    logging_dir="./logs",
    logging_steps=10,
)

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

trainer.train()

说法正确性 #

这个说法是正确的。Hugging Face Trainer的设计逻辑确实是期望数据集的张量都在CPU上，训练过程中会自动将批次数据移动到模型所在设备（GPU/CPU），同时默认开启pin_memory=True来加速数据传输——而内存固定操作仅支持CPU张量，这就是报错的直接原因：你提前把张量移到GPU，导致pin_memory操作失败。

解决方法 #

移除手动将输入移至GPU的代码即可，具体修改如下：
将原代码中的：

inputs = [tokenizer(text, return_tensors="pt").to(model.device) for text in texts]

改为：

inputs = [tokenizer(text, return_tensors="pt") for text in texts]

补充说明 #

• Trainer内部会通过DataLoader自动处理设备迁移，无需手动干预；
• 如果特殊场景下需要关闭pin_memory（不推荐，会降低数据传输效率），可以在初始化Trainer时添加dataloader_pin_memory=False参数，但更合理的做法还是让数据留在CPU。

内容的提问来源于stack exchange，提问作者FluidMechanics Potential Flows