SFT trainer on colab T4

[savour-it-last]

I would like to know the extent to which we can use SFT trainer to train something that actually gives decent results on google colab's T4.
The code I used:

!pip install transformers accelerate datasets bitsandbytes einops trl huggingface_hub torch
import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

torch.cuda.empty_cache()

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' #Reduce usage

dataset = load_dataset("imdb", split="train[:10000]")
# dataset =  dataset[:1000]

train_size = int(0.8 * len(dataset))  # Use 80% for training
train_dataset = dataset.select(range(train_size))
eval_dataset = dataset.select(range(train_size, len(dataset)))



model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
model.gradient_checkpointing_enable() #reduces memory usage


SFTconfig = SFTConfig(
    output_dir="/content/checkpoints",  # The directory to save model checkpoints
    dataset_text_field="text",
    packing=True,  # Set to False for instruction tuning
    per_device_train_batch_size=5,
    fp16=True,  # Use GPU
    save_strategy="epoch",  # Save the model after every epoch
    evaluation_strategy="epoch",  # Evaluate at the end of each epoch
    # logging_strategy="steps",  # Log training loss every X steps
    # logging_steps=10, # Log every 10 steps
    save_total_limit=1,  # Keep only the last 3 checkpoints (v1, v2, v3)
    load_best_model_at_end= True,
    learning_rate=0.01,
    # gradient_accumulation_steps=1 # you can use if you cant feed amny train_batches
)
trainer = SFTTrainer(
    model,
    train_dataset=train_dataset,
    eval_dataset = eval_dataset,
    args=SFTconfig,
)

trainer.train()
trainer.save_model("/content/final_finetuned_model")

def compare_models(user_prompt: str):
    # Load the default pre-trained model and tokenizer
    default_model_name = "facebook/opt-350m"
    tokenizer = AutoTokenizer.from_pretrained(default_model_name)
    default_model = AutoModelForCausalLM.from_pretrained(default_model_name)

    # Load the fine-tuned model
    finetuned_model_name = "/content/final_finetuned_model"  # Path to the saved fine-tuned model
    finetuned_model = AutoModelForCausalLM.from_pretrained(finetuned_model_name)

    # Tokenize the input prompt
    inputs = tokenizer(user_prompt, return_tensors="pt")

    # Generate outputs from both models
    with torch.no_grad():
        default_output = default_model.generate(**inputs, max_length=200)
        finetuned_output = finetuned_model.generate(**inputs, max_length=200)

    # Decode the outputs
    default_response = tokenizer.decode(default_output[0], skip_special_tokens=True)
    finetuned_response = tokenizer.decode(finetuned_output[0], skip_special_tokens=True)

    return {"default_model": default_response, "finetuned_model": finetuned_response}

So when I increase some of these values in the config it is giving cuda out of memory errors. Now I was thinking of training a 1b or 2b model to do a particular extraction task. But if this itself is not giving decent results/wont run then there is no point in looking furthur. This is my last attempt with SFT and was thinking of exploring PEFT, LORA etc next.