Databricks用QLoRA微调LLM遇PyTorch报错：'torch.dtype'无'itemsize'属性

问题背景 #

我在做自定义数据集上的LLM QLoRA微调，为解决之前的问题，在TrainingArguments对象中添加了gradient_checkpointing_kwargs={'use_reentrant':False}，但调用peft_trainer.train()时出现错误。

相关代码 #

peft_training_args = TrainingArguments(
    output_dir = output_dir,
    warmup_steps=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    max_steps=100, #1000
    learning_rate=2e-4,
    optim="paged_adamw_8bit",
    logging_steps=25,
    logging_dir="./logs",
    save_strategy="steps",
    save_steps=25,
    evaluation_strategy="steps",
    eval_steps=25,
    do_eval=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={'use_reentrant':False},
    report_to="none",
    overwrite_output_dir = 'True',
    group_by_length=True,
)

peft_model.config.use_cache = False

peft_trainer = transformers.Trainer(
    model=peft_model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    args=peft_training_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

错误信息 #

AttributeError: 'torch.dtype' object has no attribute 'itemsize'

环境信息 #

• 运行环境：Databricks
• PyTorch版本：2.0.1+cu118

解决方法 #

1. 升级依赖库到兼容版本
   PyTorch 2.0.x需要配合较新版本的transformers和peft才能支持use_reentrant=False的梯度检查点配置：

   • 确保transformers版本 >= 4.28.0
   • 确保peft版本 >= 0.4.0
   • 确保bitsandbytes版本 >= 0.39.0（适配paged_adamw_8bit优化器）

2. 调整梯度检查点的配置方式
   将gradient_checkpointing_kwargs的配置从TrainingArguments移到模型层面，避免参数传递冲突：

   # 移除TrainingArguments中的gradient_checkpointing_kwargs配置
   peft_training_args = TrainingArguments(
       # ... 其他参数保持不变
       gradient_checkpointing=True,
       # 去掉gradient_checkpointing_kwargs这一行
   )
   
   # 在模型上直接启用梯度检查点并传入参数
   peft_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})
   peft_model.config.use_cache = False
   

3. 验证8bit优化器的兼容性
   如果仍有问题，暂时替换paged_adamw_8bit为普通的adamw_torch，排除8bit优化器的兼容性问题：

   peft_training_args = TrainingArguments(
       # ... 其他参数
       optim="adamw_torch",
   )
   

内容的提问来源于stack exchange，提问作者Lidor Eliyahu Shelef