diff --git a/LRP_Merge.ipynb b/LRP_Merge.ipynb new file mode 100644 index 00000000..d0d96a8b --- /dev/null +++ b/LRP_Merge.ipynb @@ -0,0 +1,5321 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "-Es53aQofY5J" + }, + "source": [ + "# LRP-Merge: Layer-wise Relevance Propagation for Model Merging\n", + "\n", + "This notebook demonstrates the LRP-Merge method, which uses Layer-wise Relevance Propagation to identify and preserve the most important weights from multiple task-specific models during a merge.\n", + "\n", + "### Workflow:\n", + "1. **Setup**: Mount Drive and install dependencies.\n", + "2. **Training**: Fine-tune LoRA adapters for specific tasks.\n", + "3. **LRP Analysis**: Compute relevance scores for model weights.\n", + "4. **Model Preparation**: Reconstruct full models from adapters.\n", + "5. **Merging**: Perform LRP-Merge using a custom Mergekit implementation.\n", + "6. **Evaluation**: Test the merged model and optimize merge density." + ] + }, + { + "cell_type": "markdown", + "source": [ + "### **Note**-\n", + "1. Install 'mergekit' inorder to execute this notebook.\n", + "2. The paths mentioned in the code depend on your folder structure in google drive(if you are using colab) or your local system.\n", + "3. Here, I have used \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\" for this experiment only.\n", + "4. I had use FAKE_NEWS_DETECTION dataset to test the merged model only for this experiment." + ], + "metadata": { + "id": "nI6NsNydiyxV" + } + }, + { + "cell_type": "markdown", + "metadata": { + "id": "a220f855" + }, + "source": [ + "## Download Base Model\n", + "Before merging, we need to download the base model weights from Hugging Face and store them in our project folder." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "6507a290" + }, + "source": [ + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "import os\n", + "LRP_PATH=\"your_folder_path\"\n", + "# Define the base model ID and local storage path\n", + "base_model_id = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\" #or use any other model you like\n", + "local_base_path = os.path.join(LRP_PATH, \"base_model\")\n", + "\n", + "if not os.path.exists(local_base_path):\n", + " print(f\"Downloading {base_model_id}...\")\n", + " model = AutoModelForCausalLM.from_pretrained(base_model_id)\n", + " tokenizer = AutoTokenizer.from_pretrained(base_model_id)\n", + "\n", + " os.makedirs(local_base_path, exist_ok=True)\n", + " model.save_pretrained(local_base_path)\n", + " tokenizer.save_pretrained(local_base_path)\n", + " print(f\"Base model saved to: {local_base_path}\")\n", + "else:\n", + " print(f\"Base model already exists at: {local_base_path}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "#command to install mergekit\n", + "!pip install -q -U mergekit" + ], + "metadata": { + "id": "4f1eC6KBjR_6" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "c4546118", + "outputId": "b32b267a-9403-4df5-c835-a45805880fa4" + }, + "source": [ + "import os\n", + "\n", + "# --- PORTABLE PATH SETUP ---\n", + "# If running in Colab with Drive: \"/content/drive/MyDrive/LRP Merge method\"\n", + "# If running locally or after cloning: \".\"\n", + "BASE_DIR = \".\"\n", + "os.chdir(BASE_DIR)\n", + "\n", + "# Define relative subdirectories\n", + "MODEL_DIR = \"models\"\n", + "DATA_DIR = \"datasets\"\n", + "REPO_DIR = \"mergekit_repo\"\n", + "\n", + "# Create directories if they don't exist\n", + "for d in [MODEL_DIR, DATA_DIR]:\n", + " os.makedirs(d, exist_ok=True)\n", + "\n", + "print(f\"✅ Project root set to: {os.getcwd()}\")\n", + "print(f\"Models will be saved to: {os.path.abspath(MODEL_DIR)}\")" + ], + "execution_count": 14, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "✅ Project root set to: /content\n", + "Models will be saved to: /content/models\n" + ] + } + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "6a8cb502", + "outputId": "34cb13b4-7153-4246-d98f-ce7b4541e950" + }, + "source": [ + "import os\n", + "\n", + "# Path to your project folder\n", + "LRP_PATH = \"your_path\"\n", + "gitignore_path = os.path.join(LRP_PATH, \".gitignore\")\n", + "\n", + "gitignore_content = \"\"\"\n", + "# Byte-compiled / optimized / DLL files\n", + "__pycache__/\n", + "*.py[cod]\n", + "*$py.class\n", + "\n", + "# Model weights and Large Files (Crucial for GitHub)\n", + "models/\n", + "base_model/\n", + "mergekit_repo/venv/\n", + "*.bin\n", + "*.safetensors\n", + "*.pt\n", + "*.pth\n", + "*.zip\n", + "*.tar.gz\n", + "\n", + "# Colab/Notebook specific\n", + ".ipynb_checkpoints/\n", + ".virtual_documents/\n", + "\n", + "# Environments\n", + ".env\n", + ".venv\n", + "env/\n", + "venv/\n", + "\n", + "# OS generated files\n", + ".DS_Store\n", + "ehthumbs.db\n", + "Thumbs.db\n", + "\"\"\"\n", + "\n", + "# Ensure the directory exists (it should, but safety first)\n", + "os.makedirs(LRP_PATH, exist_ok=True)\n", + "\n", + "with open(gitignore_path, \"w\") as f:\n", + " f.write(gitignore_content.strip())\n", + "\n", + "print(f\"✅ Created .gitignore at: {gitignore_path}\")\n", + "print(\"This file will prevent large models and temporary files from being tracked by Git.\")" + ], + "execution_count": 19, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "✅ Created .gitignore at: /content/drive/MyDrive/LRP Merge method/.gitignore\n", + "This file will prevent large models and temporary files from being tracked by Git.\n" + ] + } + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "0b98312b", + "outputId": "1f30cdf0-615c-448b-f814-ec60e4f09149" + }, + "source": [ + "gitignore_content = \"\"\"\n", + "# Byte-compiled / optimized / DLL files\n", + "__pycache__/\n", + "*.py[cod]\n", + "*$py.class\n", + "\n", + "# Model weights (Crucial: do not push these!)\n", + "models/\n", + "base_model/\n", + "*.bin\n", + "*.safetensors\n", + "*.zip\n", + "\n", + "# Environments\n", + ".env\n", + ".venv\n", + "env/\n", + "venv/\n", + "\n", + "# Notebook checkpoints\n", + ".ipynb_checkpoints\n", + "\"\"\"\n", + "\n", + "with open(\".gitignore\", \"w\") as f:\n", + " f.write(gitignore_content.strip())\n", + "\n", + "print(\"✅ Created .gitignore to protect your repo from large model files.\")" + ], + "execution_count": 15, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "✅ Created .gitignore to protect your repo from large model files.\n" + ] + } + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "374741e0" + }, + "source": [ + "import subprocess\n", + "\n", + "def run_portable_merge(config_name, output_name):\n", + " # Use relative paths for the merge command\n", + " output_path = os.path.join(MODEL_DIR, output_name)\n", + " os.makedirs(output_path, exist_ok=True)\n", + "\n", + " cmd = [\n", + " \"mergekit-yaml\",\n", + " config_name,\n", + " output_path,\n", + " \"--copy-tokenizer\",\n", + " \"--allow-crimes\",\n", + " \"--lazy-unpickle\"\n", + " ]\n", + "\n", + " print(f\"Executing portable merge: {' '.join(cmd)}\")\n", + " res = subprocess.run(cmd, capture_output=True, text=True)\n", + " if res.returncode == 0:\n", + " print(f\"✅ Successfully merged to {output_path}\")\n", + " else:\n", + " print(f\"✗ Error: {res.stderr}\")" + ], + "execution_count": 16, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "16cf0b1d", + "outputId": "654fe4bd-2c29-4809-be10-818b23881438" + }, + "source": [ + "import os\n", + "\n", + "# Define path for this cell scope\n", + "LRP_PATH = \"your_folder_path\"\n", + "\n", + "print(\"--- Files and Directories in LRP Project ---\")\n", + "if os.path.exists(LRP_PATH):\n", + " # List files and indicate if they are directories\n", + " for item in os.listdir(LRP_PATH):\n", + " full_path = os.path.join(LRP_PATH, item)\n", + " if os.path.isdir(full_path):\n", + " print(f\"[DIR] {item}/\")\n", + " else:\n", + " print(f\"[FILE] {item}\")\n", + "else:\n", + " print(f\"Error: {LRP_PATH} not found.\")\n", + "\n", + "print(\"\\n--- Inside 'models/' (Checking for weights) ---\")\n", + "models_path = os.path.join(LRP_PATH, \"models\")\n", + "if os.path.exists(models_path):\n", + " !ls -F \"{models_path}\"\n" + ], + "execution_count": 13, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "--- Files and Directories in LRP Project ---\n", + "[FILE] finetune_fakenews.py\n", + "[FILE] lrp_computer.py\n", + "[FILE] INSTRUCTIONS_FOR_COLAB.md\n", + "[DIR] datasets/\n", + "[DIR] models/\n", + "[DIR] __pycache__/\n", + "[FILE] lrp_config_colab.yaml\n", + "[DIR] mergekit_repo/\n", + "[DIR] base_model/\n", + "[FILE] lrp_merge_pipeline.py\n", + "[FILE] LRP_Merge_Colab_Training.ipynb\n", + "\n", + "--- Inside 'models/' (Checking for weights) ---\n", + "compare_lrp/\tmerged-model/\t tinyllama-global-full/\n", + "compare_slerp/\tmerged-model-d0.5/ tinyllama-local/\n", + "compare_ties/\tmerged-model-d0.7/ tinyllama-local-full/\n", + "lrp-global/\tmerged-model-d0.9/\n", + "lrp-local/\ttinyllama-global/\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "4DdPOPSnfY5L" + }, + "source": [ + "##Mount Google Drive\n", + "\n", + "This connects your Google Drive to access the LRP merge method files." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "SuyAw89LfY5M", + "outputId": "1c70faaa-c4f7-486a-a007-90a8b62db93f" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Warning: Mount point /content/drive is not empty. Attempting forceful cleanup...\n", + "Drive not mounted, so nothing to flush and unmount.\n", + "Successfully unmounted Google Drive.\n", + "Successfully removed directory: /content/drive\n", + "Mounted at /content/drive\n", + "Google Drive mounted successfully!\n" + ] + } + ], + "source": [ + "# Step 1: Environment Setup\n", + "import os\n", + "from google.colab import drive\n", + "\n", + "drive.mount('/content/drive') #should be used only once for colab\n", + "LRP_PATH = \"your_folder_path\"\n", + "os.chdir(LRP_PATH)\n", + "\n", + "# Install standard dependencies\n", + "!pip install -q transformers datasets accelerate peft bitsandbytes safetensors numpy\n", + "\n", + "# Install the custom LRP-enabled Mergekit from your Drive\n", + "mergekit_repo_path = os.path.join(LRP_PATH, \"mergekit_repo\")\n", + "!pip install -e \"{mergekit_repo_path}\"\n", + "\n", + "print(\"\\n✅ Environment ready and Custom Mergekit installed!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "m7giRqKGfY5Q" + }, + "source": [ + "## Installing Dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "kPBMV65LfY5R", + "outputId": "3a192e03-e9f8-4a44-c812-589e4a75daa9" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.7/60.7 MB\u001b[0m \u001b[31m15.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25h\n", + "All dependencies installed!\n" + ] + } + ], + "source": [ + "!pip install -q transformers datasets accelerate peft bitsandbytes safetensors\n", + "!pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n", + "!pip install -q pandas scikit-learn\n", + "\n", + "print(\"\\nAll dependencies installed!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0Yis2LeRfY5S" + }, + "source": [ + "## Training Task-Specific Adapters\n", + "In this step, we fine-tune two LoRA adapters (Global and Local) on fake news detection datasets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "PpiAn5MOfY5T", + "outputId": "9d470042-b9fe-44d7-8884-0bd01ea79290" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "============================================================\n", + "Checking for datasets...\n", + "============================================================\n", + "\n", + "✓ Synthetic dataset found:\n", + " Train samples: 800\n", + " Test samples: 200\n", + " Columns: ['text', 'label']\n", + " Labels in train: {'REAL': 404, 'FAKE': 396}\n", + "\n", + " Sample data:\n", + " text label\n", + "0 Report shows economy has increased by 15% over past decade. REAL\n", + "1 Government announces new policy regarding AI after public consultation. REAL\n", + "2 Study finds correlation between AI and health benefits in new research. REAL\n" + ] + } + ], + "source": [ + "# Verify datasets are available\n", + "import os\n", + "import pandas as pd\n", + "\n", + "print(\"=\"*60)\n", + "print(\"Checking for datasets...\")\n", + "print(\"=\"*60)\n", + "\n", + "# Check for synthetic dataset\n", + "if os.path.exists(\"datasets/synthetic/train.csv\"): #these paths depend on your folder structure\n", + " df_train = pd.read_csv(\"datasets/synthetic/train.csv\")\n", + " df_test = pd.read_csv(\"datasets/synthetic/test.csv\") if os.path.exists(\"datasets/synthetic/test.csv\") else None\n", + "\n", + " print(\"\\n✓ Synthetic dataset found:\")\n", + " print(f\" Train samples: {len(df_train)}\")\n", + " if df_test is not None:\n", + " print(f\" Test samples: {len(df_test)}\")\n", + " print(f\" Columns: {list(df_train.columns)}\")\n", + " print(f\" Labels in train: {df_train['label'].value_counts().to_dict()}\")\n", + "\n", + " # Show sample\n", + " print(\"\\n Sample data:\")\n", + " print(df_train.head(3).to_string())\n", + "else:\n", + " print(\"\\n✗ Synthetic dataset not found!\")\n", + " print(\" Expected: datasets/synthetic/train.csv\")\n", + " print(\"\\nTo create sample data, run:\")\n", + " print(\" !python download_fakenews_datasets.py --dataset synthetic --output ./datasets\")\n", + " print(\"\\nOr create the folder structure manually with your CSV files.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "nHpFGULJfY5T" + }, + "source": [ + "## Step 8: Train GLOBAL Model\n", + "\n", + "This trains the general knowledge model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "cm6LGAehfY5T", + "outputId": "c98da997-db25-4387-dbf4-e05cc59000e9" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Training GLOBAL model:\n", + " Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0\n", + " Dataset: datasets/synthetic/train.csv\n", + " Epochs: 3\n", + " Batch size: 4\n", + " Max samples: 1000\n", + " Output: models/tinyllama-global\n" + ] + } + ], + "source": [ + "# Configuration for GLOBAL model (using your existing dataset)\n", + "MODEL_NAME = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\" #Use larger model with GPU!\n", + "# MODEL_NAME = \"gpt2\" # Or use smaller model if memory is limited\n", + "\n", + "#Use your synthetic dataset\n", + "DATASET = \"datasets/synthetic/train.csv\" # Your downloaded dataset\n", + "OUTPUT = \"models/tinyllama-global\"\n", + "\n", + "# Training parameters\n", + "EPOCHS = 3\n", + "BATCH_SIZE = 4 # Can increase with GPU\n", + "MAX_SAMPLES = 1000 # Use all 800 samples or limit\n", + "\n", + "print(f\"Training GLOBAL model:\")\n", + "print(f\" Model: {MODEL_NAME}\")\n", + "print(f\" Dataset: {DATASET}\")\n", + "print(f\" Epochs: {EPOCHS}\")\n", + "print(f\" Batch size: {BATCH_SIZE}\")\n", + "print(f\" Max samples: {MAX_SAMPLES}\")\n", + "print(f\" Output: {OUTPUT}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "eyyQpTx_fY5U", + "outputId": "3ffa4644-e645-409f-a2f0-85c2a39889c5" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "2026-03-29 07:38:50,334 - INFO - Starting fine-tuning...\n", + "2026-03-29 07:38:50,334 - INFO - Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0\n", + "2026-03-29 07:38:50,334 - INFO - Dataset: datasets/synthetic/train.csv\n", + "2026-03-29 07:38:50,334 - INFO - Output: models/tinyllama-global\n", + "2026-03-29 07:38:50,334 - INFO - Epochs: 3\n", + "2026-03-29 07:38:50,334 - INFO - LoRA: True\n", + "2026-03-29 07:38:50,334 - INFO - 8-bit: False\n", + "2026-03-29 07:38:50,334 - INFO - 4-bit: False\n", + "2026-03-29 07:38:50,334 - INFO - Loading dataset from datasets/synthetic/train.csv...\n", + "2026-03-29 07:38:50,340 - INFO - Loaded 800 samples\n", + "2026-03-29 07:38:50,340 - INFO - FAKE: 396\n", + "2026-03-29 07:38:50,340 - INFO - REAL: 404\n", + "2026-03-29 07:38:50,340 - INFO - Loading tokenizer and model...\n", + "2026-03-29 07:38:50,494 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:38:50,505 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:50,516 - INFO - HTTP Request: GET https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "config.json: 100% 608/608 [00:00<00:00, 2.95MB/s]\n", + "2026-03-29 07:38:50,604 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer_config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "2026-03-29 07:38:50,605 - WARNING - Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "2026-03-29 07:38:50,615 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer_config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:50,625 - INFO - HTTP Request: GET https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer_config.json \"HTTP/1.1 200 OK\"\n", + "tokenizer_config.json: 1.29kB [00:00, 4.22MB/s]\n", + "2026-03-29 07:38:50,711 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer_config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:38:50,721 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer_config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:50,811 - INFO - HTTP Request: GET https://huggingface.co/api/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/tree/main/additional_chat_templates?recursive=false&expand=false \"HTTP/1.1 404 Not Found\"\n", + "2026-03-29 07:38:50,901 - INFO - HTTP Request: GET https://huggingface.co/api/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/tree/main?recursive=true&expand=false \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:50,987 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:38:50,998 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:51,010 - INFO - HTTP Request: GET https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer.json \"HTTP/1.1 200 OK\"\n", + "tokenizer.json: 1.84MB [00:00, 73.7MB/s]\n", + "2026-03-29 07:38:51,130 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer.model \"HTTP/1.1 302 Found\"\n", + "2026-03-29 07:38:51,258 - INFO - HTTP Request: GET https://huggingface.co/api/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/xet-read-token/fe8a4ea1ffedaf415f4da2f062534de366a451e6 \"HTTP/1.1 200 OK\"\n", + "tokenizer.model: 100% 500k/500k [00:00<00:00, 680kB/s] \n", + "2026-03-29 07:38:52,085 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/added_tokens.json \"HTTP/1.1 404 Not Found\"\n", + "2026-03-29 07:38:52,172 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/special_tokens_map.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:38:52,182 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/special_tokens_map.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:52,193 - INFO - HTTP Request: GET https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/special_tokens_map.json \"HTTP/1.1 200 OK\"\n", + "special_tokens_map.json: 100% 551/551 [00:00<00:00, 3.50MB/s]\n", + "2026-03-29 07:38:52,281 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/chat_template.jinja \"HTTP/1.1 404 Not Found\"\n", + "2026-03-29 07:38:52,497 - INFO - CUDA available: True\n", + "2026-03-29 07:38:52,581 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:38:52,591 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:52,691 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/adapter_config.json \"HTTP/1.1 404 Not Found\"\n", + "`torch_dtype` is deprecated! Use `dtype` instead!\n", + "2026-03-29 07:38:52,789 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:38:52,800 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:38:52,893 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/model.safetensors \"HTTP/1.1 302 Found\"\n", + "model.safetensors: 100% 2.20G/2.20G [00:15<00:00, 147MB/s]\n", + "Loading weights: 100% 201/201 [00:04<00:00, 48.29it/s, Materializing param=model.norm.weight]\n", + "2026-03-29 07:39:12,645 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/generation_config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:39:12,655 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/generation_config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:39:12,666 - INFO - HTTP Request: GET https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/generation_config.json \"HTTP/1.1 200 OK\"\n", + "generation_config.json: 100% 124/124 [00:00<00:00, 658kB/s]\n", + "2026-03-29 07:39:12,670 - INFO - Gradient checkpointing enabled\n", + "2026-03-29 07:39:12,670 - INFO - Applying LoRA (r=16)...\n", + "trainable params: 4,505,600 || all params: 1,104,553,984 || trainable%: 0.4079\n", + "2026-03-29 07:39:19,059 - INFO - Creating dataset...\n", + "2026-03-29 07:39:19,082 - INFO - Initializing trainer...\n", + "2026-03-29 07:39:19,096 - INFO - Starting training...\n", + "{'loss': '2.963', 'grad_norm': '5.442', 'learning_rate': '1.8e-06', 'epoch': '0.2'}\n", + "{'loss': '2.922', 'grad_norm': '5.519', 'learning_rate': '3.8e-06', 'epoch': '0.4'}\n", + "{'loss': '2.871', 'grad_norm': '5.688', 'learning_rate': '5.8e-06', 'epoch': '0.6'}\n", + "{'loss': '2.698', 'grad_norm': '6.25', 'learning_rate': '7.8e-06', 'epoch': '0.8'}\n", + "{'loss': '2.509', 'grad_norm': '3.484', 'learning_rate': '9.8e-06', 'epoch': '1'}\n", + "{'loss': '2.294', 'grad_norm': '3.453', 'learning_rate': '1.18e-05', 'epoch': '1.2'}\n", + "{'loss': '1.999', 'grad_norm': '3.869', 'learning_rate': '1.38e-05', 'epoch': '1.4'}\n", + "{'loss': '1.592', 'grad_norm': '3.431', 'learning_rate': '1.58e-05', 'epoch': '1.6'}\n", + "{'loss': '1.355', 'grad_norm': '1.376', 'learning_rate': '1.78e-05', 'epoch': '1.8'}\n", + "{'loss': '1.212', 'grad_norm': '1.36', 'learning_rate': '1.98e-05', 'epoch': '2'}\n", + "{'loss': '1.087', 'grad_norm': '1.369', 'learning_rate': '1.64e-05', 'epoch': '2.2'}\n", + "{'loss': '0.9363', 'grad_norm': '1.481', 'learning_rate': '1.24e-05', 'epoch': '2.4'}\n", + "{'loss': '0.8383', 'grad_norm': '1.337', 'learning_rate': '8.4e-06', 'epoch': '2.6'}\n", + "{'loss': '0.8105', 'grad_norm': '1.799', 'learning_rate': '4.4e-06', 'epoch': '2.8'}\n", + "{'loss': '0.769', 'grad_norm': '2.05', 'learning_rate': '4e-07', 'epoch': '3'}\n", + "100% 150/150 [04:51<00:00, 1.98s/it]2026-03-29 07:44:10,966 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:44:10,976 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:44:11,061 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:44:11,072 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "{'train_runtime': '292.1', 'train_samples_per_second': '8.215', 'train_steps_per_second': '0.513', 'train_loss': '1.79', 'epoch': '3'}\n", + "100% 150/150 [04:52<00:00, 1.95s/it]\n", + "2026-03-29 07:44:11,539 - INFO - Saving model to models/tinyllama-global...\n", + "2026-03-29 07:44:11,651 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:44:11,661 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:44:11,748 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:44:11,758 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:44:11,985 - INFO - Fine-tuning complete!\n", + "\n", + "============================================================\n", + "Fine-tuned model saved to: models/tinyllama-global\n", + "============================================================\n" + ] + } + ], + "source": [ + "# (Assuming datasets are present in LRP_PATH/datasets)\n", + "# Train Global and Local models using the finetune_fakenews.py script\n", + "!python finetune_fakenews.py --dataset datasets/synthetic/train.csv --output models/tinyllama-global --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --epochs 3 --use-lora\n", + "!python finetune_fakenews.py --dataset datasets/synthetic/train.csv --output models/tinyllama-local --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --epochs 3 --use-lora" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "V-uP2FOFfY5U" + }, + "source": [ + "##Train LOCAL Model\n", + "\n", + "This trains the task-specific model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "Uc3SMHrCfY5U", + "outputId": "70ee41cd-1106-4c4e-a966-332e64a6f62f" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Training LOCAL model:\n", + " Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0\n", + " Dataset: datasets/synthetic/train.csv\n", + " Epochs: 3\n", + " Batch size: 4\n", + " Max samples: 1000\n", + " Output: models/tinyllama-local\n", + "\n", + "Note: Both models currently use the same dataset.\n", + "For a true LRP-Merge experiment, train on different datasets:\n", + " - GLOBAL: General knowledge/tasks\n", + " - LOCAL: Specific task/domain\n" + ] + } + ], + "source": [ + "# Configuration for LOCAL model\n", + "# Using the same synthetic dataset for task-specific training\n", + "# In a real scenario, you might have a different dataset for the local model\n", + "\n", + "LOCAL_OUTPUT = \"models/tinyllama-local\"\n", + "LOCAL_DATASET = \"datasets/synthetic/train.csv\" # Can be same or different dataset\n", + "\n", + "print(f\"Training LOCAL model:\")\n", + "print(f\" Model: {MODEL_NAME}\")\n", + "print(f\" Dataset: {LOCAL_DATASET}\")\n", + "print(f\" Epochs: {EPOCHS}\")\n", + "print(f\" Batch size: {BATCH_SIZE}\")\n", + "print(f\" Max samples: {MAX_SAMPLES}\")\n", + "print(f\" Output: {LOCAL_OUTPUT}\")\n", + "print(f\"\\nNote: Both models currently use the same dataset.\")\n", + "print(\"For a true LRP-Merge experiment, train on different datasets:\")\n", + "print(\" - GLOBAL: General knowledge/tasks\")\n", + "print(\" - LOCAL: Specific task/domain\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "jW042sK3fY5V", + "outputId": "bb5507aa-f105-4427-d517-8ca6a72b16aa" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "2026-03-29 07:45:23,041 - INFO - Starting fine-tuning...\n", + "2026-03-29 07:45:23,041 - INFO - Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0\n", + "2026-03-29 07:45:23,041 - INFO - Dataset: datasets/synthetic/train.csv\n", + "2026-03-29 07:45:23,041 - INFO - Output: models/tinyllama-local\n", + "2026-03-29 07:45:23,041 - INFO - Epochs: 3\n", + "2026-03-29 07:45:23,041 - INFO - LoRA: True\n", + "2026-03-29 07:45:23,041 - INFO - 8-bit: False\n", + "2026-03-29 07:45:23,041 - INFO - 4-bit: False\n", + "2026-03-29 07:45:23,041 - INFO - Loading dataset from datasets/synthetic/train.csv...\n", + "2026-03-29 07:45:23,049 - INFO - Loaded 800 samples\n", + "2026-03-29 07:45:23,049 - INFO - FAKE: 396\n", + "2026-03-29 07:45:23,049 - INFO - REAL: 404\n", + "2026-03-29 07:45:23,049 - INFO - Loading tokenizer and model...\n", + "2026-03-29 07:45:23,409 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:45:23,429 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:45:23,563 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer_config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "2026-03-29 07:45:23,564 - WARNING - Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "2026-03-29 07:45:23,577 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer_config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:45:23,663 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer_config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:45:23,673 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/tokenizer_config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:45:23,761 - INFO - HTTP Request: GET https://huggingface.co/api/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/tree/main/additional_chat_templates?recursive=false&expand=false \"HTTP/1.1 404 Not Found\"\n", + "2026-03-29 07:45:23,849 - INFO - HTTP Request: GET https://huggingface.co/api/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/tree/main?recursive=true&expand=false \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:45:24,058 - INFO - CUDA available: True\n", + "2026-03-29 07:45:24,142 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:45:24,152 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "`torch_dtype` is deprecated! Use `dtype` instead!\n", + "2026-03-29 07:45:24,245 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:45:24,255 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "Loading weights: 100% 201/201 [00:02<00:00, 92.70it/s, Materializing param=model.norm.weight] \n", + "2026-03-29 07:45:26,932 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/generation_config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:45:26,943 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/generation_config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:45:26,945 - INFO - Gradient checkpointing enabled\n", + "2026-03-29 07:45:26,945 - INFO - Applying LoRA (r=16)...\n", + "trainable params: 4,505,600 || all params: 1,104,553,984 || trainable%: 0.4079\n", + "2026-03-29 07:45:29,991 - INFO - Creating dataset...\n", + "2026-03-29 07:45:30,013 - INFO - Initializing trainer...\n", + "2026-03-29 07:45:30,027 - INFO - Starting training...\n", + "{'loss': '2.963', 'grad_norm': '5.224', 'learning_rate': '1.8e-06', 'epoch': '0.2'}\n", + "{'loss': '2.922', 'grad_norm': '5.146', 'learning_rate': '3.8e-06', 'epoch': '0.4'}\n", + "{'loss': '2.873', 'grad_norm': '5.326', 'learning_rate': '5.8e-06', 'epoch': '0.6'}\n", + "{'loss': '2.703', 'grad_norm': '5.718', 'learning_rate': '7.8e-06', 'epoch': '0.8'}\n", + "{'loss': '2.513', 'grad_norm': '3.289', 'learning_rate': '9.8e-06', 'epoch': '1'}\n", + "{'loss': '2.297', 'grad_norm': '3.386', 'learning_rate': '1.18e-05', 'epoch': '1.2'}\n", + "{'loss': '2.002', 'grad_norm': '3.824', 'learning_rate': '1.38e-05', 'epoch': '1.4'}\n", + "{'loss': '1.596', 'grad_norm': '3.423', 'learning_rate': '1.58e-05', 'epoch': '1.6'}\n", + "{'loss': '1.358', 'grad_norm': '1.35', 'learning_rate': '1.78e-05', 'epoch': '1.8'}\n", + "{'loss': '1.212', 'grad_norm': '1.324', 'learning_rate': '1.98e-05', 'epoch': '2'}\n", + "{'loss': '1.086', 'grad_norm': '1.324', 'learning_rate': '1.64e-05', 'epoch': '2.2'}\n", + "{'loss': '0.9342', 'grad_norm': '1.382', 'learning_rate': '1.24e-05', 'epoch': '2.4'}\n", + "{'loss': '0.8404', 'grad_norm': '1.275', 'learning_rate': '8.4e-06', 'epoch': '2.6'}\n", + "{'loss': '0.8157', 'grad_norm': '1.746', 'learning_rate': '4.4e-06', 'epoch': '2.8'}\n", + "{'loss': '0.7748', 'grad_norm': '1.881', 'learning_rate': '4e-07', 'epoch': '3'}\n", + "100% 150/150 [04:59<00:00, 2.00s/it]2026-03-29 07:50:29,517 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:50:29,529 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:50:29,614 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:50:29,625 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "{'train_runtime': '299.7', 'train_samples_per_second': '8.008', 'train_steps_per_second': '0.501', 'train_loss': '1.793', 'epoch': '3'}\n", + "100% 150/150 [04:59<00:00, 2.00s/it]\n", + "2026-03-29 07:50:30,000 - INFO - Saving model to models/tinyllama-local...\n", + "2026-03-29 07:50:30,114 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:50:30,125 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:50:30,217 - INFO - HTTP Request: HEAD https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/config.json \"HTTP/1.1 307 Temporary Redirect\"\n", + "2026-03-29 07:50:30,228 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0/fe8a4ea1ffedaf415f4da2f062534de366a451e6/config.json \"HTTP/1.1 200 OK\"\n", + "2026-03-29 07:50:30,418 - INFO - Fine-tuning complete!\n", + "\n", + "============================================================\n", + "Fine-tuned model saved to: models/tinyllama-local\n", + "============================================================\n" + ] + } + ], + "source": [ + "!python finetune_fakenews.py \\\n", + " --dataset {LOCAL_DATASET} \\\n", + " --output {LOCAL_OUTPUT} \\\n", + " --model {MODEL_NAME} \\\n", + " --epochs {EPOCHS} \\\n", + " --batch-size {BATCH_SIZE} \\\n", + " --use-lora \\\n", + " --lora-r 16 \\\n", + " --max-samples {MAX_SAMPLES} \\\n", + " --max-length 256" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "P6elKTCgfY5V" + }, + "source": [ + "## Computing LRP Scores\n", + "We calculate the importance of each weight using the LRP epsilon rule." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "mT3EWmZIfY5V", + "outputId": "787d7d3d-86a7-45ac-cbae-fb2024223b0b" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Sun Mar 29 07:52:39 2026 \n", + "+-----------------------------------------------------------------------------------------+\n", + "| NVIDIA-SMI 580.82.07 Driver Version: 580.82.07 CUDA Version: 13.0 |\n", + "+-----------------------------------------+------------------------+----------------------+\n", + "| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\n", + "| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\n", + "| | | MIG M. |\n", + "|=========================================+========================+======================|\n", + "| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\n", + "| N/A 60C P8 10W / 70W | 3MiB / 15360MiB | 0% Default |\n", + "| | | N/A |\n", + "+-----------------------------------------+------------------------+----------------------+\n", + "\n", + "+-----------------------------------------------------------------------------------------+\n", + "| Processes: |\n", + "| GPU GI CI PID Type Process name GPU Memory |\n", + "| ID ID Usage |\n", + "|=========================================================================================|\n", + "| No running processes found |\n", + "+-----------------------------------------------------------------------------------------+\n", + "Computing LRP scores for models/tinyllama-global...\n", + "Output will be saved to models/lrp-global/lrp_scores\n", + "Loading model from models/tinyllama-global...\n", + " Using device: cuda\n", + " Using dtype: torch.float16\n", + "`torch_dtype` is deprecated! Use `dtype` instead!\n", + "Loading weights: 100% 201/201 [00:02<00:00, 99.55it/s, Materializing param=model.norm.weight] \n", + "Loading weights: 100% 176/176 [00:00<00:00, 869.85it/s, Materializing param=model.layers.21.self_attn.v_proj.lora_B.default.weight] \n", + "Computing LRP scores using epsilon rule...\n", + "Processing model.layers.0.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.0.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.0.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.0.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.o_proj.lora_B.default.weight...\n", + "LRP scores saved to models/lrp-global/lrp_scores\n", + "LRP scores computed successfully! Saved to models/lrp-global/lrp_scores\n", + "\n", + "✓ LRP scores computed for GLOBAL model\n" + ] + } + ], + "source": [ + "!python lrp_merge_pipeline.py --compute-lrp --model models/tinyllama-global --output models/lrp-global\n", + "!python lrp_merge_pipeline.py --compute-lrp --model models/tinyllama-local --output models/lrp-local" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "P36cXjdUfY5V", + "outputId": "5e182196-b403-4d8c-efc0-51aa3f289ff0" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Sun Mar 29 07:53:02 2026 \n", + "+-----------------------------------------------------------------------------------------+\n", + "| NVIDIA-SMI 580.82.07 Driver Version: 580.82.07 CUDA Version: 13.0 |\n", + "+-----------------------------------------+------------------------+----------------------+\n", + "| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\n", + "| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\n", + "| | | MIG M. |\n", + "|=========================================+========================+======================|\n", + "| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\n", + "| N/A 60C P8 11W / 70W | 3MiB / 15360MiB | 0% Default |\n", + "| | | N/A |\n", + "+-----------------------------------------+------------------------+----------------------+\n", + "\n", + "+-----------------------------------------------------------------------------------------+\n", + "| Processes: |\n", + "| GPU GI CI PID Type Process name GPU Memory |\n", + "| ID ID Usage |\n", + "|=========================================================================================|\n", + "| No running processes found |\n", + "+-----------------------------------------------------------------------------------------+\n", + "Computing LRP scores for models/tinyllama-local...\n", + "Output will be saved to models/lrp-local/lrp_scores\n", + "Loading model from models/tinyllama-local...\n", + " Using device: cuda\n", + " Using dtype: torch.float16\n", + "`torch_dtype` is deprecated! Use `dtype` instead!\n", + "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "Loading weights: 100% 201/201 [00:02<00:00, 73.89it/s, Materializing param=model.norm.weight]\n", + "Loading weights: 100% 176/176 [00:00<00:00, 963.85it/s, Materializing param=model.layers.21.self_attn.v_proj.lora_B.default.weight]\n", + "Computing LRP scores using epsilon rule...\n", + "Processing model.layers.0.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.0.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.0.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.0.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.0.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.1.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.1.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.2.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.2.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.3.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.3.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.4.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.4.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.5.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.5.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.6.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.6.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.7.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.7.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.8.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.8.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.9.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.9.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.10.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.10.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.11.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.11.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.12.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.12.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.13.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.13.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.14.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.14.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.15.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.15.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.16.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.16.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.17.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.17.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.18.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.18.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.19.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.19.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.20.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.20.self_attn.o_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.q_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.q_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.k_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.k_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.v_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.v_proj.lora_B.default.weight...\n", + "Processing model.layers.21.self_attn.o_proj.lora_A.default.weight...\n", + "Processing model.layers.21.self_attn.o_proj.lora_B.default.weight...\n", + "LRP scores saved to models/lrp-local/lrp_scores\n", + "LRP scores computed successfully! Saved to models/lrp-local/lrp_scores\n", + "\n", + "✓ LRP scores computed for LOCAL model\n" + ] + } + ], + "source": [ + "# Compute LRP for LOCAL model\n", + "!python lrp_merge_pipeline.py --compute-lrp \\\n", + " --model {LOCAL_OUTPUT} \\\n", + " --output models/lrp-local \\\n", + " --device cuda\n", + "\n", + "print(\"\\n✓ LRP scores computed for LOCAL model\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vUHsVGb2fY5V" + }, + "source": [ + "## Reconstructing Full Models\n", + "LRP-Merge performs best when merging full model weights rather than raw adapters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "XNQMG2a0fY5W", + "outputId": "2932e596-955f-45a5-b6c8-51436a1388de" + }, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "✓ Created lrp_config_colab.yaml with corrected structure\n", + "merge_method: lrp\n", + "\n", + "base_model:\n", + " model: \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n", + "\n", + "parameters:\n", + " density: 0.7\n", + " use_lrp: true\n", + "\n", + "models:\n", + " - model:\n", + " model: \"models/tinyllama-global\"\n", + " parameters:\n", + " weight: 1.0\n", + " lrp_scores:\n", + " value: \"./models/lrp-global/lrp_scores\"\n", + " - model:\n", + " model: \"models/tinyllama-local\"\n", + " parameters:\n", + " weight: 1.0\n", + " lrp_scores:\n", + " value: \"./models/lrp-local/lrp_scores\"\n" + ] + } + ], + "source": [ + "# Create LRP config file with corrected structure\n", + "lrp_config = f\"\"\"merge_method: lrp\n", + "\n", + "base_model:\n", + " model: \"{MODEL_NAME}\"\n", + "\n", + "parameters:\n", + " density: 0.7\n", + " use_lrp: true\n", + "\n", + "models:\n", + " - model:\n", + " model: \"{OUTPUT}\"\n", + " parameters:\n", + " weight: 1.0\n", + " lrp_scores:\n", + " value: \"./models/lrp-global/lrp_scores\"\n", + " - model:\n", + " model: \"{LOCAL_OUTPUT}\"\n", + " parameters:\n", + " weight: 1.0\n", + " lrp_scores:\n", + " value: \"./models/lrp-local/lrp_scores\"\n", + "\"\"\"\n", + "\n", + "with open(\"lrp_config_colab.yaml\", \"w\") as f:\n", + " f.write(lrp_config)\n", + "\n", + "print(\"✓ Created lrp_config_colab.yaml with corrected structure\")\n", + "!cat lrp_config_colab.yaml" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "c52adfcb", + "outputId": "520e4943-ead5-4823-d5f5-ccc505c3bce0" + }, + "source": [ + "# Configuration for LOCAL model\n", + "# Using the same synthetic dataset for task-specific training\n", + "# In a real scenario, you might have a different dataset for the local model\n", + "\n", + "LOCAL_OUTPUT = \"models/tinyllama-local\" # Corrected to local path\n", + "LOCAL_DATASET = \"datasets/synthetic/train.csv\" # Can be same or different dataset\n", + "\n", + "print(f\"Training LOCAL model:\")\n", + "print(f\" Model: {MODEL_NAME}\")\n", + "print(f\" Dataset: {LOCAL_DATASET}\")\n", + "print(f\" Epochs: {EPOCHS}\")\n", + "print(f\" Batch size: {BATCH_SIZE}\")\n", + "print(f\" Max samples: {MAX_SAMPLES}\")\n", + "print(f\" Output: {LOCAL_OUTPUT}\")\n", + "print(f\"\\nNote: Both models currently use the same dataset.\")\n", + "print(\"For a true LRP-Merge experiment, train on different datasets:\")\n", + "print(\" - GLOBAL: General knowledge/tasks\")\n", + "print(\" - LOCAL: Specific task/domain\")" + ], + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Training LOCAL model:\n", + " Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0\n", + " Dataset: datasets/synthetic/train.csv\n", + " Epochs: 3\n", + " Batch size: 4\n", + " Max samples: 1000\n", + " Output: models/tinyllama-local\n", + "\n", + "Note: Both models currently use the same dataset.\n", + "For a true LRP-Merge experiment, train on different datasets:\n", + " - GLOBAL: General knowledge/tasks\n", + " - LOCAL: Specific task/domain\n" + ] + } + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "c4f4ed04", + "outputId": "0a028412-83c1-4b36-bcad-936a75b11090" + }, + "source": [ + "import os\n", + "\n", + "# Using the verified LRP_PATH from previous steps\n", + "LRP_PATH = \"/content/drive/MyDrive/LRP Merge method\"\n", + "\n", + "def make_script_portable(filename):\n", + " # Build the full path to the script in Drive\n", + " full_path = os.path.join(LRP_PATH, filename)\n", + "\n", + " if not os.path.exists(full_path):\n", + " print(f\"\\u2717 {filename} not found at {full_path}\")\n", + " return\n", + "\n", + " with open(full_path, 'r') as f:\n", + " content = f.read()\n", + "\n", + " # Replace hardcoded Drive paths with generic relative ones for GitHub\n", + " # This ensures others can run it in their local folders\n", + " old_drive_root = '/content/drive/MyDrive/LRP Merge method'\n", + " new_content = content.replace(old_drive_root + '/models', './models')\n", + " new_content = new_content.replace(old_drive_root + '/datasets', './datasets')\n", + " new_content = new_content.replace(old_drive_root, '.')\n", + "\n", + " if new_content != content:\n", + " with open(full_path, 'w') as f:\n", + " f.write(new_content)\n", + " print(f\"\\u2705 {filename} (in Drive) has been updated for GitHub portability.\")\n", + " else:\n", + " print(f\"- {filename} is already portable or clean.\")\n", + "\n", + "# Execute on your core scripts\n", + "scripts_to_fix = [\"lrp_merge_pipeline.py\", \"finetune_fakenews.py\", \"lrp_computer.py\"]\n", + "for script in scripts_to_fix:\n", + " make_script_portable(script)" + ], + "execution_count": 18, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "✅ lrp_merge_pipeline.py (in Drive) has been updated for GitHub portability.\n", + "- finetune_fakenews.py is already portable or clean.\n", + "- lrp_computer.py is already portable or clean.\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "## Download the base model" + ], + "metadata": { + "id": "y2b9BH1VgWBC" + } + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 136, + "referenced_widgets": [ + "39cd0423589d4e30b97bc66781f355ca", + "43b4885dd82d410fb6bd848c8132665e", + "2813588097ee4a6d9e232f4ace8c9627", + "b7056c9990874958ae323c3193d785ec", + "f851195eb3d847e780fb0db248a8782b", + "6b2252d926e54bb6acaf04e1cc5628d4", + "173ee7ac871b40048cffa93911762ab9", + "04f102d8728649fabe4ea26808fe172e", + "a0378f8a7a9145bbbdd864e2b5faafe4", + "d03be308d05d420dbd6572415fc0dab6", + "9d8d202635384314b34e783ca92a80df", + "cf4d7a1b49e44205a3f0dd081058494c", + "2b653b43913e44c6a62e6b708a29cda1", + "1a5ade0e8a8942848631a51670c199e8", + "3acbc10192c347ddbecd049ebb89c733", + "5437721839cf4b938ab369a56709ff30", + "b0c54005c7af4cb7b64ed3610e3b5664", + "0968426d097d4885ad5b48945043a5dc", + "49e422c375a54818bc8a732e0a00399d", + "35481a02639345fb800614e0106f4086", + "73a6788def414bd9b25bef93c86fed9a", + "722e2f73d2854e5c914ada1062f6689a" + ] + }, + "id": "a29f5f3f", + "outputId": "003edd83-9ed1-44e9-9d45-75b42c139112" + }, + "source": [ + "import os\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "\n", + "LRP_PATH = \"/content/drive/MyDrive/LRP Merge method\" # Ensure LRP_PATH is defined\n", + "base_model_hf_id = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n", + "local_base_model_path = os.path.join(LRP_PATH, \"base_model\")\n", + "\n", + "# Ensure the directory exists\n", + "os.makedirs(local_base_model_path, exist_ok=True)\n", + "\n", + "print(f\"Loading base model from Hugging Face: {base_model_hf_id}\")\n", + "model = AutoModelForCausalLM.from_pretrained(base_model_hf_id)\n", + "tokenizer = AutoTokenizer.from_pretrained(base_model_hf_id)\n", + "\n", + "print(f\"Saving base model to local path: {local_base_model_path}\")\n", + "model.save_pretrained(local_base_model_path)\n", + "tokenizer.save_pretrained(local_base_model_path)\n", + "\n", + "print(\"Base model and tokenizer saved to Drive successfully!\")" + ], + "execution_count": 9, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Loading base model from Hugging Face: TinyLlama/TinyLlama-1.1B-Chat-v1.0\n" + ] + }, + { + "output_type": "display_data", + "data": { + "text/plain": [ + "Loading weights: 0%| | 0/201 [00:00=2.0.0 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (2.10.0+cu128)\n", + "Requirement already satisfied: tqdm==4.67.1 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (4.67.1)\n", + "Requirement already satisfied: click==8.2.1 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (8.2.1)\n", + "Requirement already satisfied: safetensors~=0.5.2 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (0.5.3)\n", + "Requirement already satisfied: accelerate~=1.6.0 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (1.6.0)\n", + "Requirement already satisfied: pydantic~=2.10.6 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (2.10.6)\n", + "Requirement already satisfied: immutables==0.21 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (0.21)\n", + "Requirement already satisfied: transformers>=4.45.2 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (5.0.0)\n", + "Requirement already satisfied: tokenizers>=0.20.1 in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (0.22.2)\n", + "Requirement already satisfied: huggingface_hub in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (1.7.1)\n", + "Requirement already satisfied: peft in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (0.18.1)\n", + "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (4.15.0)\n", + "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (0.2.1)\n", + "Requirement already satisfied: protobuf in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (5.29.6)\n", + "Requirement already satisfied: scipy in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (1.16.3)\n", + "Requirement already satisfied: datasets in /usr/local/lib/python3.12/dist-packages (from mergekit==0.1.4) (4.0.0)\n", + "Requirement already satisfied: numpy<3.0.0,>=1.17 in /usr/local/lib/python3.12/dist-packages (from accelerate~=1.6.0->mergekit==0.1.4) (2.0.2)\n", + "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.12/dist-packages (from accelerate~=1.6.0->mergekit==0.1.4) (26.0)\n", + "Requirement already satisfied: psutil in /usr/local/lib/python3.12/dist-packages (from accelerate~=1.6.0->mergekit==0.1.4) (5.9.5)\n", + "Requirement already satisfied: pyyaml in /usr/local/lib/python3.12/dist-packages (from accelerate~=1.6.0->mergekit==0.1.4) (6.0.3)\n", + "Requirement already satisfied: filelock>=3.10.0 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub->mergekit==0.1.4) (3.25.2)\n", + "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub->mergekit==0.1.4) (2025.3.0)\n", + "Requirement already satisfied: hf-xet<2.0.0,>=1.4.2 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub->mergekit==0.1.4) (1.4.2)\n", + "Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub->mergekit==0.1.4) (0.28.1)\n", + "Requirement already satisfied: typer in /usr/local/lib/python3.12/dist-packages (from huggingface_hub->mergekit==0.1.4) (0.24.1)\n", + "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.12/dist-packages (from pydantic~=2.10.6->mergekit==0.1.4) (0.7.0)\n", + "Requirement already satisfied: pydantic-core==2.27.2 in /usr/local/lib/python3.12/dist-packages (from pydantic~=2.10.6->mergekit==0.1.4) (2.27.2)\n", + "Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (75.2.0)\n", + "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (1.14.0)\n", + "Requirement already satisfied: networkx>=2.5.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (3.6.1)\n", + "Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (3.1.6)\n", + "Requirement already satisfied: cuda-bindings==12.9.4 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.9.4)\n", + "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.8.93)\n", + "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.8.90)\n", + "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.8.90)\n", + "Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (9.10.2.21)\n", + "Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.8.4.1)\n", + "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (11.3.3.83)\n", + "Requirement already satisfied: nvidia-curand-cu12==10.3.9.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (10.3.9.90)\n", + "Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (11.7.3.90)\n", + "Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.5.8.93)\n", + "Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (0.7.1)\n", + "Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (2.27.5)\n", + "Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (3.4.5)\n", + "Requirement already satisfied: nvidia-nvtx-cu12==12.8.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.8.90)\n", + "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (12.8.93)\n", + "Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (1.13.1.3)\n", + "Requirement already satisfied: triton==3.6.0 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0.0->mergekit==0.1.4) (3.6.0)\n", + "Requirement already satisfied: cuda-pathfinder~=1.1 in /usr/local/lib/python3.12/dist-packages (from cuda-bindings==12.9.4->torch>=2.0.0->mergekit==0.1.4) (1.4.3)\n", + "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.12/dist-packages (from transformers>=4.45.2->mergekit==0.1.4) (2025.11.3)\n", + "Requirement already satisfied: typer-slim in /usr/local/lib/python3.12/dist-packages (from transformers>=4.45.2->mergekit==0.1.4) (0.24.0)\n", + "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.12/dist-packages (from datasets->mergekit==0.1.4) (18.1.0)\n", + "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.12/dist-packages (from datasets->mergekit==0.1.4) (0.3.8)\n", + "Requirement already satisfied: pandas in /usr/local/lib/python3.12/dist-packages (from datasets->mergekit==0.1.4) (2.2.2)\n", + "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.12/dist-packages (from datasets->mergekit==0.1.4) (2.32.4)\n", + "Requirement already satisfied: xxhash in /usr/local/lib/python3.12/dist-packages (from datasets->mergekit==0.1.4) (3.6.0)\n", + "Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.12/dist-packages (from datasets->mergekit==0.1.4) (0.70.16)\n", + "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.12/dist-packages (from fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (3.13.3)\n", + "Requirement already satisfied: anyio in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface_hub->mergekit==0.1.4) (4.12.1)\n", + "Requirement already satisfied: certifi in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface_hub->mergekit==0.1.4) (2026.2.25)\n", + "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface_hub->mergekit==0.1.4) (1.0.9)\n", + "Requirement already satisfied: idna in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface_hub->mergekit==0.1.4) (3.11)\n", + "Requirement already satisfied: h11>=0.16 in /usr/local/lib/python3.12/dist-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface_hub->mergekit==0.1.4) (0.16.0)\n", + "Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets->mergekit==0.1.4) (3.4.6)\n", + "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets->mergekit==0.1.4) (2.5.0)\n", + "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch>=2.0.0->mergekit==0.1.4) (1.3.0)\n", + "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch>=2.0.0->mergekit==0.1.4) (3.0.3)\n", + "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets->mergekit==0.1.4) (2.9.0.post0)\n", + "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets->mergekit==0.1.4) (2025.2)\n", + "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets->mergekit==0.1.4) (2025.3)\n", + "Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface_hub->mergekit==0.1.4) (1.5.4)\n", + "Requirement already satisfied: rich>=12.3.0 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface_hub->mergekit==0.1.4) (13.9.4)\n", + "Requirement already satisfied: annotated-doc>=0.0.2 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface_hub->mergekit==0.1.4) (0.0.4)\n", + "Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (2.6.1)\n", + "Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (1.4.0)\n", + "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (25.4.0)\n", + "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (1.8.0)\n", + "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (6.7.1)\n", + "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (0.4.1)\n", + "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets->mergekit==0.1.4) (1.23.0)\n", + "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas->datasets->mergekit==0.1.4) (1.17.0)\n", + "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.12/dist-packages (from rich>=12.3.0->typer->huggingface_hub->mergekit==0.1.4) (4.0.0)\n", + "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.12/dist-packages (from rich>=12.3.0->typer->huggingface_hub->mergekit==0.1.4) (2.19.2)\n", + "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.12/dist-packages (from markdown-it-py>=2.2.0->rich>=12.3.0->typer->huggingface_hub->mergekit==0.1.4) (0.1.2)\n", + "Building wheels for collected packages: mergekit\n", + " Building editable for mergekit (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", + " Created wheel for mergekit: filename=mergekit-0.1.4-0.editable-py3-none-any.whl size=13739 sha256=f31b6550c62ceb245bb585adcad39ca2e971e5417c687335412eb746279d62c8\n", + " Stored in directory: /tmp/pip-ephem-wheel-cache-vpbngun1/wheels/02/0c/3b/52585b782f267fb895caf8fddcee3dfacd3df6743ef269b5a3\n", + "Successfully built mergekit\n", + "Installing collected packages: mergekit\n", + " Attempting uninstall: mergekit\n", + " Found existing installation: mergekit 0.1.4\n", + " Uninstalling mergekit-0.1.4:\n", + " Successfully uninstalled mergekit-0.1.4\n", + "Successfully installed mergekit-0.1.4\n", + "✓ Custom mergekit_repo installed.\n", + "Checking available merge methods in mergekit after custom installation...\n", + "Note: A full check of merge methods might cause a SystemExit if the custom method isn't fully integrated yet. Proceeding with script fixes.\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5-ItEziCfY5X" + }, + "source": [ + "## Test the Merged Model" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "2b0e07fc", + "outputId": "6938fc20-daeb-4bee-cf91-cf73a3be26bd" + }, + "source": [ + "\n", + "# It uses CPU mode to ensure stability across different environment configurations\n", + "import os\n", + "import shutil\n", + "import subprocess\n", + "\n", + "def run_final_merge():\n", + " config_path = \"lrp_config.yaml\"\n", + " output_dir = os.path.join(LRP_PATH, \"models/merged-model\")\n", + " os.makedirs(output_dir, exist_ok=True)\n", + "\n", + " cmd = [\n", + " \"mergekit-yaml\",\n", + " config_path,\n", + " output_dir,\n", + " \"--copy-tokenizer\",\n", + " \"--allow-crimes\",\n", + " \"--lazy-unpickle\"\n", + " ]\n", + "\n", + " print(f\"Executing: {' '.join(cmd)}\")\n", + " process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n", + " for line in process.stdout: print(line, end=\"\")\n", + " process.wait()\n", + " if process.returncode == 0: print(\"\\n✅ LRP MERGE SUCCESSFUL\")\n", + "\n", + "run_final_merge()" + ], + "execution_count": 76, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "\n", + "Step 3: Running merge (LRP-Merge / CPU Mode)...\n", + "\n", + "Running command: /usr/local/bin/mergekit-yaml lrp_config.yaml /content/drive/MyDrive/LRP Merge method/models/merged-model --copy-tokenizer --allow-crimes --lazy-unpickle\n", + "WARNING:torchao.kernel.intmm:Warning: Detected no triton, on systems without Triton certain kernels will not work\n", + "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", + "\n", + "Warmup loader cache: 0%| | 0/3 [00:00 +``` + +### Tips for Success + +1. **Start small:** Run with `MAX_SAMPLES=500` first to test +2. **Monitor GPU:** Click the RAM/Disk indicator to see GPU usage +3. **Save frequently:** The notebook saves to Drive automatically +4. **Don't close browser:** Keep the Colab tab open during training +5. **Use Chrome:** Works best with Google Chrome browser + +### Advanced: Training Multiple Models + +To train multiple variations, duplicate the training cells: + +```python +# Model A with different learning rate +!python finetune_fakenews.py \ + --dataset datasets/global_train.csv \ + --output models/global-v2 \ + --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ + --epochs 5 \ + --lr 5e-5 # Different learning rate + +# Model B with more data +!python finetune_fakenews.py \ + --dataset datasets/global_train_large.csv \ + --output models/global-v3 \ + --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ + --epochs 5 \ + --max-samples 5000 +``` + +### Next Steps After Training + +1. Download the merged model +2. Use it locally with: +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model = AutoModelForCausalLM.from_pretrained("path/to/merged-model") +tokenizer = AutoTokenizer.from_pretrained("path/to/merged-model") +``` +3. Or deploy it to Hugging Face Hub + +### Questions? + +- Check the notebook's error messages +-- Look at the output logs in each cell +- Make sure your Drive is properly mounted +- Verify GPU is enabled in Runtime settings + +--- + +**Remember:** Colab sessions are temporary. Always save your results to Google Drive before closing! + ## Citation If you find `mergekit` useful in your research, please consider citing the [paper](https://aclanthology.org/2024.emnlp-industry.36/): diff --git a/examples/lrp.yml b/examples/lrp.yml new file mode 100644 index 00000000..4772afa5 --- /dev/null +++ b/examples/lrp.yml @@ -0,0 +1,12 @@ +merge_method: lrp +base_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +parameters: + density: 0.7 +models: + - model: psmathur/orca_mini_v3_13b + parameters: + weight: 1.0 + - model: garage-bAInd/Platypus2-13B + parameters: + weight: 1.0 +dtype: float16 diff --git a/finetune_fakenews.py b/finetune_fakenews.py new file mode 100644 index 00000000..69594620 --- /dev/null +++ b/finetune_fakenews.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +""" +Fine-tune models for fake news detection. +This script fine-tunes models on fake news datasets that can then be merged. +Optimized for CPU-only training environments with small models. + +Recommended small models for CPU training: +- gpt2 (124M params) - Fastest +- gpt2-medium (355M params) - Good balance +- TinyLlama/TinyLlama-1.1B-Chat-v1.0 (1.1B params) - Requires more memory +""" + +import argparse +import logging +import os +import sys +from pathlib import Path +from typing import Dict, List, Optional + +import pandas as pd +import torch +from torch.utils.data import Dataset +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, +) + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + + +class FakeNewsDataset(Dataset): + """Dataset for fake news classification.""" + + def __init__( + self, texts: List[str], labels: List[str], tokenizer, max_length: int = 128 + ): + self.texts = texts + self.labels = labels + self.tokenizer = tokenizer + self.max_length = max_length + + # Create prompts + self.prompts = [ + self._create_prompt(text, label) for text, label in zip(texts, labels) + ] + + def _create_prompt(self, text: str, label: str) -> str: + """Create a classification prompt.""" + return f"""Classify the following news article as FAKE or REAL. + +Article: {text} + +Classification: {label}""" + + def __len__(self): + return len(self.prompts) + + def __getitem__(self, idx) -> Dict[str, torch.Tensor]: + prompt = self.prompts[idx] + + encoding = self.tokenizer( + prompt, + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt", + ) + + return { + "input_ids": encoding["input_ids"].flatten(), + "attention_mask": encoding["attention_mask"].flatten(), + "labels": encoding["input_ids"].flatten(), # For causal LM + } + + +def load_dataset(dataset_path: str, max_samples: Optional[int] = None) -> tuple: + """Load dataset from CSV.""" + logger.info(f"Loading dataset from {dataset_path}...") + + df = pd.read_csv(dataset_path) + + # Handle different column names + text_column = None + label_column = None + + for possible_text in ["text", "content", "statement", "article", "title"]: + if possible_text in df.columns: + text_column = possible_text + break + + for possible_label in ["label", "category", "truth", "class"]: + if possible_label in df.columns: + label_column = possible_label + break + + if text_column is None or label_column is None: + logger.error( + f"Could not find text/label columns. Available: {list(df.columns)}" + ) + sys.exit(1) + + texts = df[text_column].tolist() + labels = df[label_column].tolist() + + # Normalize labels + labels = [str(l).upper().strip() for l in labels] + labels = [ + "FAKE" if l in ["FAKE", "0", "FALSE", "FALSE", "F"] else "REAL" for l in labels + ] + + if max_samples: + texts = texts[:max_samples] + labels = labels[:max_samples] + + logger.info(f"Loaded {len(texts)} samples") + logger.info(f" FAKE: {labels.count('FAKE')}") + logger.info(f" REAL: {labels.count('REAL')}") + + return texts, labels + + +def setup_lora(model, r: int = 8, alpha: int = 32, dropout: float = 0.1): + """Setup LoRA for efficient fine-tuning.""" + try: + from peft import LoraConfig, TaskType, get_peft_model + + config = LoraConfig( + r=r, + lora_alpha=alpha, + target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], + lora_dropout=dropout, + bias="none", + task_type=TaskType.CAUSAL_LM, + ) + + model = get_peft_model(model, config) + model.print_trainable_parameters() + return model + + except ImportError: + logger.error("PEFT not installed. Run: pip install peft") + sys.exit(1) + + +def fine_tune_model( + dataset_path: str, + output_dir: str, + model_name: str = "gpt2", # Default to small model for CPU + epochs: int = 3, + batch_size: int = 1, + learning_rate: float = 2e-5, + use_lora: bool = True, + lora_r: int = 16, + max_samples: Optional[int] = None, + use_8bit: bool = False, + use_4bit: bool = False, + max_length: int = 128, +): + """Fine-tune model on fake news detection. + + Args: + use_8bit: Use 8-bit quantization (saves memory, slower on CPU) + use_4bit: Use 4-bit quantization (saves more memory, slower on CPU) + max_length: Maximum sequence length (lower = less memory) + """ + + logger.info(f"Starting fine-tuning...") + logger.info(f" Model: {model_name}") + logger.info(f" Dataset: {dataset_path}") + logger.info(f" Output: {output_dir}") + logger.info(f" Epochs: {epochs}") + logger.info(f" LoRA: {use_lora}") + logger.info(f" 8-bit: {use_8bit}") + logger.info(f" 4-bit: {use_4bit}") + + # Load data + texts, labels = load_dataset(dataset_path, max_samples) + + # Load tokenizer and model + logger.info("Loading tokenizer and model...") + tokenizer = AutoTokenizer.from_pretrained(model_name) + tokenizer.pad_token = tokenizer.eos_token + + # CPU/GPU compatibility logic + has_cuda = torch.cuda.is_available() + logger.info(f" CUDA available: {has_cuda}") + + # Model loading kwargs + model_kwargs = { + "low_cpu_mem_usage": True, + } + + # Handle dtype - CPU doesn't support float16 well + if has_cuda: + model_kwargs["torch_dtype"] = torch.float16 + model_kwargs["device_map"] = "auto" + else: + # CPU training - use float32 and explicit device + model_kwargs["torch_dtype"] = torch.float32 + # device_map not recommended for CPU-only + logger.info("Using CPU with float32 precision") + + # Quantization for memory saving + if use_8bit and has_cuda: + model_kwargs["load_in_8bit"] = True + logger.info("Loading model in 8-bit mode") + elif use_4bit and has_cuda: + model_kwargs["load_in_4bit"] = True + logger.info("Loading model in 4-bit mode") + elif (use_8bit or use_4bit) and not has_cuda: + logger.warning( + "8-bit/4-bit quantization requires CUDA. Using full precision on CPU." + ) + + try: + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + except Exception as e: + logger.error(f"Failed to load model: {e}") + logger.info("Trying fallback with trust_remote_code=True...") + model_kwargs["trust_remote_code"] = True + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + + # Move to CPU explicitly if not using device_map + if not has_cuda: + model = model.to("cpu") + + # Enable gradient checkpointing for memory efficiency (if supported) + if hasattr(model, "gradient_checkpointing_enable"): + model.gradient_checkpointing_enable() + logger.info("Gradient checkpointing enabled") + + # Apply LoRA if requested + if use_lora: + logger.info(f"Applying LoRA (r={lora_r})...") + model = setup_lora(model, r=lora_r) + + # Create dataset + logger.info("Creating dataset...") + dataset = FakeNewsDataset(texts, labels, tokenizer, max_length=max_length) + + # Training arguments - optimized for CPU + training_args = TrainingArguments( + output_dir=output_dir, + num_train_epochs=epochs, + per_device_train_batch_size=batch_size, + gradient_accumulation_steps=8 if not has_cuda else 4, # Higher for CPU + per_device_eval_batch_size=batch_size, + learning_rate=learning_rate, + warmup_steps=100, + weight_decay=0.01, + logging_steps=10, + save_steps=500, + save_total_limit=2, + # FP16 only on GPU + fp16=has_cuda and not (use_8bit or use_4bit), + # BF16 not on CPU + bf16=False, + dataloader_pin_memory=has_cuda, + report_to="none", + remove_unused_columns=False, + use_cpu=not has_cuda, + # Disable some optimizations that don't work well on CPU + dataloader_num_workers=0, # Avoid multiprocessing issues on Windows + disable_tqdm=False, + ) + + # Data collator + data_collator = DataCollatorForLanguageModeling( + tokenizer=tokenizer, + mlm=False, + ) + + # Trainer + logger.info("Initializing trainer...") + trainer = Trainer( + model=model, + args=training_args, + train_dataset=dataset, + data_collator=data_collator, + ) + + # Train + logger.info("Starting training...") + trainer.train() + + # Save + logger.info(f"Saving model to {output_dir}...") + trainer.save_model(output_dir) + tokenizer.save_pretrained(output_dir) + + logger.info("Fine-tuning complete!") + return output_dir + + +def main(): + parser = argparse.ArgumentParser( + description="Fine-tune TinyLlama for fake news detection" + ) + + parser.add_argument("--dataset", required=True, help="Path to training dataset CSV") + + parser.add_argument( + "--output", required=True, help="Output directory for fine-tuned model" + ) + + parser.add_argument( + "--model", + default="gpt2", + help="Base model to fine-tune (default: gpt2 for CPU, ~124M params)", + ) + + parser.add_argument( + "--epochs", type=int, default=3, help="Number of training epochs" + ) + + parser.add_argument( + "--batch-size", + type=int, + default=1, # Default to 1 for CPU stability + help="Training batch size", + ) + + parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate") + + parser.add_argument( + "--use-lora", action="store_true", help="Use LoRA for efficient fine-tuning" + ) + + parser.add_argument("--lora-r", type=int, default=16, help="LoRA rank") + + parser.add_argument( + "--max-samples", type=int, help="Limit training samples (for testing)" + ) + + parser.add_argument( + "--use-8bit", + action="store_true", + help="Use 8-bit quantization (requires CUDA, saves memory)", + ) + + parser.add_argument( + "--use-4bit", + action="store_true", + help="Use 4-bit quantization (requires CUDA, saves more memory)", + ) + + parser.add_argument( + "--max-length", + type=int, + default=128, + help="Maximum sequence length (default: 128, lower = less memory)", + ) + + args = parser.parse_args() + + # Validate dataset exists + if not os.path.exists(args.dataset): + logger.error(f"Dataset not found: {args.dataset}") + sys.exit(1) + + # Create output directory + os.makedirs(args.output, exist_ok=True) + + # Run fine-tuning + fine_tune_model( + dataset_path=args.dataset, + output_dir=args.output, + model_name=args.model, + epochs=args.epochs, + batch_size=args.batch_size, + learning_rate=args.lr, + use_lora=args.use_lora, + lora_r=args.lora_r, + max_samples=args.max_samples, + use_8bit=args.use_8bit, + use_4bit=args.use_4bit, + max_length=args.max_length, + ) + + print("\n" + "=" * 60) + print("Fine-tuned model saved to:", args.output) + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/git b/git new file mode 100644 index 00000000..e69de29b diff --git a/lrp_computer.py b/lrp_computer.py new file mode 100644 index 00000000..dbfc84e4 --- /dev/null +++ b/lrp_computer.py @@ -0,0 +1,144 @@ +import json +import logging +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import torch +from tqdm import tqdm +from transformers import AutoModelForCausalLM, AutoTokenizer + +logger = logging.getLogger(__name__) + + +class LRPComputer: + """ + Supercharged LRP Computer with Multimodal support and Tied-Tensor handling. + """ + + def __init__(self, config): + self.config = config + self.model = None + self.tokenizer = None + self.relevance_scores: Dict[str, torch.Tensor] = {} + + def load_model(self) -> None: + """Load the model and tokenizer.""" + logger.info(f"Loading model from {self.config.model_path}...") + self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path) + self.model = AutoModelForCausalLM.from_pretrained( + self.config.model_path, + torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, + device_map="auto" if torch.cuda.is_available() else None, + trust_remote_code=True, + ) + self.model.eval() + + def compute_layer_relevance( + self, + layer_input: torch.Tensor, + layer_output: torch.Tensor, + relevance_output: torch.Tensor, + layer_name: str, + ) -> torch.Tensor: + """Compute relevance using autograd for the chosen rule.""" + # Simple epsilon rule implementation using gradients as a proxy for LRP + # (Compatible with any architecture) + layer_input.requires_grad_(True) + with torch.enable_grad(): + # Local forward pass + # We use a simple dot product as a surrogate for relevance propagation + surrogate = (layer_output * relevance_output).sum() + surrogate.backward(retain_graph=True) + + relevance_input = layer_input.grad * layer_input + return relevance_input.detach() + + def compute_all_relevance_scores( + self, prompts: List[str] + ) -> Dict[str, torch.Tensor]: + """Compute LRP scores across all parameters.""" + if self.model is None: + self.load_model() + + importance_scores = {} + activations = {} + + # Register hooks to capture activations + def get_hook(name): + def hook(module, input, output): + activations[name] = input[0].detach() + + return hook + + hooks = [] + for name, module in self.model.named_modules(): + hooks.append(module.register_forward_hook(get_hook(name))) + + # Forward pass to collect activations + inputs = self.tokenizer(prompts, return_tensors="pt", padding=True).to( + self.model.device + ) + outputs = self.model(**inputs, output_hidden_states=True) + + # Remove hooks + for h in hooks: + h.remove() + + # Get last token logits for relevance signal + last_logits = outputs.logits[:, -1, :].abs().mean(dim=0) + relevance_vocab = last_logits # (vocab_size,) + + # Project relevance from vocab_size -> hidden_size via lm_head weight. + hidden_relevance = None + for name, param in self.model.named_parameters(): + # MULTIMODAL FIX: Support 'language_model.lm_head' and similar prefixes + if "lm_head" in name and param.dim() == 2: + if param.shape[0] == relevance_vocab.shape[0]: + hidden_relevance = ( + (relevance_vocab.unsqueeze(0) @ param.float()).squeeze(0).abs() + ) + break + + # Propagate relevance through each named parameter + for name, param in tqdm(self.model.named_parameters(), desc="Computing LRP"): + layer_name = ".".join(name.split(".")[:-1]) + act = activations.get(layer_name) + if act is None: + importance_scores[name] = param.data.abs().cpu().clone() + continue + + # Flatten act to (tokens, features) + act_flat = act.reshape(-1, act.shape[-1]) if act.dim() > 1 else act + feat_dim = act_flat.shape[-1] + + if hidden_relevance is not None and hidden_relevance.numel() == feat_dim: + rel_signal = hidden_relevance.expand_as(act_flat) + else: + rel_signal = act_flat.abs().mean() * torch.ones_like(act_flat) + + # Assign score based on magnitude x activation + importance_scores[name] = ( + param.data.abs().cpu() * act_flat.abs().mean(dim=0).cpu() + ).clone() + + self.relevance_scores = importance_scores + return importance_scores + + def save_relevance_scores(self) -> None: + """Save computed relevance scores to disk.""" + output_path = Path(self.config.output_path) + output_path.mkdir(parents=True, exist_ok=True) + + # TIED-TENSOR FIX: Clone scores before saving to prevent share storage errors + scores_path = output_path / "lrp_scores.pt" + save_dict = {k: v.clone() for k, v in self.relevance_scores.items()} + torch.save(save_dict, scores_path) + logger.info(f"Saved LRP scores to {scores_path}") + + # Save metadata + metadata = { + "model_path": self.config.model_path, + "num_tensors": len(self.relevance_scores), + } + with open(output_path / "lrp_metadata.json", "w") as f: + json.dump(metadata, f, indent=2) diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py new file mode 100644 index 00000000..eda8598c --- /dev/null +++ b/lrp_merge_pipeline.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 + +import os +import shutil +import subprocess +import sys + +# Config file is auto injected + +BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # model can be user-specific + +# These placeholders will be replaced by absolute paths from the notebook +MODEL_1 = "./models/tinyllama-global-full" +MODEL_2 = "./models/tinyllama-local-full" + +OUTPUT_DIR = "./models/merged-model" + +# Toggle this if GPU available +USE_CUDA = True + +# validation + + +def validate(): + print("Step 1: Checking local model paths...") + + for path in [MODEL_1, MODEL_2]: + if os.path.exists(path): + print(f"Found: {path}") + else: + print(f" Missing model: {path}") + # We won't raise error here to see the full debug output + + +# YAML GENERATION + + +def generate_yaml(): + print("Step 2: Generating YAML...") + + yaml = f""" +merge_method: lrp + +base_model: + model: \"{BASE_MODEL}\" + +parameters: + density: 0.7 + use_lrp: true + +models: + - model: \"{MODEL_1}\" + parameters: + weight: 1.0 + + - model: \"{MODEL_2}\" + parameters: + weight: 1.0 +""" + + with open("lrp_config.yaml", "w") as f: + f.write(yaml.strip()) + + print("\nYAML Generated:\n") + print(yaml) + + +# MERGE EXECUTION + + +def run_merge(): + print("\nStep 3: Running merge...\n") + + # Use the verified path + mergekit_exec = "/usr/local/bin/mergekit-yaml" + if not os.path.exists(mergekit_exec): + mergekit_exec = shutil.which("mergekit-yaml") + + if not mergekit_exec: + raise RuntimeError("✗ mergekit-yaml not found. Install mergekit.") + + cmd = [ + mergekit_exec, + "lrp_config.yaml", + OUTPUT_DIR, + "--copy-tokenizer", + "--allow-crimes", + ] + + if USE_CUDA: + cmd.append("--cuda") + + print("Running command:") + print(" ".join(cmd), "\n") + + # Run and show output in real-time + res = subprocess.run(cmd) + + if res.returncode != 0: + raise RuntimeError(f"✗ Merge failed with code {res.returncode}") + + print("\nMerge completed successfully!") + print(f"📁 Output: {OUTPUT_DIR}") + + +# MAIN + + +def main(): + print("=== LRP MERGE PIPELINE START ===\n") + + validate() + generate_yaml() + os.makedirs(OUTPUT_DIR, exist_ok=True) + run_merge() + + print("\n✓ ALL DONE") + + +if __name__ == "__main__": + main() diff --git a/mergekit/architecture/auto.py b/mergekit/architecture/auto.py index 5515b44c..72a81205 100644 --- a/mergekit/architecture/auto.py +++ b/mergekit/architecture/auto.py @@ -157,15 +157,38 @@ def infer_architecture_info( f" {repr(prefix or 'default')} with {module_layer_counts[prefix]} layers, {len(module_templates[prefix])} templates, and {len(module_loose_weights[prefix])} loose weights" ) - def _wi(template: str, prefix: str) -> WeightInfo: + def _wi(template: str, prefix: str, num_layers: int = 1) -> WeightInfo: full_name = prefix + template - optional = (full_name.replace("${layer_index}", "0") not in in_all_models) or ( + # A template is optional if ANY of its layer instantiations are missing + # from in_all_models. Required for hybrid architectures like Qwen3.5 + # (alternating self_attn / linear_attn per layer) where layer 0 may have + # both kinds of weights but later layers may have only one. Without this, + # layer 0 satisfies the lookup, but the planner emits required LoadTensor + # for `linear_attn.norm.weight` at every layer index 0..N-1 — and a + # self_attn-only layer raises at execute time. + if "${layer_index}" in full_name: + layer_optional = any( + full_name.replace("${layer_index}", str(i)) not in in_all_models + for i in range(num_layers) + ) + else: + layer_optional = full_name not in in_all_models + + optional = layer_optional or ( + tied_keys is not None + and any( + re.search(pat, full_name.replace("${layer_index}", "0")) + for pat in tied_keys + ) + ) + + is_embed = (full_name in embed_names) or ( tied_keys is not None - and any(re.search(pat, full_name) for pat in tied_keys) + and any( + re.search(pat, full_name.replace("${layer_index}", "0")) + for pat in tied_keys + ) ) - is_embed = (full_name in embed_names) or any( - re.search(pat, full_name) for pat in tied_keys - ) # strictly speaking you can have tied non-embedding/lm-head weights # but i've never seen it so let's not worry about it until this breaks something return WeightInfo( name=template, @@ -174,15 +197,17 @@ def _wi(template: str, prefix: str) -> WeightInfo: ) module_archs = {} - for prefix in module_prefixes: + for prefix in module_layer_counts.keys(): num_layers = module_layer_counts[prefix] module_archs[prefix or "default"] = JsonModuleArchitecture( definition=JsonModuleArchDef( model_type="", architectures=[], - pre_weights=[_wi(t, "") for t in module_loose_weights[prefix]], + pre_weights=[ + _wi(t, "", num_layers) for t in module_loose_weights[prefix] + ], layer_templates=JsonLayerTemplates( - weights=[_wi(t, "") for t in module_templates[prefix]] + weights=[_wi(t, "", num_layers) for t in module_templates[prefix]] ), post_weights=[], num_layers_config_key=None, diff --git a/mergekit/architecture/base.py b/mergekit/architecture/base.py index 738fb8f6..aad01ce9 100644 --- a/mergekit/architecture/base.py +++ b/mergekit/architecture/base.py @@ -38,6 +38,7 @@ class WeightInfo(BaseModel, frozen=True): def _prefix_weight(weight: WeightInfo, prefix: Optional[str] = None) -> WeightInfo: if prefix is None: return weight + return WeightInfo( name=prefix + weight.name, aliases=tuple(prefix + alias for alias in weight.aliases or ()) or None, @@ -66,7 +67,7 @@ def layer_weights( ... def num_layers_config_key(self) -> str: - """Key in config that represents number of layers""" + """Key in config that represents number of layers.""" return "num_hidden_layers" def num_layers(self, config: PretrainedConfig) -> int: @@ -76,9 +77,12 @@ def num_layers(self, config: PretrainedConfig) -> int: def all_weights(self, config: PretrainedConfig) -> List[WeightInfo]: """Return all weights associated with a model.""" num_layers = self.num_layers(config) + res = list(self.pre_weights(config)) + for layer_idx in range(num_layers): res.extend(self.layer_weights(layer_idx, config)) + res.extend(self.post_weights(config)) return res @@ -133,13 +137,17 @@ class ModelArchitecture(BaseModel, frozen=True): def all_weights(self, config: PretrainedConfig) -> List[WeightInfo]: res = [] + for module in self.modules.values(): for weight_info in module.architecture.all_weights(config=config): res.append(_prefix_weight(weight_info, module.weight_prefix)) + return res -class ConfiguredModelArchitecture(BaseModel, frozen=True, arbitrary_types_allowed=True): +class ConfiguredModelArchitecture( + BaseModel, frozen=True, arbitrary_types_allowed=True +): info: ModelArchitecture config: PretrainedConfig @@ -155,4 +163,6 @@ def get_module(self, module_name: str) -> ConfiguredModuleArchitecture: ConfiguredModuleArchitecture.model_rebuild() +ModuleDefinition.model_rebuild() ConfiguredModelArchitecture.model_rebuild() +WeightInfo.model_rebuild() diff --git a/mergekit/common.py b/mergekit/common.py index c1f3de77..05e151dd 100644 --- a/mergekit/common.py +++ b/mergekit/common.py @@ -291,6 +291,12 @@ def __iter__(self): def __getitem__(self, key: T_K) -> T_V: return self.data[key] + def __contains__(self, key: Any) -> bool: + return key in self.data + + def get(self, key: T_K, default: Optional[T_V] = None) -> Optional[T_V]: + return self.data.get(key, default) + def __len__(self) -> int: return len(self.data) diff --git a/mergekit/config.py b/mergekit/config.py index fc237137..d6cf805d 100644 --- a/mergekit/config.py +++ b/mergekit/config.py @@ -10,7 +10,7 @@ from mergekit.common import ModelReference from mergekit.tokenizer.config import TokenizerConfig -ScalarOrGradient: TypeAlias = Union[float, List[float]] +ScalarOrGradient: TypeAlias = Union[float, str, List[float], List[str]] class ConditionalParameter(BaseModel): diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py new file mode 100644 index 00000000..8b231ac3 --- /dev/null +++ b/mergekit/merge_methods/lrp.py @@ -0,0 +1,210 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only + +import gc +import logging +import time +from typing import Any, Dict, List, Optional + +import torch +from typing_extensions import override + +from mergekit.architecture import WeightInfo +from mergekit.common import ImmutableMap, ModelReference +from mergekit.graph import Task +from mergekit.merge_methods.base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from mergekit.merge_methods.rectify_embed import rectify_embed_sizes +from mergekit.sparsify import build_mask + +_TASK_COUNTER = 0 +logger = logging.getLogger(__name__) + + +class LRPMergeTask(Task[torch.Tensor], frozen=True): + """ + Supercharged LRP Merge Task: + - Multimodal support (optional tensor handling) + - Turbo optimizations (Iron-Man stabilization, in-place math) + - Resource-efficient (GC management for Windows) + """ + + gather_tensors: MergeTensorInput + base_model: Optional[ModelReference] + model_weights: ImmutableMap[ModelReference, float] + density: float + weight_info: WeightInfo + lrp_scores: Optional[ImmutableMap[str, str]] = None + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: + global _TASK_COUNTER + _TASK_COUNTER += 1 + + # IRON-MAN STABILIZATION: Prevent SSD/Memory thrashing on Windows + if _TASK_COUNTER % 100 == 0: + gc.collect() + time.sleep(1.0) + + # Get base tensor + base_tensor = tensors.get(self.base_model) if self.base_model else None + + if base_tensor is None: + first_tensor = next(iter(tensors.values())) if tensors else None + if first_tensor is None: + raise ValueError("No tensors provided for merging") + base_tensor = torch.zeros_like(first_tensor) + + # Collect and rectify non-base tensors + weight_tensors = { + ref: t for ref, t in tensors.items() if ref != self.base_model + } + if not weight_tensors: + return base_tensor + + # Rectification for embedding size mismatches + refs = list(weight_tensors.keys()) + all_tensors = [base_tensor] + [weight_tensors[r] for r in refs] + rectify_embed_sizes(self.weight_info, all_tensors) + base_tensor = all_tensors[0] + weight_tensors = {r: all_tensors[i + 1] for i, r in enumerate(refs)} + + # Initialize merged deltas (using in-place additions later) + merged_deltas = torch.zeros_like(base_tensor) + + total_weight = sum(self.model_weights.values()) + if total_weight == 0: + total_weight = 1.0 + + _lrp_cache: Dict[str, Any] = {} + + # CRITICAL LAYER PROTECTION: 100% density for norms, heads, etc. + name = self.weight_info.name.lower() + is_critical = any(x in name for x in ["norm", "embed", "ln_", "head", "bias"]) + current_density = 1.0 if is_critical else self.density + + for ref, fine_tuned_weight in weight_tensors.items(): + # MULTIMODAL FIX: Handle optional tensors (dt_bias, etc.) in hybrid architectures + if fine_tuned_weight is None: + continue + + # Validate tensor shape + if fine_tuned_weight.shape != base_tensor.shape: + continue + + # Compute delta (task vector) - in-place subtraction to save memory + delta = fine_tuned_weight.sub(base_tensor) + + importance = None + ref_str = str(ref) + if self.lrp_scores is not None and ref_str in self.lrp_scores: + lrp_path = self.lrp_scores[ref_str] + if lrp_path not in _lrp_cache: + try: + if lrp_path.endswith(".safetensors"): + from safetensors.torch import load_file + + _lrp_cache[lrp_path] = load_file(lrp_path) + else: + _lrp_cache[lrp_path] = torch.load( + lrp_path, map_location="cpu" + ) + except Exception as e: + logger.warning( + f"Failed to load LRP scores from {lrp_path}: {e}" + ) + _lrp_cache[lrp_path] = {} + importance = _lrp_cache[lrp_path].get(self.weight_info.name) + if importance is not None: + importance = importance.to(delta.device) + + if importance is None: + # MULTIMODAL FIX: Treat missing importance as passthrough for this model + continue + + if importance.shape != delta.shape: + importance = delta.abs() + + # Sparsify based on importance + mask = build_mask(importance, current_density) + + # Weighted addition to merged_deltas + weight = self.model_weights.get(ref, 1.0) + normalized_weight = weight / total_weight + + # In-place accumulation to save memory + merged_deltas.add_(delta.mul_(mask), alpha=normalized_weight) + + # Clean up to keep memory footprint low + del delta, mask, importance + if _TASK_COUNTER % 10 == 0: + gc.collect() + + return base_tensor.add_(merged_deltas) + + def uses_accelerator(self) -> bool: + return True + + def group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + def priority(self) -> int: + return 0 + + +class LRPMerge(MergeMethod): + """ + Explainable Layer-wise Relevance Propagation (eX-LRP) Merge Method. + Optimized for multimodal support and high-performance execution. + """ + + @override + def name(self) -> str: + return "lrp" + + @override + def pretty_name(self) -> Optional[str]: + return "LRP Merge (Supercharged)" + + def parameters(self) -> List[ConfigParameterDef]: + return [ + ConfigParameterDef(name="density", required=False, default_value=0.7), + ] + + def tensor_parameters(self) -> List[ConfigParameterDef]: + return [ConfigParameterDef(name="weight", required=False, default_value=1.0)] + + @override + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + parameters: ImmutableMap[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + base_model: Optional[ModelReference], + lrp_scores: Optional[Dict[str, str]] = None, + **_kwargs, + ) -> Task: + model_weights = {} + for model_ref, params in tensor_parameters.items(): + if model_ref != base_model: + model_weights[model_ref] = params.get("weight", 1.0) + + density = parameters.get("density", 0.7) + if not 0 <= density <= 1: + raise ValueError(f"density must be between 0 and 1, got {density}") + + return LRPMergeTask( + gather_tensors=tensors, + base_model=base_model, + model_weights=ImmutableMap(model_weights), + density=density, + weight_info=output_weight, + lrp_scores=ImmutableMap(lrp_scores) if lrp_scores else None, + ) diff --git a/mergekit/merge_methods/registry.py b/mergekit/merge_methods/registry.py index 86d9f906..ff996065 100644 --- a/mergekit/merge_methods/registry.py +++ b/mergekit/merge_methods/registry.py @@ -12,6 +12,7 @@ from mergekit.merge_methods.karcher import KarcherMerge from mergekit.merge_methods.linear import LinearMerge from mergekit.merge_methods.model_stock import ModelStockMerge +from mergekit.merge_methods.lrp import LRPMerge from mergekit.merge_methods.nuslerp import NuSlerpMerge from mergekit.merge_methods.passthrough import PassthroughMerge from mergekit.merge_methods.slerp import SlerpMerge @@ -25,6 +26,7 @@ ModelStockMerge(), ArceeFusionMerge(), KarcherMerge(), + LRPMerge(), # generalized task arithmetic methods GeneralizedTaskArithmeticMerge( consensus_method=None, diff --git a/mergekit/sparsify.py b/mergekit/sparsify.py index b6173321..e80dbca0 100644 --- a/mergekit/sparsify.py +++ b/mergekit/sparsify.py @@ -197,3 +197,27 @@ def sparsify( ) else: raise NotImplementedError(method) + + +def build_mask(importance: torch.Tensor, density: float) -> torch.Tensor: + """Create a binary mask based on importance scores and desired density.""" + if density >= 1.0: + return torch.ones_like(importance) + if density <= 0.0: + return torch.zeros_like(importance) + + k = int(density * importance.numel()) + if k <= 0: + return torch.zeros_like(importance) + if k >= importance.numel(): + return torch.ones_like(importance) + + mask = torch.zeros_like(importance) + w = importance.abs().view(-1) + if w.device.type == "cpu": + w = w.float() + + # Use topk for performance on large tensors + _, topk_indices = torch.topk(w, k, sorted=False) + mask.view(-1)[topk_indices] = 1 + return mask