{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hYw25Om2abp1"
      },
      "source": [
        "# Fine-tune Deepseek-R1 1.5B with LoRA on Colab T4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "msHV5EAhmQBH"
      },
      "source": [
        "- Fine-Tune DeepSeek R1 1.5B on Free GCP Colab T4: A Hands-On Guide with LoRA\n",
        "- SFT Trainer: https://huggingface.co/docs/trl/en/sft_trainer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2qMs2zGMs2IZ"
      },
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81xLAPOBsHDn"
      },
      "outputs": [],
      "source": [
        "# Install required packages\n",
        "!pip install -q transformers datasets peft accelerate bitsandbytes trl torch\n",
        "!pip install -q sentencepiece protobuf\n",
        "\n",
        "print(\"✅ All packages installed successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "imports"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "imports_code"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import os\n",
        "import gc\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    BitsAndBytesConfig,\n",
        "    TrainingArguments,\n",
        "    pipeline\n",
        ")\n",
        "from peft import (\n",
        "    LoraConfig,\n",
        "    get_peft_model,\n",
        "    prepare_model_for_kbit_training,\n",
        "    PeftModel,\n",
        "    PeftConfig\n",
        ")\n",
        "from trl import SFTTrainer\n",
        "from datasets import Dataset, load_dataset\n",
        "import pandas as pd\n",
        "\n",
        "print(\"✅ All imports successful!\")\n",
        "print(f\"PyTorch version: {torch.__version__}\")\n",
        "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
        "if torch.cuda.is_available():\n",
        "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
        "    print(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "model_config"
      },
      "source": [
        "## Model Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "model_config_code"
      },
      "outputs": [],
      "source": [
        "# Model and tokenizer names\n",
        "model_name = \"deepseek-ai/deepseek-r1-1.5b\"  # or use \"deepseek-ai/deepseek-r1-7b\" for larger model\n",
        "\n",
        "# Quantization config for 4-bit to fit in T4 memory\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        ")\n",
        "\n",
        "# Load tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\"\n",
        "\n",
        "# Load model with 4-bit quantization\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_name,\n",
        "    quantization_config=bnb_config,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True,\n",
        ")\n",
        "\n",
        "# Enable gradient checkpointing to save memory\n",
        "model.gradient_checkpointing_enable()\n",
        "\n",
        "# Prepare model for k-bit training\n",
        "model = prepare_model_for_kbit_training(model)\n",
        "\n",
        "print(\"✅ Model and tokenizer loaded successfully!\")\n",
        "print(f\"Model parameters: {model.num_parameters():,}\")\n",
        "print(f\"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lora_config"
      },
      "source": [
        "## LoRA Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lora_config_code"
      },
      "outputs": [],
      "source": [
        "# LoRA configuration\n",
        "lora_config = LoraConfig(\n",
        "    r=16,  # rank\n",
        "    lora_alpha=32,\n",
        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")\n",
        "\n",
        "# Apply LoRA to model\n",
        "model = get_peft_model(model, lora_config)\n",
        "model.print_trainable_parameters()\n",
        "\n",
        "print(\"✅ LoRA configuration applied!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset"
      },
      "source": [
        "## Dataset Preparation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dataset_code"
      },
      "outputs": [],
      "source": [
        "# Create a sample dataset for fine-tuning\n",
        "# Replace this with your own dataset\n",
        "\n",
        "# Example dataset: Math QA pairs\n",
        "dataset_data = [\n",
        "    {\n",
        "        \"instruction\": \"What is 5 + 3?\",\n",
        "        \"response\": \"5 + 3 = 8. The sum of 5 and 3 is 8.\"\n",
        "    },\n",
        "    {\n",
        "        \"instruction\": \"Solve for x: 2x + 3 = 7\",\n",
        "        \"response\": \"2x + 3 = 7\\nSubtract 3 from both sides: 2x = 4\\nDivide both sides by 2: x = 2\"\n",
        "    },\n",
        "    {\n",
        "        \"instruction\": \"What is the square root of 16?\",\n",
        "        \"response\": \"The square root of 16 is 4, because 4 × 4 = 16.\"\n",
        "    },\n",
        "    {\n",
        "        \"instruction\": \"Calculate the area of a circle with radius 3.\",\n",
        "        \"response\": \"Area = π × r²\\nArea = π × 3²\\nArea = 9π ≈ 28.27 square units\"\n",
        "    },\n",
        "    {\n",
        "        \"instruction\": \"What is 15% of 200?\",\n",
        "        \"response\": \"15% of 200 = 0.15 × 200 = 30\"\n",
        "    },\n",
        "]\n",
        "\n",
        "# Convert to dataset\n",
        "dataset = Dataset.from_list(dataset_data)\n",
        "\n",
        "# Format function for instruction-response pairs\n",
        "def format_example(example):\n",
        "    return {\n",
        "        \"text\": f\"### Instruction:\\n{example['instruction']}\\n\\n### Response:\\n{example['response']}\"\n",
        "    }\n",
        "\n",
        "# Apply formatting\n",
        "dataset = dataset.map(format_example)\n",
        "\n",
        "# Split into train and validation\n",
        "dataset = dataset.train_test_split(test_size=0.2, seed=42)\n",
        "train_dataset = dataset[\"train\"]\n",
        "eval_dataset = dataset[\"test\"]\n",
        "\n",
        "print(f\"✅ Dataset loaded successfully!\")\n",
        "print(f\"Training samples: {len(train_dataset)}\")\n",
        "print(f\"Validation samples: {len(eval_dataset)}\")\n",
        "print(f\"\\nSample:\")\n",
        "print(train_dataset[0][\"text\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "training_args"
      },
      "source": [
        "## Training Arguments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "training_args_code"
      },
      "outputs": [],
      "source": [
        "# Training arguments optimized for T4\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./deepseek-r1-lora\",\n",
        "    num_train_epochs=3,\n",
        "    per_device_train_batch_size=2,\n",
        "    per_device_eval_batch_size=2,\n",
        "    gradient_accumulation_steps=4,\n",
        "    eval_strategy=\"steps\",\n",
        "    eval_steps=50,\n",
        "    save_strategy=\"steps\",\n",
        "    save_steps=100,\n",
        "    logging_steps=10,\n",
        "    learning_rate=2e-4,\n",
        "    warmup_ratio=0.03,\n",
        "    lr_scheduler_type=\"constant\",\n",
        "    bf16=True,\n",
        "    tf32=True,\n",
        "    max_grad_norm=0.3,\n",
        "    report_to=\"none\",\n",
        "    optim=\"paged_adamw_8bit\",\n",
        "    load_best_model_at_end=True,\n",
        "    metric_for_best_model=\"eval_loss\",\n",
        "    greater_is_better=False,\n",
        "    ddp_find_unused_parameters=False,\n",
        "    push_to_hub=False,\n",
        ")\n",
        "\n",
        "print(\"✅ Training arguments configured!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "trainer"
      },
      "source": [
        "## SFT Trainer Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "trainer_code"
      },
      "outputs": [],
      "source": [
        "# Initialize SFT Trainer\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    train_dataset=train_dataset,\n",
        "    eval_dataset=eval_dataset,\n",
        "    args=training_args,\n",
        "    max_seq_length=512,\n",
        "    dataset_text_field=\"text\",\n",
        ")\n",
        "\n",
        "print(\"✅ SFT Trainer initialized!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "training"
      },
      "source": [
        "## Start Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "training_code",
        "outputId": "f3a2b1c4-5d6e-7f8a-9b0c-1d2e3f4a5b6c"
      },
      "outputs": [],
      "source": [
        "# Clear cache before training\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "# Start training\n",
        "print(\"🚀 Starting training...\")\n",
        "trainer.train()\n",
        "\n",
        "print(\"✅ Training completed!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "save_model"
      },
      "source": [
        "## Save Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "save_model_code"
      },
      "outputs": [],
      "source": [
        "# Save the model\n",
        "model_path = \"./deepseek-r1-lora-final\"\n",
        "trainer.save_model(model_path)\n",
        "tokenizer.save_pretrained(model_path)\n",
        "\n",
        "print(f\"✅ Model saved to {model_path}\")\n",
        "\n",
        "# Optionally, push to Hugging Face Hub\n",
        "# trainer.push_to_hub(\"your-username/deepseek-r1-lora\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "inference"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "inference_code"
      },
      "outputs": [],
      "source": [
        "# Load the trained model\n",
        "def load_trained_model(model_path=\"./deepseek-r1-lora-final\"):\n",
        "    # Load base model\n",
        "    base_model = AutoModelForCausalLM.from_pretrained(\n",
        "        model_name,\n",
        "        quantization_config=bnb_config,\n",
        "        device_map=\"auto\",\n",
        "        trust_remote_code=True,\n",
        "    )\n",
        "    \n",
        "    # Load LoRA weights\n",
        "    model = PeftModel.from_pretrained(base_model, model_path)\n",
        "    \n",
        "    # Merge LoRA weights into base model (optional)\n",
        "    model = model.merge_and_unload()\n",
        "    \n",
        "    return model\n",
        "\n",
        "# Load model for inference\n",
        "inference_model = load_trained_model()\n",
        "print(\"✅ Inference model loaded!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "inference_test",
        "outputId": "g1h2i3j4-5k6l-7m8n-9o0p-1q2r3s4t5u6v"
      },
      "outputs": [],
      "source": [
        "# Test the model\n",
        "def generate_response(prompt, max_length=200):\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\", truncation=True, max_length=512)\n",
        "    inputs = {k: v.to(inference_model.device) for k, v in inputs.items()}\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        outputs = inference_model.generate(\n",
        "            **inputs,\n",
        "            max_new_tokens=max_length,\n",
        "            temperature=0.7,\n",
        "            do_sample=True,\n",
        "            top_p=0.95,\n",
        "            pad_token_id=tokenizer.eos_token_id,\n",
        "        )\n",
        "    \n",
        "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "    return response\n",
        "\n",
        "# Test with a sample prompt\n",
        "test_prompt = \"### Instruction:\\nWhat is 2 + 2?\\n\\n### Response:\\n\"\n",
        "response = generate_response(test_prompt)\n",
        "print(\"Prompt:\", test_prompt)\n",
        "print(\"Response:\", response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "troubleshooting"
      },
      "source": [
        "## Troubleshooting Tips\n",
        "\n",
        "### Common Issues and Fixes:\n",
        "\n",
        "1. **CUDA Out of Memory**:\n",
        "   - Reduce batch size: `per_device_train_batch_size=1`\n",
        "   - Increase gradient accumulation: `gradient_accumulation_steps=8`\n",
        "   - Reduce max sequence length: `max_seq_length=256`\n",
        "\n",
        "2. **Import Errors**:\n",
        "   - Restart runtime after installation\n",
        "   - Update packages: `!pip install --upgrade transformers datasets peft`\n",
        "\n",
        "3. **Tokenizer Issues**:\n",
        "   - Ensure pad token is set: `tokenizer.pad_token = tokenizer.eos_token`\n",
        "   - Set padding side: `tokenizer.padding_side = \"right\"`\n",
        "\n",
        "4. **Slow Training**:\n",
        "   - Use smaller dataset\n",
        "   - Reduce epochs\n",
        "   - Use gradient checkpointing\n",
        "\n",
        "5. **Poor Performance**:\n",
        "   - Increase LoRA rank (r=32 or r=64)\n",
        "   - Increase training epochs\n",
        "   - Use larger dataset\n",
        "   - Adjust learning rate\n",
        "\n",
        "### Alternative Model Sizes:\n",
        "- `deepseek-ai/deepseek-r1-1.5b` (1.5B params) - Fits T4 well\n",
        "- `deepseek-ai/deepseek-r1-7b` (7B params) - May need more memory\n",
        "- `deepseek-ai/deepseek-r1-14b` (14B params) - T4 may not have enough memory"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}