pythongoogle-colaboratoryhuggingface-transformersmistral-7bpeft

TypeError in SFTTrainer Initialization: Unexpected Keyword Argument 'tokenizer'


Question: I am trying to fine-tune the Mistral-7B-Instruct-v0.1-GPTQ model using SFTTrainer from trl. However, when running my script in Google Colab, I encounter the following error:

TypeError: SFTTrainer.__init__() got an unexpected keyword argument 'tokenizer'

Here is the relevant portion of my code:

import torch
from datasets import load_dataset, Dataset
from peft import LoraConfig, AutoPeftModelForCausalLM, prepare_model_for_kbit_training, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig, TrainingArguments
from trl import SFTTrainer, SFTConfig
import os

# Load and preprocess dataset
data = load_dataset("tatsu-lab/alpaca", split="train")
data_df = data.to_pandas()
data_df = data_df[:5000]
data_df["text"] = data_df[["input", "instruction", "output"]].apply(
    lambda x: "###Human: " + x["instruction"] + " " + x["input"] + " ###Assistant: " + x["output"], axis=1
)
data = Dataset.from_pandas(data_df)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.1-GPTQ")
tokenizer.pad_token = tokenizer.eos_token

# Load and prepare model
quantization_config_loading = GPTQConfig(bits=4, disable_exllama=True, tokenizer=tokenizer)
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ",
    device_map="auto"
)
model.config.use_cache = False
model.config.pretraining_tp = 1
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

# LoRA configuration
peft_config = LoraConfig(
    r=16, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", target_modules=["q_proj", "v_proj"]
)
model = get_peft_model(model, peft_config)

# Training configuration
training_arguments = SFTConfig(
    output_dir="mistral-finetuned-alpaca",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=1,
    optim="paged_adamw_32bit",
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    save_strategy="epoch",
    logging_steps=100,
    num_train_epochs=1,
    max_steps=250,
    fp16=True,
    packing=False,
    max_seq_length=512,
    dataset_text_field="text",
    push_to_hub=True
)

# Initialize trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=data,
    peft_config=peft_config,
    args=training_arguments,
    tokenizer=tokenizer,  # This causes the error
)

trainer.train()

What I Have Tried:

Question: Is SFTTrainer expecting the tokenizer to be handled differently in the latest versions of trl? How should I correctly pass the tokenizer to ensure training works?

I would appreciate any insights!


Solution

  • In the 0.12.0 release it is explained that the tokenzier argument is now called the processing_class parameter.

    You should be able to run your code as before by replacing tokenizer with processing_class:

    trainer = SFTTrainer(
        model=model,
        train_dataset=data,
        peft_config=peft_config,
        args=training_arguments,
        processing_class=tokenizer,  # This causes the error
    )