PyTorch: Training TinyLlama 1.1B in Google Colab | ML Engineering

Hi All

Today we'll be loading a pre-trained model and train it on Colab. Here's a minimal example using Huggingface for transformers and accelerate for memory efficiency. View source below.

Let's break down the above code. Block by block.

I'm assuming you're using a notebook for this example. However in the full example in my GitHub I've included the Ops files too (Docker and k8s for GKE or AKS). I don't go into those here, I'll have a series of articles tacking MLOps in depth, this is about fine-tuning a pre-trained small LLM in Colab. 

Prep: Setup Environment 

!pip install -q transformers accelerate datasets peft bitsandbytes

 

*    transformers provides pre-trained models and training utilities

*    peft enables LoRA (low-rank Adaptation), reducing memory usage

*     bitsandbytes allows 8-bit precision to shrink model size

*     accelerate helps manage GPU memory and multi-GPU setups

 

Step 1: Load Model and Tokenizer

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Model name (TinyLlama-1.1B)
model_name = "TinyLlama/TinyLlama-1.1B"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load model with 8-bit quantization and auto device mapping
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # Use mixed precision
    device_map="auto",         # Automatically use available GPU
    load_in_8bit=True,         # Quantize to 8-bit
)

*    device_map="auto" : Automatically offloads memory to CPU if GPU memoy is full (useful for T4 16GB limit).

*    load_in_8bit=True :  Reduces model size from ~2.5GB (fp16) to 1.2GB(8-bit), freeing up memory.

*    torch_dtype=torch_fpoint16 : Uses half precision (fp16), for faster training and lower memory usage.

*    Memory Impact : Without quantization ~16GB of VRAM needed. With 8-bit and fp16 memory requirements shrink to ~8-10GB of VRAM (fits on T4).

 

Step 2: Apply LoRA (Low rank Adaptation)

LoRA freezes the base model and trains only low-rank matrices, drastically reducing memory useage.

from peft import LoraConfig, get_peft_model

# LoRA configuration
lora_config = LoraConfig(
    r=8,                  # Rank of low-rank matrices (smaller = less memory)
    lora_alpha=32,        # Scaling factor for LoRA weights
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Target attention layers
    lora_dropout=0.05,    # Dropout for regularization
    bias="none",          # No bias terms
    task_type="CAUSAL_LM", # For causal language modeling
)

# Apply LoRA to the model
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # Check how many parameters are trainable

 Output example:

trainable params: 1,179,648 || all params: 1,115,452,416 || trainable%: 0.1058%

 *    Intepretation: Only 1.2M parameters are trainable (vs 1.1B total), reducing memory usage further.

 

Step 3: Load and Preprocess Dataset

 We'll use the tiny_shakespear dataset (a small text corpus for testing).

from datasets import load_dataset

# Load dataset
dataset = load_dataset("tiny_shakespeare", split="train")

# Tokenize the dataset
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,  # Limit sequence length
        return_overflowing_tokens=True,
        return_length=True,
    )

# Apply tokenization
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"],  # Remove original text column
)

 *    max_length=512: Limits input sequence length to fit in GPU memory.

*    truncation=true: Truncates long sequences to max_length. 

*    return_overflowing_tokens=True: Handles sequences longer than max_length by splitting them.

*    Memory Tip: Smaller max_length=less memory usage. Start with 256 or 512 and adjust based on GPU limits.

 

Step 4: Define Training Arguments 

Configure hyperparameters for training. 

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",       # Directory to save checkpoints
    per_device_train_batch_size=4, # Batch size per GPU (adjust based on memory)
    gradient_accumulation_steps=4, # Simulate larger batch size
    num_train_epochs=1,           # Number of epochs
    fp16=True,                    # Use mixed precision
    save_steps=500,               # Save checkpoint every 500 steps
    logging_steps=100,            # Log metrics every 100 steps
    learning_rate=2e-5,           # Learning rate
    weight_decay=0.01,            # L2 regularization
    warmup_steps=100,             # Learning rate warmup
    optim="paged_adamw_8bit",     # Optimizer with 8-bit states
)

 *    per_device_train_batch_size=4: Start with 4; increase if GPU memory allows.

 *    gradient_accumulation_steps=4: Effectively makes batch size 4 * 4 - 16.

*    optim="paged_adamw_8bit": Uses 8-bit AdamW optimizer to save memory.

*    fp16=True: Enables mixed-precision training (faster and lower memory).

 

Step 5: Initialize Trainer and Train

from transformers import Trainer

# Initialize Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)

# Start training
trainer.train()

  The trainer:

    * Splits dataset into batches

    * Computes gradients and updates LoRA weights

    * Saves checkpoints to ./results

Progress is logged every logging_steps=100. 

Expected Output:

{'loss': 3.5678, 'learning_rate': 2e-05, 'epoch': 0.1}
{'loss': 3.4567, 'learning_rate': 2e-05, 'epoch': 0.2}
...

 Step 6: Save and Test the Model

After training, save the model and the tokenizer.

Output Example:

To be or not to be, that is the question:
Whether 'tis nobler in the mind to suffer
The slings and arrows of outrageous fortune...

 The repo has the full source which goes much more in depth and covers some SRE work (k8s). Which is not covered here. So yes, we can train a small model like Llama 1.1B in Kaggle/Colab but it comes with limitations, work within those and you'll be fine. That's all folks.

Later

- Ed