I have been trying to fine tune a QLora version of Llama3-8B-IT model on kaggle notebook on a custom dataset of about 44 questions. However I am not getting good results in all of the responses. The training script is provided below.

I have tried to ensure that the chat template and formatting is okay. I tried to include eos_token in the training arguments but it keeps on producing output until it reaches max_length values.

Probably there is some tweaking that needs to be done with LORA rank, learning rate or other hyperparameters that I am not able to assign probably. I typically train for 5-7 epochs and terminate when loss is around 0.9 and then testing the output by using the same prompts and even then the output is not satisfactory.

Any LLM expert over here can take a look and help me out? Thanks!

!pip install -U bitsandbytes
!pip install -U transformers
!pip install -U accelerate
!pip install -U peft
!pip install -U trl
!pip install -U datasets

from kaggle_secrets import UserSecretsClient

HF_READ = UserSecretsClient().get_secret("HF_TOKEN")
HF_WRITE = UserSecretsClient().get_secret("HF_TOKEN_WRITE")
WANDB_API = UserSecretsClient().get_secret("WANDB_API")

!pip install huggingface_hub

import subprocess
import pandas as pd

subprocess.run(["huggingface-cli", "login", "--token", HF_READ])

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import (
    LoraConfig,
    PeftModel,
    prepare_model_for_kbit_training,
    get_peft_model,
)
import os, torch, wandb
from datasets import load_dataset
from datasets import Dataset
from trl import DataCollatorForCompletionOnlyLM
from trl import SFTTrainer

base_model = "meta-llama/Meta-Llama-3-8B-Instruct"
dataset_name = "APaul1/Inmation_QA"
new_model = "meta-llama/Meta-Llama-3-8B-Instruct-Inmation_QA"

dataset = load_dataset(dataset_name, split="train")
dataset["text"][10]

tokenizer = AutoTokenizer.from_pretrained(base_model)

# tokenizer.eos_token_id
# tokenizer.pad_token_id

tokenizer = AutoTokenizer.from_pretrained(base_model)
tokenizer.padding_side = 'right'
tokenizer.pad_token = tokenizer.eos_token
tokenizer.add_eos_token = True
tokenizer.bos_token, tokenizer.eos_token

tokenizer.convert_tokens_to_ids('<|end_of_text|>')

bnbConfig = BitsAndBytesConfig(
    load_in_4bit = True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

bnbConfig_8bit = BitsAndBytesConfig(load_in_8bit=True,
                                         llm_int8_threshold=200.0)

model = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=bnbConfig_8bit,
        device_map="auto"
)

model.config.use_cache = False # silence the warnings. Please re-enable for inference!
model.config.pretraining_tp = 1
model.gradient_checkpointing_enable()

def formatting_chat_template_func(example, tokenizer=tokenizer):
    messages = example["text"]
    example["text"] = tokenizer.apply_chat_template(messages, tokenize=False) + '<|end_of_text|>'
    
    return example

dataset = dataset.map(formatting_chat_template_func,
                        fn_kwargs={"tokenizer": tokenizer})

dataset['text'][30]

model = prepare_model_for_kbit_training(model)
peft_config = LoraConfig(
    lora_alpha=8,
    lora_dropout=0.1,
    r=16,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=['o_proj', 'q_proj', 'up_proj', 'v_proj', 'k_proj', 'down_proj', 'gate_proj']
)
model = get_peft_model(model, peft_config)

instruction_template = "<|start_header_id|>user<|end_header_id|>"
response_template = "<|start_header_id|>assistant<|end_header_id|>"
collator = DataCollatorForCompletionOnlyLM(instruction_template=instruction_template, response_template=response_template, tokenizer=tokenizer, mlm=False)

wandb.login(key = WANDB_API)
run = wandb.init(
    project='Fine tuning Llama3', 
    job_type="training", 
    anonymous="allow"
)

training_arguments = TrainingArguments(
    output_dir="./Llama3-8B-IT-Inmation_QA", #change here for changing where the files get uploaded in huggingface
    num_train_epochs=5,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=1,
    optim="paged_adamw_32bit",
    save_strategy="epoch",
    logging_steps=10,
    logging_strategy="steps",
    learning_rate=2e-4,
    fp16=False,
    bf16=False,
    group_by_length=True,
    report_to="wandb"
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    max_seq_length= 512,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_arguments,
    packing= False,
    data_collator=collator
)

trainer.train()

### Test a sample output

index = 20
dataset['text'][index]

prompt = dataset['text'][index].split('<|start_header_id|>assistant<|end_header_id|>')[0] + '<|start_header_id|>assistant<|end_header_id|>n'
prompt

inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True).to("cuda")

terminators = [
    tokenizer.eos_token_id,
    tokenizer.convert_tokens_to_ids("<|eot_id|>"),
    tokenizer.convert_tokens_to_ids("<|end_of_text|>")
]

outputs = model.generate(**inputs, eos_token_id=terminators, max_length=256, num_return_sequences=1, do_sample=True, top_k=10)

text = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(text)