使用Hugging Face微调TheBloke/Llama-2-13B-chat-GPTQ遇Exllama错误

你遇到的错误核心原因是：TheBloke的这款GPTQ模型默认启用Exllama后端，该后端要求所有模型模块必须直接加载到GPU上，但model.to('cuda')无法对Exllama加载的模块生效——Exllama的权重加载逻辑独立于PyTorch常规设备迁移方法。

以下是两种可行的解决方案，以及代码中需要修正的其他问题：

方案1：禁用Exllama后端 #

直接在加载模型时通过量化配置关闭Exllama，改用常规GPTQ推理/微调逻辑，这样model.to('cuda')或device_map='auto'就能正常工作。

修改模型加载代码：

from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

model_name = "TheBloke/Llama-2-13B-chat-GPTQ"

# 配置量化参数，禁用Exllama
quantization_config = GPTQConfig(
    disable_exllama=True
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map='auto'  # 自动分配模型到GPU
)

方案2：强制Exllama后端加载所有模块到GPU #

如果想保留Exllama的性能优势，需要在加载模型时就强制所有模块分配到GPU，而非事后迁移。修改模型加载代码：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "TheBloke/Llama-2-13B-chat-GPTQ"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='cuda'  # 强制所有模块直接加载到GPU
)

额外代码修正点 #

除了上述报错，你的代码还有几个潜在问题需要调整：

• 数据集加载后是DatasetDict，必须取对应的split（比如train_data['train']），否则Trainer会报错
• formatting_func需要遵循Llama-2的聊天格式，同时返回包含labels的字典（因果语言模型训练要求）
• 13B模型的batch_size设置过大（32/64），即使GPTQ量化也会显存溢出，建议调小到2-8之间

修改后的完整代码示例（方案1） #

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, GPTQConfig
from datasets import load_dataset
import torch

# 检查GPU状态
print("可用GPU数量:", torch.cuda.device_count())
print("第一块GPU名称:", torch.cuda.get_device_name(0))

model_name = "TheBloke/Llama-2-13B-chat-GPTQ"

# 加载tokenizer并设置pad token（Llama默认无pad token）
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 加载禁用Exllama的模型
quantization_config = GPTQConfig(disable_exllama=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map='auto'
)

# 加载并提取数据集split
train_data = load_dataset('json', data_files='train_data.jsonl')['train']
val_data = load_dataset('json', data_files='val_data.jsonl')['train']

# 格式化数据：遵循Llama-2官方聊天格式
def formatting_func(example):
    text = f"<s>[INST] {example['input']} [/INST] {example.get('output', '')} </s>"
    tokenized = tokenizer(
        text,
        truncation=True,
        max_length=512,  # 根据你的数据长度调整
        padding='max_length'
    )
    # 设置labels，与input_ids一致（因果LM训练要求）
    tokenized['labels'] = tokenized['input_ids'].copy()
    return tokenized

# 处理数据集并移除原始文本列
train_data = train_data.map(formatting_func, remove_columns=['input', 'output'])
val_data = val_data.map(formatting_func, remove_columns=['input', 'output'])

# 调整训练参数：降低batch_size并开启混合精度
training_args = TrainingArguments(
    output_dir="./output",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    save_steps=1000,
    save_total_limit=2,
    logging_steps=10,
    fp16=True,  # 开启混合精度节省显存
)

# 初始化Trainer并开始训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=val_data,
)

trainer.train()

# 保存模型和tokenizer
model.save_pretrained("./output")
tokenizer.save_pretrained("./output")

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