Train TinyLlama 1.1B Locally on Own Custom Dataset

 This video explains in easy and simple tutorial as how to train or fine-tune TinyLlama model locally by using unsloth on your own data.

Code Used:

import torch

major_version, minor_version = torch.cuda.get_device_capability()

!pip install "unsloth[colab] @ git+https://github.com/unslothai/unsloth.git"

from unsloth import FastLanguageModel

import torch

max_seq_length = 4096

dtype = None

load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(

    model_name = "unsloth/tinyllama-bnb-4bit",

    max_seq_length = max_seq_length,

    dtype = dtype,

    load_in_4bit = load_in_4bit,

)

model = FastLanguageModel.get_peft_model(

    model,

    r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128

    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",

                      "gate_proj", "up_proj", "down_proj",],

    lora_alpha = 32,

    lora_dropout = 0,

    bias = "none",   

    use_gradient_checkpointing = False,

    random_state = 3407,

    max_seq_length = max_seq_length,

)

from trl import SFTTrainer

from transformers import TrainingArguments

from transformers.utils import logging

logging.set_verbosity_info()

trainer = SFTTrainer(

    model = model,

    train_dataset = dataset,

    dataset_text_field = "text",

    max_seq_length = max_seq_length,

    packing = True, 

    args = TrainingArguments(

        per_device_train_batch_size = 2,

        gradient_accumulation_steps = 4,

        warmup_ratio = 0.1,

        num_train_epochs = 1,

        learning_rate = 2e-5,

        fp16 = not torch.cuda.is_bf16_supported(),

        bf16 = torch.cuda.is_bf16_supported(),

        logging_steps = 1,

        optim = "adamw_8bit",

        weight_decay = 0.1,

        lr_scheduler_type = "linear",

        seed = 3407,

        output_dir = "outputs",

    ),

)

trainer_stats = trainer.train()