How to Fine-tune an LLM Part 3: The HuggingFace Trainer

from trl.train import SFTTrainer, TrainingArguments

os.environ["WANDB_PROJECT"] = "alpaca_ft"  # name your W&B project
os.environ["WANDB_LOG_MODEL"] = "checkpoint"  # log all model checkpoints

training_args = TrainingArguments(
    report_to="wandb", # enables logging to W&B 😎
    per_device_train_batch_size=16,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    num_train_epochs=3,
    gradient_accumulation_steps=2, # simulate larger batch sizes
)

trainer = SFTTrainer(
    model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    packing=True, # pack samples together for efficient training
    max_seq_length=1024, # maximum packed length 
    args=training_args,
    formatting_func=formatting_func, # format samples with a model schema
)
trainer.train()

def prompt_no_input(row):
    return ("Below is an instruction that describes a task. "
            "Write a response that appropriately completes the request.\n\n"
            "### Instruction:\n{instruction}\n\n### Response:\n{output}").format_map(row)

def prompt_input(row):
    return ("Below is an instruction that describes a task, paired with an input that provides further context. "
            "Write a response that appropriately completes the request.\n\n"
            "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n{output}").format_map(row)

def create_alpaca_prompt(row):
    return prompt_no_input(row) if row["input"] == "" else prompt_input(row)

train_dataset = ConstantLengthDataset(
    tokenizer,
    train_data,
    formatting_func=create_alpaca_prompt,
    seq_length=1024,
)

sample = {"instruction": "### Instruction: List three fruits\n### Response:\n",
          "output": "- Apple\n- Orange\n- Strawberry"}

instruction = tokenizer(sample["instruction"])["input_ids"]
output = tokenizer(sample["output"], add_special_tokens=False)["input_ids"] + [tokenizer.eos_token_id]

def printl(mylist):
    return print(*mylist, sep=' ')

printl(instruction)
printl(instruction + output)
# 1 835 2799 4080 29901 2391 2211 285 21211 13 2277 29937 13291 29901 13. # tokenized instruction
# 1 835 2799 4080 29901 2391 2211 285 21211 13 2277 29937 13291 29901 13 29899 12113 13 29899 26048 13 29899 624 1610 16344 2

{"input_ids":instruction + output, "labels":instruction + output} # the HF model will take care of the shift + 1

{"input_ids":instruction + output, "labels":[-100]*len(instruction) + output}

trainer = SFTTrainer(
    model,  # the same model as before, were we froze 75% of the layers
    args=training_args, # the parameters of the training: batch_size, report_to="wandb", max_steps, etc...
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    packing=True, # this tells the trainer to pack sequences of `max_seq_lenght`
    max_seq_length=1024, # the desired input sequences, we could increase this even more
    formatting_func=create_alpaca_prompt, # The instruction template to apply to the examples
)

trainer.train()

batch_size = 16
gradient_accumulation_steps = 2
num_train_epochs = 3

os.environ["WANDB_LOG_MODEL"] = "end" # the model will be uploaded to W&B at the end of training

training_args = TrainingArguments(
    output_dir="./output/",
    report_to="wandb",  # this tells the Trainer to log the metrics to W&B
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size//2,
    bf16=True,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio = 0.1,
    gradient_accumulation_steps=gradient_accumulation_steps,
    gradient_checkpointing=True,
    evaluation_strategy="epoch",
    num_train_epochs=num_train_epochs,
    # logging strategies 
    logging_strategy="steps",
    logging_steps=1,
    save_strategy="epoch", # saving is done at the end of each epoch
)

  class LLMSampleCB(WandbCallback):
    def __init__(self, trainer, test_dataset, num_samples=10, max_new_tokens=256, log_model="checkpoint"):
	"A CallBack to log samples a wandb.Table during training"
        super().__init__()
        self._log_model = log_model
        self.sample_dataset = test_dataset.select(range(num_samples))
        self.model, self.tokenizer = trainer.model, trainer.tokenizer
        self.gen_config = GenerationConfig.from_pretrained(trainer.model.name_or_path,
                                                           max_new_tokens=max_new_tokens)
    def generate(self, prompt):
        tokenized_prompt = self.tokenizer(prompt, return_tensors='pt')['input_ids'].cuda()
        with torch.inference_mode():
            output = self.model.generate(tokenized_prompt, generation_config=self.gen_config)
        return self.tokenizer.decode(output[0][len(tokenized_prompt[0]):], skip_special_tokens=True)
    
    def samples_table(self, examples):
	"Create a wandb.Table to store the generations"
        records_table = wandb.Table(columns=["prompt", "generation"] + list(self.gen_config.to_dict().keys()))
        for example in tqdm(examples, leave=False):
            prompt = example["text"]
            generation = self.generate(prompt=prompt)
            records_table.add_data(prompt, generation, *list(self.gen_config.to_dict().values()))
        return records_table
        
    def on_evaluate(self, args, state, control,  **kwargs):
	"Log the wandb.Table after calling trainer.evaluate"
        super().on_evaluate(args, state, control, **kwargs)
        records_table = self.samples_table(self.sample_dataset)
        self._wandb.log({"sample_predictions":records_table})

trainer = SFTTrainer(....)

wandb.init(project=..., config=config)

# we instantiate the W&B callback with the trainer object and the dataset we want to sample from
wandb_callback = LLMSampleCB(trainer, test_dataset, num_samples=10, max_new_tokens=256)
trainer.add_callback(wandb_callback)

# we then continue as regular
trainer.train()
 
# good practice to end your run
wandb.finish()

from peft import LoraConfig

peft_config = LoraConfig(
    r=64,  # the rank of the LoRA matrices
    lora_alpha=16, # the weight
    lora_dropout=0.1, # dropout to add to the LoRA layers
    bias="none", # add bias to the nn.Linear layers?
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj","v_proj","o_proj"], # the name of the layers to add LoRA
    modules_to_save=None, # layers to unfreeze and train from the original pre-trained model
)

trainer = SFTTrainer(
  ...
  peft_config = peft_config,
  ...
)

model = AutoModelForCausalLM.from_pretrained('meta-llama/Llama-2-7b-hf',  load_in_4bit=True, ...)