🦾 Pythia-160M LoRA Fine-Tuning

Causal Language Model fine-tuned with LoRA (PEFT) on a diverse 41K-sample corpus spanning news, reviews, stories, sentiment, and encyclopedic text.

---

📌 Project Overview

This project fine-tunes EleutherAI/Pythia-160M using Low-Rank Adaptation (LoRA) for causal language modeling. By combining 10 diverse text datasets, the model develops robust next-token prediction capabilities across multiple domains — from news articles to tweet sentiment — without modifying base model weights.

Why LoRA?

LoRA inserts trainable low-rank matrices into the transformer's attention layers, drastically reducing trainable parameters (from ~160M to ~300k) while retaining the expressiveness of the full model. This enables fine-tuning on a single Tesla T4 GPU in under 12 minutes(725 seconds).

---

🏗️ Architecture

EleutherAI/Pythia-160M (frozen)
         │
         ▼
   LoRA Adapters (r=8)
   ┌─────────────────────┐
   │  q_proj  ← rank-8   │
   │  v_proj  ← rank-8   │
   │  alpha = 16         │
   │  dropout = 0.05     │
   └─────────────────────┘
         │
         ▼
  Causal LM Head

---

📊 Dataset Composition

Total: 41,430 samples → Train: 37,287 · Validation: 4,143

---

⚙️ Training Configuration

# LoRA
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
target_modules: ["query_key_value"]

# Training
num_train_epochs: 1
learning_rate: 2e-4
per_device_train_batch_size: 4
gradient_accumulation_steps: 4
effective_batch_size: 16
max_seq_length: 128
fp16: true

Training Results

---

🚀 Quick Start

1. Installation

git clone https://github.com/kriish446/pythia-160m-lora.git
cd pythia-160m-lora
pip install -r requirements.txt

2. Prepare Data

python scripts/prepare_data.py --output_dir data/processed

3. Train

python scripts/train.py --config configs/training_config.yaml

4. Inference

from src.model.inference import LoRAInference

model = LoRAInference.from_pretrained("outputs/final_model")
output = model.generate("Once upon a time", max_new_tokens=50)
print(output)

5. Evaluate

python scripts/evaluate.py --model_path outputs/final_model --split validation

---

📁 Repository Structure

pythia-160m-lora/
├── configs/                    # YAML configuration files
│   ├── training_config.yaml    # Training hyperparameters
│   ├── lora_config.yaml        # LoRA adapter configuration
│   └── data_config.yaml        # Dataset configuration
├── src/
│   ├── data/                   # Data loading & processing
│   │   ├── __init__.py
│   │   ├── dataset_loader.py   # Loads all 10 datasets
│   │   ├── preprocessor.py     # Tokenization & formatting
│   │   └── data_collator.py    # Batching utilities
│   ├── model/                  # Model definition
│   │   ├── __init__.py
│   │   ├── lora_model.py       # LoRA model setup
│   │   └── inference.py        # Inference utilities
│   ├── training/               # Training loop
│   │   ├── __init__.py
│   │   ├── trainer.py          # Custom trainer
│   │   └── callbacks.py        # W&B, checkpointing callbacks
│   ├── evaluation/             # Evaluation metrics
│   │   ├── __init__.py
│   │   └── metrics.py          # Perplexity, loss metrics
│   └── utils/                  # Shared utilities
│       ├── __init__.py
│       ├── logger.py
│       └── helpers.py
├── scripts/                    # Runnable entry points
│   ├── prepare_data.py
│   ├── train.py
│   └── evaluate.py
├── notebooks/                  # Jupyter exploration
│   ├── 01_data_exploration.ipynb
│   └── 02_model_inference.ipynb
├── tests/                      # Unit & integration tests
│   ├── test_data.py
│   ├── test_model.py
│   └── test_training.py
├── docs/                       # Extended documentation
│   └── training_details.md
├── .github/
│   └── workflows/
│       └── ci.yml              # GitHub Actions CI
├── requirements.txt
├── setup.py
├── .gitignore
└── LICENSE

---

🔬 Experiment Tracking

All runs are tracked with Weights & Biases. To view:

wandb login
# Metrics are auto-logged during training

Key tracked metrics:

• train/loss per step and epoch
• eval/loss per epoch
• eval/perplexity
• Learning rate schedule
• GPU memory utilization

---

📈 Reproducing Results

# Set seed for reproducibility
export PYTHONHASHSEED=42

# Full pipeline
python scripts/prepare_data.py
python scripts/train.py --config configs/training_config.yaml --seed 42
python scripts/evaluate.py --model_path outputs/final_model

---

🧪 Running Tests

pytest tests/ -v --tb=short

---

📄 License

This project is licensed under the MIT License. See LICENSE for details.

---

🙏 Acknowledgements

• EleutherAI for the Pythia model family
• HuggingFace for Transformers, PEFT & Datasets
• Microsoft Research for the LoRA paper