From 90a9d869de2f6a0ed835ab7e8f5280bf8fdfb63a Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sat, 25 Apr 2026 21:31:13 +0530 Subject: [PATCH 01/13] docs: Update README with eX-LRP documentation and features --- README.md | 825 ++++++++++++++++-------------------------------------- 1 file changed, 243 insertions(+), 582 deletions(-) diff --git a/README.md b/README.md index 7881816a..0bb69a50 100644 --- a/README.md +++ b/README.md @@ -1,701 +1,362 @@ -# mergekit +# eX-LRP + +### eXplainable Layer-wise Relevance Propagation for Transformer Model Merging +**Intelligent, faithful weight preservation for merging Transformer LLMs with a single backward pass** [![License: LGPL v3](https://img.shields.io/badge/License-LGPL_v3-blue.svg)](https://www.gnu.org/licenses/lgpl-3.0) [![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/arcee-ai/mergekit/pre-commit.yml?label=Tests)](https://github.com/arcee-ai/mergekit/actions/workflows/pre-commit.yml) -[![Arcee Discord](https://img.shields.io/badge/Arcee%20Discord-Arcee%20Discord?logo=discord&logoColor=white&color=5865F2)](https://discord.gg/arceeai) - -`mergekit` is a toolkit for merging pre-trained language models. `mergekit` uses an out-of-core approach to perform unreasonably elaborate merges in resource-constrained situations. Merges can be run entirely on CPU or accelerated with as little as 8 GB of VRAM. Many merging algorithms are supported, with more coming as they catch my attention. - -## LRP-Merge: Layer-wise Relevance Propagation for LLM Merging +[![Built on Mergekit](https://img.shields.io/badge/Built%20on-Mergekit-orange)](https://github.com/arcee-ai/mergekit) +[![Python 3.10+](https://img.shields.io/badge/Python-3.10%2B-blue)](https://www.python.org/) +[![PyTorch](https://img.shields.io/badge/PyTorch-2.0%2B-red)](https://pytorch.org/) -LRP-Merge is a custom model merging method built on top of [Mergekit](https://github.com/arcee-ai/mergekit). It uses **Layer-wise Relevance Propagation (LRP)** scores to identify and preserve functionally critical weights during model merging. - -### The Core Concept +--- -#### The Problem with Standard Merging -Traditional merging (linear averaging) dilutes fine-tuned knowledge because it treats all weights equally. In reality, only a small fraction of weights drive new capabilities. +## Table of Contents -#### The LRP Solution -LRP-Merge applies an XAI technique to score each weight's contribution to correct predictions: -1. **Compute Task Vector:** `δ = θ_fine_tuned - θ_base` -2. **LRP-Based Trimming:** Keep weights with highest LRP relevance scores -3. **Weighted Averaging:** Merge sparse task vectors, add back to base +- [What is eX-LRP?](#-what-is-ex-lrp) +- [How It Works](#-how-ex-lrp-works) +- [Supported Transformer Architectures](#-supported-transformer-architectures) +- [Getting Started](#-getting-started) +- [Quickstart](#-quickstart) +- [Architecture Diagram](#-architecture-diagram) +- [Merge Configuration](#-merge-configuration) +- [All Mergekit Methods](#-all-mergekit-merge-methods) +- [Running on Google Colab](#-running-on-google-colab) +- [Acknowledgements](#-acknowledgements) +- [Citation](#-citation) --- -## Contents - -- [Why Merge Models?](#why-merge-models) -- [Features](#features) -- [Installation](#installation) -- [Community & Support](#community--support) - - [Contributing](#contributing) - - [Community Tools](#community-tools) -- [Usage](#usage) -- [Merge Configuration](#merge-configuration) - - [Parameter Specification](#parameter-specification) - - [Tokenizer Configuration](#tokenizer-configuration) - - [Chat Template Configuration](#chat-template-configuration) - - [Examples](#examples) -- [Merge Methods](#merge-methods) -- [LoRA Extraction](#lora-extraction) -- [Mixture of Experts Merging](#mixture-of-experts-merging) -- [Evolutionary Merge Methods](#evolutionary-merge-methods) -- [Multi-Stage Merging (`mergekit-multi`)](#multi-stage-merging-mergekit-multi) -- [Raw PyTorch Model Merging (`mergekit-pytorch`)](#raw-pytorch-model-merging-mergekit-pytorch) -- [Tokenizer Transplantation (`mergekit-tokensurgeon`)](#tokenizer-transplantation-mergekit-tokensurgeon) -- [LRP-Merge Colab Guide](#instructions-for-running-lrp-merge-on-google-colab) -- [Citation](#citation) - -## Why Merge Models? - -Model merging is a powerful technique that allows combining the strengths of different models without the computational overhead of ensembling or the need for additional training. By operating directly in the weight space of models, merging can: - -- Combine multiple specialized models into a single versatile model -- Transfer capabilities between models without access to training data -- Find optimal trade-offs between different model behaviors -- Improve performance while maintaining inference costs -- Create new capabilities through creative model combinations - -Unlike traditional ensembling which requires running multiple models, merged models maintain the same inference cost as a single model while often achieving comparable or superior performance. - -## Features - -Key features of `mergekit` include: - -- Supports Llama, Mistral, GPT-NeoX, StableLM, and more -- Many [merge methods](#merge-methods) -- GPU or CPU execution -- Lazy loading of tensors for low memory use -- Interpolated gradients for parameter values (inspired by Gryphe's [BlockMerge_Gradient](https://github.com/Gryphe/BlockMerge_Gradient) script) -- Piecewise assembly of language models from layers ("Frankenmerging") -- [Mixture of Experts merging](#mixture-of-experts-merging) -- [LORA extraction](#lora-extraction) -- [Evolutionary merge methods](#evolutionary-merge-methods) -- [Multi-stage merging](#multi-stage-merging-mergekit-multi) for complex workflows. -- [Merging of raw PyTorch models (`mergekit-pytorch`)](#raw-pytorch-model-merging-mergekit-pytorch). - -## Installation +## ✨ What is eX-LRP? -```sh -git clone https://github.com/arcee-ai/mergekit.git -cd mergekit +**eX-LRP** is an advanced model merging method built on top of [Mergekit](https://github.com/arcee-ai/mergekit) that uses **Explainable AI (XAI)** to intelligently preserve the most functionally critical weights when merging Transformer-based Large Language Models. -pip install -e . # install the package and make scripts available -``` +Unlike standard merging methods (linear averaging, TIES, DARE) that rely on magnitude or random pruning, eX-LRP uses **AttnLRP (Attention-Aware Layer-wise Relevance Propagation)** — a backpropagation-based technique — to precisely determine how much each individual weight in the Transformer contributes to the model's actual predictions. -If the above fails with the error of: - -``` -ERROR: File "setup.py" or "setup.cfg" not found. Directory cannot be installed in editable mode: -(A "pyproject.toml" file was found, but editable mode currently requires a setuptools-based build.) -``` +### System Architecture +![alt text](image.png) -You may need to upgrade pip to > 21.3 with the command `python3 -m pip install --upgrade pip`. +### The Problem -## Community & Support +| Standard Merging | eX-LRP Merging | +|:-:|:-:| +| Treats all weights equally or uses magnitude as a proxy | Scores every weight by its **true functional contribution** | +| Blindly averages or randomly prunes task vectors | Keeps only **XAI-verified critical weights** | +| Often causes "catastrophic forgetting" | **Preserves specialized knowledge** with mathematical guarantees | -- **Issues**: [GitHub Issues](https://github.com/arcee-ai/mergekit/issues) -- **Discussions**: [Arcee Discord](https://discord.gg/arceeai) +### Why eX-LRP? -### Contributing +> *"Magnitude does not equal importance."* -We welcome contributions to `mergekit`! If you have ideas for new merge methods, features, or other improvements, please check out our [contributing guide](CONTRIBUTING.md) for details on how to get started. +A weight with a large magnitude might be irrelevant to the task, while a small weight deep in an Attention head might be the critical link that makes the model perform well. eX-LRP solves this by **tracing the prediction backward** through the entire Transformer graph — through Self-Attention, MLPs, LayerNorms, and Residual connections — to find the weights that truly matter. -### Community Tools +--- -- **[FrankensteinAI](https://frankenstein-ai.com/)**: For those who prefer a browser-based experience without local setup or hardware wrangling, the team at FrankensteinAI has built a hosted platform powered by `mergekit`. Also features a community gallery and leaderboard for sharing and comparing merged models. +## 🔬 How eX-LRP Works -## Usage +eX-LRP implements a **Gradient × Input** formulation of Layer-wise Relevance Propagation, inspired by the [AttnLRP paper (ICML 2024)](https://proceedings.mlr.press/v235/achtibat24a.html). This enables faithful relevance attribution in a single backward pass through the entire Transformer. -The script `mergekit-yaml` is the main entry point for `mergekit`. It takes a YAML configuration file and an output path, like so: +### The Pipeline -```sh -mergekit-yaml path/to/your/config.yml ./output-model-directory [--cuda] [--lazy-unpickle] [--allow-crimes] [... other options] ``` +1. TASK VECTOR EXTRACTION + δ = θ_fine_tuned - θ_base + (Isolate what the model learned) -This will run the merge and write your merged model to `./output-model-directory`. +2. FORWARD PASS WITH HOOKS + Run calibration prompts → Capture activations at every + Self-Attention, MLP, and LayerNorm block -For more information on the arguments accepted by `mergekit-yaml` run the command `mergekit-yaml --help`. +3. RELEVANCE SEEDING + Extract logits of the predicted token → Set initial R_out -### Uploading to Huggingface +4. AttnLRP BACKWARD PASS (The Core Innovation) + Propagate relevance backward using PyTorch autograd: + ┌─────────────────────────────────────────────────┐ + │ • Linear Layers: Gradient × Input rule │ + │ • LayerNorm/RMSNorm: Gradient × Input rule │ + │ • Residual Connections: Proportional splitting │ + │ • Non-linearities (SiLU/GELU): Autograd-based │ + └─────────────────────────────────────────────────┘ + Output: Importance Score Tensor for every weight -When you have a merged model you're happy with, you may want to share it on the Hugging Face Hub. `mergekit` generates a `README.md` for your merge with some basic information for a model card. You can edit it to include more details about your merge, like giving it a good name or explaining what it's good at; rewrite it entirely; or use the generated `README.md` as-is. It is also possible to edit your `README.md` online once it has been uploaded to the Hub. +5. SPARSIFICATION + Apply density threshold → Generate binary relevance mask -Once you're happy with your model card and merged model, you can upload it to the Hugging Face Hub using the [huggingface_hub](https://huggingface.co/docs/huggingface_hub/index) Python library. - -```sh -# log in to huggingface with an access token (must have write permission) -huggingface-cli login -# upload your model -huggingface-cli upload your_hf_username/my-cool-model ./output-model-directory . +6. MERGE + Masked task vector × weight → Add back to base model ``` -The [documentation](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-upload) for `huggingface_hub` goes into more detail about other options for uploading. +### Key Technical Features -## Merge Configuration +| Feature | Description | +|---|---| +| 🧠 **AttnLRP Rules** | Uses `Gradient × Input` via PyTorch `autograd` for mathematically faithful relevance through any differentiable operation | +| 🔀 **Residual Split Logic** | Correctly divides relevance between skip-connections and Transformer blocks based on relative activation magnitudes | +| 📐 **Numerical Stabilization** | Epsilon-stabilized denominators prevent division-by-zero in deep networks | +| 🎯 **Prediction-Seeded** | Relevance starts from the model's actual top-predicted token, not arbitrary uniform initialization | +| ⚡ **Single Backward Pass** | Entire importance computation runs in one backward pass — no iterative probing | -Merge configurations are YAML documents specifying the operations to perform in order to produce your merged model. -Below are the primary elements of a configuration file: - -- `merge_method`: Specifies the method to use for merging models. See [Merge Methods](#merge-methods) for a list. -- `slices`: Defines slices of layers from different models to be used. This field is mutually exclusive with `models`. -- `models`: Defines entire models to be used for merging. This field is mutually exclusive with `slices`. -- `base_model`: Specifies the base model used in some merging methods. -- `parameters`: Holds various parameters such as weights and densities, which can also be specified at different levels of the configuration. -- `dtype`: Specifies the data type used for the merging operation. -- `tokenizer` or `tokenizer_source`: Determines how to construct a tokenizer for the merged model. -- `chat_template`: Specifies a chat template for the merged model. - -### Parameter Specification - -Parameters are flexible and can be set with varying precedence. They can be specified conditionally using tensor name filters, which allows finer control such as differentiating between attention heads and fully connected layers. +--- -Parameters can be specified as: +## 🧩 Supported Transformer Architectures -- **Scalars**: Single floating-point values. -- **Gradients**: List of floating-point values, specifying an interpolated gradient. +eX-LRP is designed for **Transformer-based** language models. The current implementation targets decoder-only causal LMs, which represent the vast majority of models merged with `mergekit`. -The parameters can be set at different levels, with decreasing precedence as follows: +| Architecture | Model Family | Status | Notes | +|---|---|:---:|---| +| 🦙 Decoder-Only | LLaMA 2 / 3, Mistral, TinyLlama | ✅ Full Support | Primary target. All layers fully traced. | +| 🤖 Decoder-Only | Qwen 2 / 2.5 | ✅ Full Support | Standard `self_attn` / `mlp` structure. | +| 🧠 Decoder-Only | GPT-NeoX, StableLM | ✅ Full Support | Compatible naming conventions. | +| 🔤 Encoder-Only | BERT, RoBERTa | ⚠️ Partial | Linear/Norm layers traced; residual heuristics may need tuning. | +| 🔁 Encoder-Decoder | T5, BART, FLAN | ❌ Not Yet | Cross-attention and dual-stack propagation not implemented. | +| 🧬 Mixture of Experts | Mixtral, DeepSeek-MoE | ❌ Not Yet | Expert routing logic not yet handled. | -1. `slices.*.sources.parameters` - applying to a specific input slice -2. `slices.*.parameters` - applying to a specific output slice -3. `models.*.parameters` or `input_model_parameters` - applying to any tensors coming from specific input models -4. `parameters` - catchall +> **Note:** For `mergekit`, ~95% of all merges are between models in the same decoder-only family (e.g., two LLaMA fine-tunes). eX-LRP is built precisely for this use case. -### Tokenizer Configuration +--- -The tokenizer behavior can be configured in two ways: using the new `tokenizer` field (recommended) or the legacy `tokenizer_source` field (maintained for backward compatibility). These fields are mutually exclusive - you should use one or the other, not both. +## 🛠️ Getting Started -#### Modern Configuration (tokenizer) +### Installation -The `tokenizer` field provides fine-grained control over vocabulary and embeddings: +```sh +git clone https://github.com/Tusm11/mergekit.git +cd mergekit +git checkout feat/lrp-merge-v3 -```yaml -tokenizer: - source: "union" # or "base" or a specific model path - tokens: # Optional: configure specific tokens - : - source: ... # Specify embedding source - force: false # Optional: force this embedding for all models - pad_to_multiple_of: null # Optional: pad vocabulary size +pip install -e . # Install mergekit with eX-LRP support ``` -##### Tokenizer Source - -The `source` field determines the vocabulary of the output model: - --- `base`: Use vocabulary from the base model --- `"path/to/model"`: Use vocabulary from a specific model +**Requirements:** Python 3.10+, PyTorch 2.0+, Transformers 4.36+ -##### Token Embedding Handling +### Pre-computing LRP Scores (Optional) -When a tokenizer is configured, each input model's embedding matrix is adjusted to match the output vocabulary before being passed to the merge method. For tokens a model already has, its own embedding is used. For tokens a model is *missing*, a fallback embedding is assigned using these rules: +For maximum control, you can pre-compute relevance scores before merging: -- If the base model has the token, use the base model's embedding -- If only one model has the token, use that model's embedding -- Otherwise, use an average of all available embeddings - -The merge method then combines these per-model embeddings (original and filled-in) to produce the final output. This means the final embedding for a token present in multiple models is determined by your merge method (SLERP, linear, TIES, etc.), not simply taken from one model. - -You can override these defaults for specific tokens. Any tokens listed here that don't already exist in the output vocabulary will be added automatically, making this useful for introducing new special tokens. - -```yaml -tokenizer: - source: union - tokens: - # Use embedding from a specific model - <|im_start|>: - source: "path/to/chatml/model" - - # Force a specific embedding for all models - <|special|>: - source: "path/to/model" - force: true - - # Map a token to another model's token embedding - <|renamed_token|>: - source: - kind: "model_token" - model: "path/to/model" - token: "<|original_token|>" # or use token_id: 1234 - - # Use a zero embedding - <|unused|>: - source: - kind: "zero" +```sh +python lrp_computer.py \ + your-model-path \ + ./lrp-scores-output \ + --rule epsilon \ + --device cuda \ + --prompts "The capital of France is" "Artificial intelligence can" ``` -##### Practical Example +This generates a `lrp_scores.safetensors` file that can be passed to the merge configuration. -Here's how you might preserve both Llama 3 Instruct and ChatML prompt formats when merging models: - -```yaml -tokenizer: - source: union - tokens: - # ChatML tokens - <|im_start|>: - source: "chatml_model" - <|im_end|>: - source: "chatml_model" - - # Llama 3 tokens - force original embeddings - <|start_header_id|>: - source: "llama3_model" - force: true - <|end_header_id|>: - source: "llama3_model" - force: true - <|eot_id|>: - source: "llama3_model" - force: true -``` +--- -#### Legacy Configuration (tokenizer_source) +## 🚀 Quickstart -For backward compatibility, the `tokenizer_source` field is still supported: +### 1. Create a Merge Configuration ```yaml -tokenizer_source: "union" # or "base" or a model path -``` +# examples/lrp.yml +merge_method: lrp +base_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +parameters: + density: 0.7 # Keep top 70% most relevant weights +models: + - model: your-finetuned-model-1 + parameters: + weight: 1.0 + - model: your-finetuned-model-2 + parameters: + weight: 1.0 +dtype: float16 +``` + +### 2. Run the Merge -This provides basic tokenizer selection but lacks the fine-grained control of the modern `tokenizer` field. - -### Chat Template Configuration - -The optional `chat_template` field allows overriding the chat template used for the merged model. - -```yaml -chat_template: "auto" # or a template name or Jinja2 template +```sh +mergekit-yaml examples/lrp.yml ./output-merged-model \ + --cuda \ + --copy-tokenizer \ + --allow-crimes ``` -Options include: +### 3. Use Your Merged Model -- `"auto"`: Automatically select the most common template among input models -- Built-in templates: `"alpaca"`, `"chatml"`, `"llama3"`, `"mistral"`, `"exaone"` -- A Jinja2 template string for custom formatting +```python +from transformers import AutoModelForCausalLM, AutoTokenizer -### Examples +model = AutoModelForCausalLM.from_pretrained("./output-merged-model") +tokenizer = AutoTokenizer.from_pretrained("./output-merged-model") -Several examples of merge configurations are available in [`examples/`](examples/). +inputs = tokenizer("The future of AI is", return_tensors="pt") +outputs = model.generate(**inputs, max_new_tokens=50) +print(tokenizer.decode(outputs[0])) +``` -## Merge Methods +--- -`mergekit` offers many methods for merging models, each with its own strengths and weaknesses. Choosing the right method depends on your specific goals, the relationship between the models you're merging, and the desired characteristics of the final model. +## 📐 Architecture Diagram -For detailed explanations, parameter descriptions, and use cases for each method, please see our [**Merge Method Guide**](docs/merge_methods.md). +```mermaid +graph TD + subgraph INPUTS["📥 Inputs"] + A["🧠 Base Model
(e.g. Llama-3 Base)"] + B["🎯 Fine-Tuned Model
(e.g. Llama-3 Domain Expert)"] + C["📝 Calibration Prompts
(Sample text inputs)"] + end -### Method Overview + subgraph TASKVEC["📊 Task Vector Extraction"] + D["δ = θ_finetuned - θ_base
(Isolate learned knowledge)"] + end -| Method (`value`) | Core Idea | # Models | Base Model | Key Strengths / Use Cases | -|:----------------------------------------------------------------------------------------------------------------------|:--------------------------------------------------------------------|:--------:|:----:|:---------------------------------------------------------------| -| [**Linear** (`linear`)](docs/merge_methods.md#linear-linear) | Simple weighted average of model parameters. | ≥2 | - | Averaging similar checkpoints, model soups. | -| [**SLERP** (`slerp`)](docs/merge_methods.md#slerp-slerp) | Spherical linear interpolation between two models. | 2 | ✓ | Smoothly transitioning between two models. | -| [**NuSLERP** (`nuslerp`)](docs/merge_methods.md#nuslerp-nuslerp) | Enhanced SLERP with flexible weighting. | 2 | * | More intuitive SLERP; task vector SLERP. | -| [**Multi-SLERP** (`multislerp`)](docs/merge_methods.md#multi-slerp-multislerp) | Barycentric SLERP for multiple models. | ≥2 | * | Spherical interpolation for >2 models. | -| [**Karcher Mean** (`karcher`)](docs/merge_methods.md#karcher-mean-karcher) | Riemannian barycenter of model parameters. | ≥2 | - | Geometrically sound averaging on manifolds. | -| [**Task Arithmetic** (`task_arithmetic`)](docs/merge_methods.md#task-arithmetic-task_arithmetic) | Linearly combine "task vectors" (differences from a base). | ≥2 | ✓ | Transferring/combining fine-tuned skills. | -| [**TIES** (`ties`)](docs/merge_methods.md#ties-merging-ties) | Task arithmetic + sparsification & sign consensus. | ≥2 | ✓ | Merging many models, reducing interference. | -| [**DARE** (`dare_linear`, `dare_ties`)](docs/merge_methods.md#dare-dare_linear-dare_ties) | Task arithmetic + random pruning & rescaling. | ≥2 | ✓ | Robust skill retention, similar to TIES. | -| [**DELLA** (`della`, `della_linear`)](docs/merge_methods.md#della-della-della_linear) | Task arithmetic + adaptive magnitude-based pruning. | ≥2 | ✓ | Prioritizing important changes, reducing interference. | -| [**Model Breadcrumbs** (`breadcrumbs`, `breadcrumbs_ties`)](docs/merge_methods.md#model-breadcrumbs-breadcrumbs_ties) | Task arithmetic + outlier removal (small & large diffs). | ≥2 | ✓ | Refining task vectors by removing extreme changes. | -| [**SCE** (`sce`)](docs/merge_methods.md#sce-sce) | Task arithmetic + adaptive matrix-level weighting based on variance. | ≥2 | ✓ | Dynamically weighting models based on parameter variance. | -| [**Model Stock** (`model_stock`)](docs/merge_methods.md#model-stock-model_stock) | Geometric weight calculation for linear interpolation. | ≥3 | ✓ | Finding good linear interpolation weights for many checkpoints. | -| [**Nearswap** (`nearswap`)](docs/merge_methods.md#nearswap-nearswap) | Interpolate where parameters are similar. | 2 | ✓ | Selective merging based on parameter similarity. | -| [**Arcee Fusion** (`arcee_fusion`)](docs/merge_methods.md#arcee-fusion-arcee_fusion) | Dynamic thresholding for fusing important changes. | 2 | ✓ | Identifying and merging salient features. | -| [**LRP-Merge** (`lrp`)](README.md#lrp-merge-layer-wise-relevance-propagation-for-llm-merging) | Relevance propagation weights preservation. | ≥2 | ✓ | Preserving functionally critical weights using XAI. | -| [**Passthrough** (`passthrough`)](docs/merge_methods.md#passthrough-passthrough) | Directly copies tensors from a single input model. | 1 | - | Frankenmerging, layer stacking, model surgery. | + subgraph EXLRP["⚡ eX-LRP Computation Engine"] + direction TB + E["Forward Pass
Run prompts through model
Capture activations at
Self-Attn / MLP / Norm"] + F["Relevance Seeding
Extract top-predicted token logit
Set initial R_out"] + G["AttnLRP Backward Pass
Gradient × Input rules
Residual split logic
Numerical stabilization"] + H["Importance Score Tensor
(Per-weight relevance map)"] -**Key for `Base Model` Column:** + E --> F --> G --> H + end -- ✓: **Required** - One of the input models *must* be designated as the `base_model`. -- *: **Optional** - One of the input models *can* be designated as the `base_model`. -- -: **Not Applicable** - `base_model` has no effect on this method. + subgraph SPARSE["🔍 Sparsification"] + I["Apply density threshold
Generate binary relevance mask"] + end -## LoRA Extraction + subgraph MERGE["🔀 Merge Execution"] + J["Masked δ × weight
Weighted averaging
Add back to base"] + end -Mergekit allows extracting PEFT-compatible low-rank approximations of finetuned models. + K["✅ eX-LRP Merged Model
(Preserves critical knowledge)"] -### Usage + A --> D + B --> D + B --> E + C --> E + D --> I + H --> I + I --> J + A --> J + J --> K -```sh -mergekit-extract-lora --model finetuned_model_id_or_path --base-model base_model_id_or_path --out-path output_path [--no-lazy-unpickle] [--cuda] [--max-rank=desired_rank] [--sv-epsilon=tol] + style EXLRP fill:#1a1a2e,stroke:#e94560,stroke-width:3px,color:#fff + style K fill:#0f3460,stroke:#16c79a,stroke-width:2px,color:#fff ``` -## Mixture of Experts Merging - -The `mergekit-moe` script supports merging multiple dense models into a mixture of experts, either for direct use or for further training. For more details see the [`mergekit-moe` documentation](docs/moe.md). - -## Evolutionary Merge Methods +--- -See [`docs/evolve.md`](docs/evolve.md) for details. +## 🔧 Merge Configuration -## Multi-Stage Merging (`mergekit-multi`) +### eX-LRP Parameters -`mergekit-multi` enables the execution of complex, multi-stage model merging workflows. You can define multiple merge configurations in a single YAML file, where later merges can use the outputs of earlier ones as inputs. This is useful for building up sophisticated models through a series of targeted merges. +| Parameter | Type | Default | Description | +|---|---|---|---| +| `density` | float (0-1) | `0.7` | Fraction of most-relevant weights to keep. Lower = more aggressive pruning. | +| `weight` | float | `1.0` | Per-model weight for weighted averaging of task vectors. | -See the [`mergekit-multi` documentation](docs/multimerge.md) for usage details and examples. +### Example: Global-to-Local Transfer -## Raw PyTorch Model Merging (`mergekit-pytorch`) +This is the primary use case for eX-LRP — merging a base "global" model with a domain-specific "local" fine-tune: -For merging arbitrary PyTorch models (not necessarily Hugging Face Transformers), `mergekit-pytorch` provides a way to apply mergekit's algorithms directly to `.pt` or `.safetensors` checkpoints. The configuration is similar to the YAML format used in `mergekit-yaml`, but does not support layer slicing or tokenizer configuration. +```yaml +merge_method: lrp +base_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +parameters: + density: 0.7 +models: + - model: ./models/tinyllama-domain-expert + parameters: + weight: 1.0 +dtype: float16 +``` -### Usage +--- -```sh -mergekit-pytorch path/to/your/raw_config.yml ./output_pytorch_model_directory [options] -``` +## 📋 All Mergekit Merge Methods -Use `mergekit-pytorch --help` for detailed options. +`mergekit` supports many merging algorithms. eX-LRP is the first to use **Explainable AI** for weight selection. -## Tokenizer Transplantation (`mergekit-tokensurgeon`) +| Method (`value`) | Core Idea | # Models | Base | Key Use Case | +|:---|:---|:---:|:---:|:---| +| [**Linear** (`linear`)](docs/merge_methods.md#linear-linear) | Weighted average | ≥2 | - | Model soups | +| [**SLERP** (`slerp`)](docs/merge_methods.md#slerp-slerp) | Spherical interpolation | 2 | ✓ | Smooth transitions | +| [**Task Arithmetic** (`task_arithmetic`)](docs/merge_methods.md#task-arithmetic-task_arithmetic) | Linear task vectors | ≥2 | ✓ | Skill transfer | +| [**TIES** (`ties`)](docs/merge_methods.md#ties-merging-ties) | Sparsify + sign consensus | ≥2 | ✓ | Reduce interference | +| [**DARE** (`dare_linear`, `dare_ties`)](docs/merge_methods.md#dare-dare_linear-dare_ties) | Random pruning + rescale | ≥2 | ✓ | Robust retention | +| [**DELLA** (`della`)](docs/merge_methods.md#della-della-della_linear) | Magnitude-based pruning | ≥2 | ✓ | Adaptive pruning | +| **eX-LRP** (`lrp`) | **AttnLRP relevance propagation** | **≥2** | **✓** | **XAI-driven weight preservation** | +| [**Passthrough** (`passthrough`)](docs/merge_methods.md#passthrough-passthrough) | Direct copy | 1 | - | Layer surgery | -`mergekit-tokensurgeon` is a specialized tool for transplanting tokenizers between models, allowing you to align the vocabulary of one model with another. This is particularly useful for cheaply producing draft models for speculative decoding or for cross-tokenizer knowledge distillation. See the [documentation](docs/tokensurgeon.md) for more details and how to use it. +For the full list and detailed docs, see the [Merge Method Guide](docs/merge_methods.md). -## Instructions for Running LRP-Merge on Google Colab +--- -This guide explains how to train your models using Google Colab's free GPU, which is much faster than training on a CPU-only laptop. +## ☁️ Running on Google Colab -### Why Use Colab? +For users without a local GPU, eX-LRP can be run entirely on Google Colab's free tier. | Feature | Your Laptop (CPU) | Google Colab (Free GPU) | -|---------|-------------------|-------------------------| +|---|---|---| | Training Time (1 epoch) | 2-4 hours | 5-15 minutes | | Model Size Supported | GPT-2 (124M) | TinyLlama (1.1B) | | Batch Size | 1 | 4-8 | | Max Sequence Length | 64-128 | 256-512 | -### Prerequisites - -1. A Google account (free) -2. Your LRP merge method files uploaded to Google Drive -3. This notebook: `LRP_Merge_Colab_Training.ipynb` - -### Required Files to Upload - -You need to upload these files from your `D:\LRP merge method` folder to Google Drive: - -#### Core Python Files (Required) -``` -LRP merge method/ -├── finetune_fakenews.py # Training script (FIXED for Colab) -├── lrp_merge_pipeline.py # LRP merge pipeline (FIXED for Colab) -├── lrp_computer.py # LRP score computation (FIXED for Colab) -``` - -#### Supporting Files -``` -├── base.py # Base classes -├── common.py # Common utilities -├── graph.py # Graph operations -├── tasks.py # Task definitions -├── sparsify.py # Sparsification -├── embed.py # Embedding utilities -├── registry.py # Registry -└── __init__.py # Package init -``` - -#### Configuration Files -``` -├── lrp_config.yaml # LRP merge configuration template -└── pyproject.toml # Python dependencies (optional) -``` - -#### Your Datasets (Important!) -``` -datasets/ -└── synthetic/ # Your downloaded datasets - ├── train.csv # Training data (~800 samples) - └── test.csv # Test data (~200 samples) -``` - -**Note:** The notebook is configured to use `datasets/synthetic/train.csv` by default. - -#### Optional: Mergekit Repository -``` -mergekit_repo/ # Only if you're using custom merge methods -└── mergekit/ - └── ... -``` - -#### Complete File List (Copy-Paste Ready) - -Create a ZIP file with these files: -``` -base.py -common.py -embed.py -finetune_fakenews.py -graph.py -lrp.py -lrp_computer.py -lrp_config.yaml -lrp_merge_pipeline.py -pyproject.toml -registry.py -sparsify.py -tasks.py -__init__.py -download_fakenews_datasets.py -datasets/synthetic/train.csv -datasets/synthetic/test.csv -``` - -### Step-by-Step Instructions - -#### Step 1: Prepare Your Files - -##### Option A: Upload to Google Drive (Recommended) - -1. Open [Google Drive](https://drive.google.com) -2. Create a folder called `LRP merge method` -3. Upload all files from your `D:\LRP merge method` folder to this Drive folder -4. Wait for upload to complete - -##### Option B: Upload Directly to Colab - -1. Create a ZIP file of your `D:\LRP merge method` folder -2. Upload this notebook to Colab -3. Skip to "Upload Files Manually" section in the notebook - -#### Step 2: Open the Notebook in Colab - -1. Go to [Google Colab](https://colab.research.google.com) -2. Click **File** → **Upload notebook** -3. Select `LRP_Merge_Colab_Training.ipynb` -4. Wait for the notebook to load - -#### Step 3: Enable GPU Runtime - -**IMPORTANT: This step is crucial!** +### Quick Steps -1. Click **Runtime** in the top menu -2. Select **Change runtime type** -3. Under "Hardware accelerator", select **GPU** -4. Click **Save** -5. The runtime will restart - -**Verify GPU is enabled:** -- Look at the top right corner --- It should say "RAM" and "Disk" with a checkmark -- If you see "Connecting" or "Busy", wait for it to connect - -#### Step 4: Run the Notebook - -1. Click **Runtime** → **Run all** (or run cells one by one) -2. When prompted, authenticate Google Drive: - - A link will appear - - Click the link and sign in to your Google account - - Copy the authorization code - - Paste it back in Colab - - Press Enter - -#### Step 5: Wait for Training +1. Upload your files to Google Drive +2. Open `LRP_Merge_Colab_Training.ipynb` in Colab +3. Enable **GPU Runtime** (Runtime → Change runtime type → GPU) +4. Run all cells The notebook will automatically: -1. Mount your Google Drive -2. Install all dependencies -3. Create sample datasets -4. Train the GLOBAL model (~10-15 minutes) -5. Train the LOCAL model (~10-15 minutes) -6. Compute LRP scores (~5 minutes) -7. Merge the models (~5 minutes) -8. Test and evaluate (~5 minutes) +- Install dependencies +- Train Global and Local models +- Compute eX-LRP relevance scores +- Merge the models +- Evaluate results **Total time: ~40-60 minutes** -#### Step 6: Save Results - -The trained models are automatically saved to your Google Drive at: -``` -My Drive/LRP merge method/models/ (varies based on user's folder's structure) -├── tinyllama-global/ (varies from model to model) -├── tinyllama-local/ -└── merged-model/ -``` - -You can also download them as a ZIP file (see the last cell). - -### Configuration Options - -#### Change the Model - -In the notebook, find this cell: -```python -MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" (varies from model to model) -``` - -Options: -- `"TinyLlama/TinyLlama-1.1B-Chat-v1.0"` - Good quality, requires ~8GB GPU -- `"gpt2"` - Fastest, works on any GPU -- `"gpt2-medium"` - Balance of speed and quality - -#### Change Training Parameters - -```python -EPOCHS = 3 # Increase for better results (try 5-10) -BATCH_SIZE = 4 # Can increase if you have more GPU memory -MAX_SAMPLES = 1000 # Increase for better training (try 5000-10000) -``` - -#### Use Your Own Dataset - -1. Upload your dataset CSV to Google Drive in the `datasets` folder -2. Update the path in the notebook: -```python -DATASET = "datasets/your_dataset.csv" -``` - -### Troubleshooting - -#### "No GPU detected" Error - -**Solution:** -1. Go to Runtime → Change runtime type -2. Select GPU -3. Click Save -4. Runtime will restart - run cells again - -#### "Out of Memory" Error - -**Solution:** Reduce memory usage: -```python -BATCH_SIZE = 2 # Reduce from 4 -MAX_SAMPLES = 500 # Reduce from 1000 -MODEL_NAME = "gpt2" # Use smaller model -``` +For detailed Colab instructions, see the [Colab Guide](docs/colab_guide.md). -#### "File not found" Error - -**Solution:** Check your Google Drive path: -```python -# Update this to match your Drive structure -LRP_PATH = "/content/drive/MyDrive/LRP merge method" -``` - -#### Training Takes Too Long - -**Solutions:** -1. Use smaller model: `MODEL_NAME = "gpt2"` -2. Reduce samples: `MAX_SAMPLES = 500` -3. Reduce epochs: `EPOCHS = 1` - -#### Runtime Disconnected - -Colab may disconnect after ~90 minutes of inactivity or 12 hours max. - -**Solutions:** -1. Click around periodically to keep it active -2. Save results to Drive frequently (the notebook does this automatically) -3. Use Colab Pro ($10/month) for longer runtimes - -### Colab Pro vs Free - -| Feature | Free | Pro ($10/month) | -|---------|------|-----------------| -| GPU | K80/T4 | T4/P100 | -| Runtime limit | 12 hours | 24 hours | -| Idle timeout | 90 min | None | -| Background execution | No | Yes | - -For LRP-Merge training, the **free tier is usually sufficient**. - -### Downloading Results - -#### Method 1: From Google Drive -1. Open [Google Drive](https://drive.google.com) -2. Navigate to `My Drive/LRP merge method/models/` -3. Right-click the model folder → Download - -#### Method 2: Direct Download from Colab -Run the last cell in the notebook which creates a ZIP file and downloads it. - -#### Method 3: Use `gdown` (Command Line) -```bash -pip install gdown -gdown -``` - -### 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? +## 🙏 Acknowledgements -- 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 +- **[Mergekit](https://github.com/arcee-ai/mergekit)** by Arcee AI — The foundational toolkit this method is built on. +- **[LRP-eXplains-Transformers (LXT)](https://github.com/rachtibat/LRP-eXplains-Transformers)** — The AttnLRP paper and codebase that inspired the Gradient × Input formulation used in eX-LRP. +- **[AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers](https://proceedings.mlr.press/v235/achtibat24a.html)** (ICML 2024) — The foundational paper by Achtibat et al. --- -**Remember:** Colab sessions are temporary. Always save your results to Google Drive before closing! - -## Citation +## 📄 Citation -If you find `mergekit` useful in your research, please consider citing the [paper](https://aclanthology.org/2024.emnlp-industry.36/): +If you use eX-LRP in your research, please cite both Mergekit and the AttnLRP paper: ```bibtex @inproceedings{goddard-etal-2024-arcees, title = "Arcee{'}s {M}erge{K}it: A Toolkit for Merging Large Language Models", - author = "Goddard, Charles and - Siriwardhana, Shamane and - Ehghaghi, Malikeh and - Meyers, Luke and - Karpukhin, Vladimir and - Benedict, Brian and - McQuade, Mark and - Solawetz, Jacob", - editor = "Dernoncourt, Franck and - Preo{\c{t}}iuc-Pietro, Daniel and - Shimorina, Anastasia", + author = "Goddard, Charles and Siriwardhana, Shamane and Ehghaghi, Malikeh and Meyers, Luke and Karpukhin, Vladimir and Benedict, Brian and McQuade, Mark and Solawetz, Jacob", booktitle = "Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track", month = nov, year = "2024", - address = "Miami, Florida, US", publisher = "Association for Computational Linguistics", url = "https://aclanthology.org/2024.emnlp-industry.36", doi = "10.18653/v1/2024.emnlp-industry.36", pages = "477--485", - abstract = "The rapid growth of open-source language models provides the opportunity to merge model checkpoints, combining their parameters to improve performance and versatility. Advances in transfer learning have led to numerous task-specific models, which model merging can integrate into powerful multitask models without additional training. MergeKit is an open-source library designed to support this process with an efficient and extensible framework suitable for any hardware. It has facilitated the merging of thousands of models, contributing to some of the world{'}s most powerful open-source model checkpoints. The library is accessible at: https://github.com/arcee-ai/mergekit.", +} + +@InProceedings{pmlr-v235-achtibat24a, + title = "{A}ttn{LRP}: Attention-Aware Layer-Wise Relevance Propagation for Transformers", + author = "Achtibat, Reduan and Hatefi, Sayed Mohammad Vakilzadeh and Dreyer, Maximilian and Jain, Aakriti and Wiegand, Thomas and Lapuschkin, Sebastian and Samek, Wojciech", + booktitle = "Proceedings of the 41st International Conference on Machine Learning", + pages = "135--168", + year = "2024", + volume = "235", + series = "Proceedings of Machine Learning Research", + publisher = "PMLR", } ``` + +--- + +## 📜 License + +This project is licensed under the [LGPL-3.0 License](https://www.gnu.org/licenses/lgpl-3.0). From 0125b2e53996a9b396431801241c6e64f95bdfc3 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 13:18:32 +0530 Subject: [PATCH 02/13] Address PR feedback: wire lrp_scores, migrate to lxt for true AttnLRP, add LRP mask test, and clean up files --- LRP_Merge.ipynb | 9193 ++++++++++++++------------------- examples/lrp.yml | 4 +- finetune_fakenews.py | 392 -- git | 0 lrp_computer.py | 292 +- mergekit/merge_methods/lrp.py | 12 +- tests/test_lrp_merge.py | 88 + 7 files changed, 4073 insertions(+), 5908 deletions(-) delete mode 100644 finetune_fakenews.py delete mode 100644 git create mode 100644 tests/test_lrp_merge.py diff --git a/LRP_Merge.ipynb b/LRP_Merge.ipynb index d0d96a8b..b87ab72f 100644 --- a/LRP_Merge.ipynb +++ b/LRP_Merge.ipynb @@ -1,5321 +1,3990 @@ { - "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." - ] + "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/" }, - { - "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" - } + "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\"\u2705 Project root set to: {os.getcwd()}\")\n", + "print(f\"Models will be saved to: {os.path.abspath(MODEL_DIR)}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "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." - ] + "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\"\u2705 Created .gitignore at: {gitignore_path}\")\n", + "print(\"This file will prevent large models and temporary files from being tracked by Git.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "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": [] + "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(\"\u2705 Created .gitignore to protect your repo from large model files.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "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\"\u2705 Successfully merged to {output_path}\")\n", + " else:\n", + " print(f\"\u2717 Error: {res.stderr}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "cell_type": "code", - "source": [ - "#command to install mergekit\n", - "!pip install -q -U mergekit" - ], - "metadata": { - "id": "4f1eC6KBjR_6" - }, - "execution_count": null, - "outputs": [] + "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": null, + "outputs": [] + }, + { + "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": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "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" - ] - } - ] + "id": "SuyAw89LfY5M", + "outputId": "1c70faaa-c4f7-486a-a007-90a8b62db93f" + }, + "outputs": [], + "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\u2705 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/" }, - { - "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" - ] - } - ] + "id": "kPBMV65LfY5R", + "outputId": "3a192e03-e9f8-4a44-c812-589e4a75daa9" + }, + "outputs": [], + "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/" }, - { - "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" - ] - } - ] + "id": "PpiAn5MOfY5T", + "outputId": "9d470042-b9fe-44d7-8884-0bd01ea79290" + }, + "outputs": [], + "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\u2713 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\u2717 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/" }, - { - "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": [] + "id": "cm6LGAehfY5T", + "outputId": "c98da997-db25-4387-dbf4-e05cc59000e9" + }, + "outputs": [], + "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/" }, - { - "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" - ] - } - ] + "id": "eyyQpTx_fY5U", + "outputId": "3ffa4644-e645-409f-a2f0-85c2a39889c5" + }, + "outputs": [], + "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/" }, - { - "cell_type": "markdown", - "metadata": { - "id": "4DdPOPSnfY5L" - }, - "source": [ - "##Mount Google Drive\n", - "\n", - "This connects your Google Drive to access the LRP merge method files." - ] + "id": "Uc3SMHrCfY5U", + "outputId": "70ee41cd-1106-4c4e-a966-332e64a6f62f" + }, + "outputs": [], + "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/" }, - { - "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!\")" - ] + "id": "jW042sK3fY5V", + "outputId": "bb5507aa-f105-4427-d517-8ca6a72b16aa" + }, + "outputs": [], + "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/" }, - { - "cell_type": "markdown", - "metadata": { - "id": "m7giRqKGfY5Q" - }, - "source": [ - "## Installing Dependencies" - ] + "id": "mT3EWmZIfY5V", + "outputId": "787d7d3d-86a7-45ac-cbae-fb2024223b0b" + }, + "outputs": [], + "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/" }, - { - "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!\")" - ] + "id": "P36cXjdUfY5V", + "outputId": "5e182196-b403-4d8c-efc0-51aa3f289ff0" + }, + "outputs": [], + "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\u2713 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/" }, - { - "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." - ] + "id": "XNQMG2a0fY5W", + "outputId": "2932e596-955f-45a5-b6c8-51436a1388de" + }, + "outputs": [], + "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(\"\u2713 Created lrp_config_colab.yaml with corrected structure\")\n", + "!cat lrp_config_colab.yaml" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "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.\")" - ] + "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": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "cell_type": "markdown", - "metadata": { - "id": "nHpFGULJfY5T" - }, - "source": [ - "## Step 8: Train GLOBAL Model\n", - "\n", - "This trains the general knowledge model." - ] + "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": null, + "outputs": [] + }, + { + "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" + ] }, - { - "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}\")" - ] + "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": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "##Merging the Global and Local model before using LRP method" + ], + "metadata": { + "id": "uwc7OAgmgfKl" + } + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 347, + "referenced_widgets": [ + "afebbe9105224e5ca383c8c752ed5775", + "754b6f10f5bb49f5b621352d1dc3470c", + "f677e65d178a4a30a0b7895814b34bc1", + "dc14a6b06cb84a4b877f42abf48cbdc2", + "19f4f141e14c45c0bd93ad7ca1ad4513", + "e6265f6e7e934e2081bea29e33bc7d15", + "f4ac3ccb34d74016830d70a464d3dbb9", + "9492bdd651ff4b01b4790ea5bf833d39", + "2e0fa57592eb44e0bfe7357197ca2a16", + "206608ad94554710af5047a3a3d08eab", + "c7e9a370c64a4c359005aee9dc72090d", + "52f24136d0ed4890bcb6c25c16730dee", + "646ab867da2446fe9791cfebae223aa7", + "7c7d14a7512846a1aea8b71f22b4a2e6", + "14fbfe3a4e5d44b7a22153db3cca1aa8", + "6d36895fe0c64bd9b47270d936b01fcd", + "0c62cea3d1bf4704a03e0580dfe9ea73", + "7301c18c7ef44c08a2135ae3d5e0d43d", + "3b65be145d344b8fab0e3e5b94d8b8bb", + "15d3e325039e4018b24f927bca7fff77", + "9d4071d6c77c422ea30317d0ee2aadc5", + "b904b21a75864d87aa71802c8870a766", + "87268e0da9904873b7910e7b5dd7a57b", + "aa12cd6df73c41c8b82efe7fb5fc05da", + "19b500f0b9704c0095db07a00ef01ef4", + "3b2e30fd08794a828ae64958ca18158d", + "d1f27967d90440018152bf8f79b43b7e", + "c04bf306209944ae9d8729691ae54377", + "641e647ba0bb409f934a9be0bed1fd86", + "ee68b3494fb64c8d87497375ed024522", + "036f981877d14a85b090c64ddd404654", + "204b3ae8fa174bd4b31679c8058fc8c1", + "a5e087211dad429fb555cc513829dd85", + "1e3d1f1fa762423f8408979bf6202de6", + "99daa9d955d847ed81d1915e49f401d4", + "14bbf4d9e37d4c99ae34b1692ebcdb5c", + "a674641859a745098ee1175b54090cdb", + "654ebccf2a9f4d42869731b9e2b85355", + "3ea5f48919974de3989c590af201e5e3", + "7bd6b68cfa2345c9aa75906b99b70284", + "5a4bfca98ca241748dcd57226c9dcfcc", + "4f786bd17a244d50bb39dcb7e83b29a7", + "dc18364f3e554575bd3fc2db300fb4d4", + "304304c7321e41ecbf48bdcd6b06fc54" + ] }, - { - "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" - ] + "id": "6caa5c8b", + "outputId": "3d552cf0-b734-4535-bde6-b2dfab7b206f" + }, + "source": [ + "import os\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "from peft import PeftModel\n", + "\n", + "# Ensure LRP_PATH is defined for this cell's scope\n", + "LRP_PATH = \"/content/drive/MyDrive/LRP Merge method\"\n", + "\n", + "# Define full paths for base and adapter models\n", + "LOCAL_BASE_MODEL_PATH = os.path.join(LRP_PATH, \"base_model\")\n", + "GLOBAL_ADAPTER_PATH = os.path.join(LRP_PATH, \"models\", \"tinyllama-global\")\n", + "LOCAL_ADAPTER_PATH = os.path.join(LRP_PATH, \"models\", \"tinyllama-local\")\n", + "\n", + "GLOBAL_FULL_MODEL_OUTPUT = os.path.join(LRP_PATH, \"models\", \"tinyllama-global-full\")\n", + "LOCAL_FULL_MODEL_OUTPUT = os.path.join(LRP_PATH, \"models\", \"tinyllama-local-full\")\n", + "\n", + "# Ensure output directories exist\n", + "os.makedirs(GLOBAL_FULL_MODEL_OUTPUT, exist_ok=True)\n", + "os.makedirs(LOCAL_FULL_MODEL_OUTPUT, exist_ok=True)\n", + "\n", + "print(f\"Loading base model from: {LOCAL_BASE_MODEL_PATH}\")\n", + "base = AutoModelForCausalLM.from_pretrained(LOCAL_BASE_MODEL_PATH)\n", + "\n", + "print(f\"Loading global adapter from: {GLOBAL_ADAPTER_PATH}\")\n", + "global_model = PeftModel.from_pretrained(base, GLOBAL_ADAPTER_PATH)\n", + "print(\"Merging global model and unloading...\")\n", + "global_model = global_model.merge_and_unload()\n", + "\n", + "print(f\"Saving merged global model to: {GLOBAL_FULL_MODEL_OUTPUT}\")\n", + "global_model.save_pretrained(GLOBAL_FULL_MODEL_OUTPUT)\n", + "\n", + "print(\"Reloading base model for local adapter merging...\")\n", + "# Reload base again (important for fresh state if base weights are modified during merge)\n", + "base = AutoModelForCausalLM.from_pretrained(LOCAL_BASE_MODEL_PATH)\n", + "\n", + "print(f\"Loading local adapter from: {LOCAL_ADAPTER_PATH}\")\n", + "local_model = PeftModel.from_pretrained(base, LOCAL_ADAPTER_PATH)\n", + "print(\"Merging local model and unloading...\")\n", + "local_model = local_model.merge_and_unload()\n", + "\n", + "print(f\"Saving merged local model to: {LOCAL_FULL_MODEL_OUTPUT}\")\n", + "local_model.save_pretrained(LOCAL_FULL_MODEL_OUTPUT)\n", + "\n", + "print(\"\u2713 Global and local models merged and saved successfully!\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "##Download the custom mergekit repo before applying LRP method" + ], + "metadata": { + "id": "BmgjayfJgx9G" + } + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "cell_type": "markdown", - "metadata": { - "id": "V-uP2FOFfY5U" - }, - "source": [ - "##Train LOCAL Model\n", - "\n", - "This trains the task-specific model." - ] + "id": "95429cf5", + "outputId": "6dc56ede-6979-4352-fa52-4d439720f7a9" + }, + "source": [ + "import os\n", + "\n", + "mergekit_repo_path = os.path.join(LRP_PATH, \"mergekit_repo\")\n", + "\n", + "if os.path.exists(mergekit_repo_path) and os.path.isdir(mergekit_repo_path):\n", + " print(f\"Found custom mergekit_repo at: {mergekit_repo_path}\")\n", + " print(\"Installing custom mergekit_repo in editable mode...\")\n", + " !pip install -e \"{mergekit_repo_path}\"\n", + " print(\"\u2713 Custom mergekit_repo installed.\")\n", + "else:\n", + " print(f\"\u2717 Custom mergekit_repo not found at: {mergekit_repo_path}\")\n", + " print(\"Please ensure 'mergekit_repo' directory containing your LRP implementation is present in your Google Drive at: {LRP_PATH}/mergekit_repo\")\n", + " #Exit if custom mergekit_repo is not found, as the LRP method will not work.\n", + " import sys\n", + " sys.exit(\"Error: Custom mergekit_repo not found. Please place it in your Google Drive as specified.\")\n", + "\n", + "# Re-check available merge methods after installation\n", + "print(\"Checking available merge methods in mergekit after custom installation...\")\n", + "print(\"Note: A full check of merge methods might cause a SystemExit if the custom method isn't fully integrated yet. Proceeding with script fixes.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5-ItEziCfY5X" + }, + "source": [ + "## Test the Merged Model" + ] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "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\")" - ] + "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\u2705 LRP MERGE SUCCESSFUL\")\n", + "\n", + "run_final_merge()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Testing the final merged model" + ], + "metadata": { + "id": "TO-1UJErhK-b" + } + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 343, + "referenced_widgets": [ + "40bf519b3bc84b8395ee4df0ef94857d", + "bd37344622db48389e249b3042745bf0", + "4f46e2b9646a4c34a7b78f262dfded5d", + "7cd3262cde474f598a9c30f98013792a", + "aacb526fff5047e280983f36fb0b23ea", + "a87d74f47b5e46da9a7d1bd22cba0d07", + "cce6ded52f8248d295018ab218f52a09", + "796da76c23b44542a0afb53d0104b465", + "541b71d665d64b788d7a1ce2b5af0a77", + "f80704a30c894717bf8925f3780b1d7c", + "d3a50b0ecae34e0981ea19adf5071973" + ] }, - { - "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" - ] + "id": "7f364544", + "outputId": "6374fbeb-c84c-49e3-b8bf-6b7a8f683d12" + }, + "source": [ + "import torch\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "import os\n", + "\n", + "# Path to your freshly merged model\n", + "MERGED_MODEL_PATH = os.path.join(LRP_PATH, \"models/merged-model\")\n", + "\n", + "print(f\"Loading merged model from: {MERGED_MODEL_PATH}\")\n", + "\n", + "try:\n", + " # Load tokenizer and model\n", + " tokenizer = AutoTokenizer.from_pretrained(MERGED_MODEL_PATH)\n", + " model = AutoModelForCausalLM.from_pretrained(\n", + " MERGED_MODEL_PATH,\n", + " torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,\n", + " device_map=\"auto\"\n", + " )\n", + "\n", + " # Prepare a test prompt\n", + " prompt = \"### Instruction: Classify this news as REAL or FAKE.\\n\\n### Input: Scientists discover new planet using James Webb telescope.\\n\\n### Response:\"\n", + "\n", + " inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n", + "\n", + " print(\"\\nGenerating response...\")\n", + " with torch.no_grad():\n", + " outputs = model.generate(**inputs, max_new_tokens=10, temperature=0.1)\n", + "\n", + " result = tokenizer.decode(outputs[0], skip_special_tokens=True)\n", + " print(\"=\"*50)\n", + " print(\"Test Result:\")\n", + " print(result)\n", + " print(\"=\"*50)\n", + "\n", + "except Exception as e:\n", + " print(f\"Error testing model: {e}\")\n", + " print(\"Make sure the merge process finished successfully and the files exist in Drive.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "7f9bac3e" + }, + "source": [ + "!pip install -q -U mergekit" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" }, - { - "cell_type": "markdown", - "metadata": { - "id": "P6elKTCgfY5V" - }, - "source": [ - "## Computing LRP Scores\n", - "We calculate the importance of each weight using the LRP epsilon rule." - ] + "id": "dadb7a2d", + "outputId": "3e657346-4ed3-4050-d54c-0f67b1846e48" + }, + "source": [ + "import os\n", + "from google.colab import drive\n", + "\n", + "# Check if drive is still mounted\n", + "if not os.path.exists('/content/drive/MyDrive'):\n", + " print(\"Drive not detected. Re-mounting...\")\n", + " drive.mount('/content/drive', force_remount=True)\n", + "\n", + "# List contents of MyDrive to verify the folder name\n", + "print(\"Contents of MyDrive:\")\n", + "!ls -d /content/drive/MyDrive/LRP*\n", + "\n", + "# Reset LRP_PATH based on verified name\n", + "# (Sometimes spaces or casing in folder names cause issues)\n", + "LRP_PATH = \"/content/drive/MyDrive/LRP Merge method\"\n", + "print(f\"\\nLRP_PATH set to: {LRP_PATH}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Evaluation of the model" + ], + "metadata": { + "id": "7dC09yLFhXT0" + } + }, + { + "cell_type": "code", + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 306, + "referenced_widgets": [ + "56149efbe9a3427ea568c5e88c59b6c0", + "87b08b3da6974dae946bf0584d64afa5", + "3191005414a543f3a8834af3153aca28", + "5b802d7678854dd2b6251c9278559695", + "5d0e827084e7445591bf56f79827175a", + "b9dca4dcaf3a48ce844e744b1570549d", + "e0d7bfc7a6e846a48203fb3616f698d9", + "893b22a44ea14f29be457a809cdbaa38", + "84f5c5abb1864fb4a44e5439e26daf9b", + "b716717de9124ae48eddba487d39679d", + "00b31c2532624ae59ba1d18acb98d198" + ] }, - { - "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" - ] - } + "id": "fe46e843", + "outputId": "0c76eb26-2719-4096-da7d-b0c9884085f4" + }, + "source": [ + "# Test Evaluation\n", + "import torch\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "import pandas as pd\n", + "\n", + "path = os.path.join(LRP_PATH, \"models/merged-model\")\n", + "tokenizer = AutoTokenizer.from_pretrained(path)\n", + "model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float32, device_map=\"auto\")\n", + "\n", + "test_data = [\n", + " {\"text\": \"NASA announces new space mission to Mars.\", \"label\": \"REAL\"},\n", + " {\"text\": \"Aliens found living on the moon.\", \"label\": \"FAKE\"}\n", + "]\n", + "\n", + "for item in test_data:\n", + " inputs = tokenizer(f\"### Input: {item['text']}\\n### Response:\", return_tensors=\"pt\").to(model.device)\n", + " out = model.generate(**inputs, max_new_tokens=5)\n", + " print(f\"Text: {item['text']} | Prediction: {tokenizer.decode(out[0], skip_special_tokens=True).split('### Response:')[-1].strip()}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TkxF67ZdfY5Y" + }, + "source": [ + "## Save Results to Google Drive\n", + "\n", + "Ensure all trained models are saved back to your Google Drive." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8X9dqxnLfY5Y" + }, + "outputs": [], + "source": [ + "import shutil\n", + "\n", + "# List what we have\n", + "!ls -lh models/\n", + "\n", + "# If models are in /content (not Drive), copy them\n", + "if not LRP_PATH.startswith(\"/content/drive\"):\n", + " DRIVE_PATH = \"/content/drive/MyDrive/LRP merge method/models\"\n", + " os.makedirs(DRIVE_PATH, exist_ok=True)\n", + "\n", + " # Copy trained models\n", + " for model_dir in [\"tinyllama-global\", \"tinyllama-local\", \"merged-model\"]:\n", + " src = f\"models/{model_dir}\"\n", + " dst = f\"{DRIVE_PATH}/{model_dir}\"\n", + " if os.path.exists(src):\n", + " print(f\"Copying {src} to Drive...\")\n", + " shutil.copytree(src, dst, dirs_exist_ok=True)\n", + "\n", + " print(f\"\\n\u2713 Models copied to: {DRIVE_PATH}\")\n", + "else:\n", + " print(\"\u2713 Models already saved to Google Drive\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "jhni4e8KfY5Z" + }, + "source": [ + "## Download Trained Models (Optional)\n", + "\n", + "Download the trained models to your local machine." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "DEJTa7lwfY5Z" + }, + "outputs": [], + "source": [ + "# Create a ZIP file of the models\n", + "!zip -r trained_models.zip models/\n", + "\n", + "# Download the ZIP\n", + "from google.colab import files\n", + "files.download(\"trained_models.zip\")\n", + "\n", + "print(\"Download started!\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + }, + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "accelerator": "GPU", + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "39cd0423589d4e30b97bc66781f355ca": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HBoxModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_43b4885dd82d410fb6bd848c8132665e", + "IPY_MODEL_2813588097ee4a6d9e232f4ace8c9627", + "IPY_MODEL_b7056c9990874958ae323c3193d785ec" ], - "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" - ] + "layout": "IPY_MODEL_f851195eb3d847e780fb0db248a8782b" + } }, - { - "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\")" - ] + "43b4885dd82d410fb6bd848c8132665e": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HTMLModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_6b2252d926e54bb6acaf04e1cc5628d4", + "placeholder": "\u200b", + "style": "IPY_MODEL_173ee7ac871b40048cffa93911762ab9", + "value": "Loading\u2007weights:\u2007100%" + } }, - { - "cell_type": "markdown", - "metadata": { - "id": "vUHsVGb2fY5V" - }, - "source": [ - "## Reconstructing Full Models\n", - "LRP-Merge performs best when merging full model weights rather than raw adapters." - ] + "2813588097ee4a6d9e232f4ace8c9627": { + "model_module": "@jupyter-widgets/controls", + "model_name": "FloatProgressModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_04f102d8728649fabe4ea26808fe172e", + "max": 201, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_a0378f8a7a9145bbbdd864e2b5faafe4", + "value": 201 + } }, - { - "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" - ] + "b7056c9990874958ae323c3193d785ec": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HTMLModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_d03be308d05d420dbd6572415fc0dab6", + "placeholder": "\u200b", + "style": "IPY_MODEL_9d8d202635384314b34e783ca92a80df", + "value": "\u2007201/201\u2007[00:00<00:00,\u2007207.27it/s,\u2007Materializing\u2007param=model.norm.weight]" + } }, - { - "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" - ] - } - ] + "f851195eb3d847e780fb0db248a8782b": { + "model_module": "@jupyter-widgets/base", + "model_name": "LayoutModel", + "model_module_version": "1.2.0", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } }, - { - "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" - ] - } - ] + "6b2252d926e54bb6acaf04e1cc5628d4": { + "model_module": "@jupyter-widgets/base", + "model_name": "LayoutModel", + "model_module_version": "1.2.0", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } }, - { - "cell_type": "markdown", - "source": [ - "## Download the base model" - ], - "metadata": { - "id": "y2b9BH1VgWBC" - } + "173ee7ac871b40048cffa93911762ab9": { + "model_module": "@jupyter-widgets/controls", + "model_name": "DescriptionStyleModel", + "model_module_version": "1.5.0", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } }, - { - "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" - ] - } - ] + "layout": "IPY_MODEL_5437721839cf4b938ab369a56709ff30" + } }, - { - "cell_type": "markdown", - "metadata": { - "id": "5-ItEziCfY5X" - }, - "source": [ - "## Test the Merged Model" - ] + "2b653b43913e44c6a62e6b708a29cda1": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HTMLModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_b0c54005c7af4cb7b64ed3610e3b5664", + "placeholder": "\u200b", + "style": "IPY_MODEL_0968426d097d4885ad5b48945043a5dc", + "value": "Writing\u2007model\u2007shards:\u2007100%" + } }, - { - "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 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 deleted file mode 100644 index e69de29b..00000000 diff --git a/lrp_computer.py b/lrp_computer.py index 5758cbd6..7ed1f563 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -42,6 +42,7 @@ def __init__(self, config: LRPConfig): self.tokenizer = None self.relevance_scores: Dict[str, torch.Tensor] = {} self.activations: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {} + self.module_relevance: Dict[str, torch.Tensor] = {} # Store propagated relevance self.hooks: List[Any] = [] def load_model(self) -> None: @@ -85,274 +86,65 @@ def load_model(self) -> None: self.model.eval() - def compute_relevance_epsilon( - self, - activations: torch.Tensor, - weights: torch.Tensor, - output_relevance: torch.Tensor, - ) -> torch.Tensor: - """ - Compute relevance using LRP-epsilon rule. - R_j = sum_k (z_jk / (sum_j z_jk + epsilon)) * R_k - """ - epsilon = self.config.epsilon - - # Ensure activations and weights are on same device - activations = activations.to(weights.device) - output_relevance = output_relevance.to(weights.device) - - # Compute forward pass contribution z = xW^T - z = F.linear(activations, weights) - - # Add epsilon for numerical stability - z_stable = z + epsilon * torch.sign(z) - - # Compute redistribution factor s = R_out / z_stable - s = output_relevance / z_stable - - # Flatten batch and sequence: (B, S, F) -> (N, F) - x_flat = activations.reshape(-1, activations.shape[-1]) - s_flat = s.reshape(-1, s.shape[-1]) - - # Relevance for weights: |W_ij * x_i * s_j| summed over batch/seq - weight_relevance = weights.abs() * (s_flat.abs().t() @ x_flat.abs()) - - return weight_relevance - - def compute_relevance_gamma( - self, - activations: torch.Tensor, - weights: torch.Tensor, - output_relevance: torch.Tensor, - ) -> torch.Tensor: - """ - Compute relevance using LRP-gamma rule. - Adds positive contributions with a gamma factor. - """ - gamma = self.config.gamma - epsilon = self.config.epsilon - - activations = activations.to(weights.device) - output_relevance = output_relevance.to(weights.device) - - # Separate positive contributions - weights_pos = torch.clamp(weights, min=0) - w_gamma = weights + gamma * weights_pos - - # Forward pass with enhanced weights - z = F.linear(activations, w_gamma) - z = z + epsilon * torch.sign(z) - - # Redistribute relevance - s = output_relevance / z - - x_flat = activations.reshape(-1, activations.shape[-1]) - s_flat = s.reshape(-1, s.shape[-1]) - - weight_relevance = weights.abs() * (s_flat.abs().t() @ x_flat.abs()) - - return weight_relevance - - def compute_relevance_alpha_beta( - self, - activations: torch.Tensor, - weights: torch.Tensor, - output_relevance: torch.Tensor, - ) -> torch.Tensor: - """ - Compute relevance using LRP-alpha_beta rule. - Separates positive and negative contributions. - """ - alpha = self.config.alpha - beta = self.config.beta - epsilon = self.config.epsilon - - activations = activations.to(weights.device) - output_relevance = output_relevance.to(weights.device) - - weights_pos = torch.clamp(weights, min=0) - weights_neg = torch.clamp(weights, max=0) - - # Positive and negative forward passes - z_pos = F.linear(activations, weights_pos) - z_neg = F.linear(activations, weights_neg) - - z = alpha * z_pos + beta * z_neg - z = z + epsilon * torch.sign(z) - - s = output_relevance / z - - x_flat = activations.reshape(-1, activations.shape[-1]) - s_flat = s.reshape(-1, s.shape[-1]) - - weight_relevance = weights.abs() * (s_flat.abs().t() @ x_flat.abs()) - - return weight_relevance - - def compute_gradcam_importance( - self, - input_ids: torch.Tensor, - target_layer: str, - ) -> torch.Tensor: - """ - Compute importance using Grad-CAM style gradients. - This is a practical alternative to full LRP. - """ - self.model.zero_grad() - - # Enable gradients for input - embedding_layer = self.model.get_input_embeddings() - inputs_embeds = embedding_layer(input_ids) - inputs_embeds.requires_grad_(True) - - # Forward pass - outputs = self.model(inputs_embeds=inputs_embeds, output_hidden_states=True) - logits = outputs.logits - - # Compute gradient of output w.r.t. embeddings - target_token_idx = logits.shape[1] - 1 - target_logit = logits[0, target_token_idx, :].max() - target_logit.backward() - - # Get gradients - gradients = inputs_embeds.grad - - # Importance = gradient magnitude - importance = torch.abs(gradients) - - return importance - - def compute_relevance_for_tensor( - self, - tensor_name: str, - tensor: torch.Tensor, - sample_activations: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - Compute relevance scores for a specific tensor. - - Args: - tensor_name: Name of the tensor (e.g., "model.layers.0.attn.q_proj.weight") - tensor: The weight tensor - sample_activations: Optional sample activations from a forward pass - - Returns: - Relevance scores for the tensor - """ - # If we have sample activations, use proper LRP rules - if sample_activations is not None: - # Create dummy output relevance (normally from backward pass) - # Match the shape of the output activations - output_relevance = torch.ones_like( - ( - F.linear(sample_activations, tensor) - if tensor.dim() == 2 - else sample_activations - ), - device=tensor.device, - ) - - if self.config.lrp_rule == "epsilon": - relevance = self.compute_relevance_epsilon( - sample_activations, tensor, output_relevance - ) - elif self.config.lrp_rule == "gamma": - relevance = self.compute_relevance_gamma( - sample_activations, tensor, output_relevance - ) - elif self.config.lrp_rule == "alpha_beta": - relevance = self.compute_relevance_alpha_beta( - sample_activations, tensor, output_relevance - ) - else: - raise ValueError(f"Unknown LRP rule: {self.config.lrp_rule}") - else: - # Fallback: use magnitude-based proxy - relevance = torch.abs(tensor) - - return relevance - - def _register_hooks(self) -> None: - """Register forward hooks to collect activations.""" - self.activations = {} - self.hooks = [] - - def get_hook(name): - def hook(module, input, output): - # Store detached tensors; move to CPU if memory is an issue - self.activations[name] = (input[0].detach(), output.detach()) - - return hook - - for name, module in self.model.named_modules(): - if isinstance(module, torch.nn.Linear): - handle = module.register_forward_hook(get_hook(name)) - self.hooks.append(handle) - - def _remove_hooks(self) -> None: - """Remove previously registered hooks.""" - for handle in self.hooks: - handle.remove() - self.hooks = [] - def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: """ - Compute relevance scores for all model weights. - - This is the main entry point for computing LRP scores. + Compute relevance scores for all model weights using true AttnLRP. """ if self.model is None: self.load_model() - print(f"Computing LRP scores using {self.config.lrp_rule} rule...") + print("Computing LRP scores using AttnLRP (via lxt)...") + + try: + from lxt.models.llama import attnlrp + # Try to handle architecture-specific registration if needed. + # For this example, we assume Llama architecture. + attnlrp.register(self.model) + except ImportError: + print("Warning: lxt not found. AttnLRP rules will not be applied.") + + # Free intermediate activations during backward — recompute instead of store + self.model.gradient_checkpointing_enable() + self.model.config.use_cache = False # required with checkpointing + + # Accumulator on CPU; only the active sample's grads live on GPU + relevance_acc = { + n: torch.zeros_like(p, device="cpu", dtype=torch.float32) + for n, p in self.model.named_parameters() if p.requires_grad + } # Tokenize sample prompts if self.config.sample_prompts: inputs = self.tokenizer( - self.config.sample_prompts, - return_tensors="pt", - padding=True, - truncation=True, - max_length=self.config.max_length, + self.config.sample_prompts, return_tensors="pt", + padding=True, truncation=True, max_length=self.config.max_length, ).to(self.config.device) - - # Get sample activations via forward pass with hooks - self._register_hooks() - with torch.no_grad(): - self.model(**inputs, output_hidden_states=False) - self._remove_hooks() else: - # No samples provided, use magnitude fallback - print("No sample prompts provided, using magnitude-based importance...") + raise ValueError("No sample prompts provided for LRP computation") - # Map module names to parameter names - # Most parameters in transformer layers follow {module_name}.weight or {module_name}.bias - parameter_to_module = {} - for mod_name, _ in self.model.named_modules(): - parameter_to_module[f"{mod_name}.weight"] = mod_name - parameter_to_module[f"{mod_name}.bias"] = mod_name + n_samples = inputs["input_ids"].shape[0] - # Compute relevance for each parameter - for name, param in self.model.named_parameters(): - if not param.requires_grad: - continue + for i in range(n_samples): # one sample at a time + print(f"Processing sample {i+1}/{n_samples}...") + ids = inputs["input_ids"][i:i+1] + embed = self.model.get_input_embeddings()(ids) + embed.requires_grad_(True) - print(f"Processing {name}...") + logits = self.model(inputs_embeds=embed).logits # full graph, checkpointed + target = logits[:, -1, :].max(dim=-1).values.sum() # seed: predicted-token logit - # Try to find captured activations for this parameter's module - module_name = parameter_to_module.get(name) - act_data = self.activations.get(module_name) - - # Pass input activations (the first element of act_data tuple) - sample_act = act_data[0] if act_data is not None else None - - # Compute relevance for this parameter - relevance = self.compute_relevance_for_tensor( - name, param.data, sample_activations=sample_act - ) + self.model.zero_grad(set_to_none=True) + target.backward() # ONE real backward pass - self.relevance_scores[name] = relevance.cpu() + with torch.no_grad(): # accumulate R_w = grad ⊙ w + for n, p in self.model.named_parameters(): + if p.grad is None: + continue + relevance_acc[n] += (p.grad.detach() * p.detach()).abs().float().cpu() + p.grad = None + torch.cuda.empty_cache() + self.relevance_scores = {n: (r / n_samples) for n, r in relevance_acc.items()} return self.relevance_scores def save_relevance_scores(self, output_format: str = "safetensors") -> None: diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index c3132e32..4362027e 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -164,7 +164,10 @@ def parameters(self) -> List[ConfigParameterDef]: ] def tensor_parameters(self) -> List[ConfigParameterDef]: - return [ConfigParameterDef(name="weight", required=False, default_value=1.0)] + return [ + ConfigParameterDef(name="weight", required=False, default_value=1.0), + ConfigParameterDef(name="lrp_scores", required=False), + ] @override def make_task( @@ -175,12 +178,12 @@ def make_task( parameters: ImmutableMap[str, Any], tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], base_model: Optional[ModelReference], - lrp_scores: Optional[Dict[str, str]] = None, **_kwargs, ) -> Task: """Create the LRP merge task with proper validation.""" # Collect model weights from non-base models model_weights = {} + lrp_scores_map = {} for model_ref, params in tensor_parameters.items(): if model_ref != base_model: try: @@ -189,6 +192,9 @@ def make_task( weight = 1.0 model_weights[model_ref] = weight + if "lrp_scores" in params and params["lrp_scores"] is not None: + lrp_scores_map[str(model_ref)] = str(params["lrp_scores"]) + if not model_weights: raise ValueError( "At least one fine-tuned model (other than base) is required for LRP merge" @@ -210,5 +216,5 @@ def make_task( model_weights=ImmutableMap(model_weights), density=density, weight_info=output_weight, - lrp_scores=ImmutableMap(lrp_scores) if lrp_scores else None, + lrp_scores=ImmutableMap(lrp_scores_map) if lrp_scores_map else None, ) diff --git a/tests/test_lrp_merge.py b/tests/test_lrp_merge.py new file mode 100644 index 00000000..f6d3719b --- /dev/null +++ b/tests/test_lrp_merge.py @@ -0,0 +1,88 @@ +import os +import torch +import tempfile +import pytest + +from mergekit.config import ( + InputModelDefinition, + MergeConfiguration, +) +from tests.common import make_picollama +from mergekit.io.lazy_tensor_loader import LazyTensorLoader +from mergekit.merge import MergeOptions, run_merge + +@pytest.fixture(scope="session") +def model_a(tmp_path_factory): + return make_picollama(tmp_path_factory.mktemp("model_a")) + +@pytest.fixture(scope="session") +def model_b(tmp_path_factory): + return make_picollama(tmp_path_factory.mktemp("model_b")) + +def test_lrp_merge_differs_from_magnitude(model_a, model_b): + """ + Test that providing LRP scores produces a different mask + than the default magnitude-based fallback. + """ + with tempfile.TemporaryDirectory() as tmpdir: + lrp_path = os.path.join(tmpdir, "lrp_scores.pt") + + loader = LazyTensorLoader.from_disk(model_b) + + lrp_scores = {} + for name in loader.index.tensor_paths: + tensor = loader.get_tensor(name) + # Use inverse of magnitude to ensure LRP mask differs from magnitude mask + lrp_scores[name] = 1.0 / (torch.abs(tensor) + 1e-6) + + torch.save(lrp_scores, lrp_path) + + # Run LRP merge + config_lrp = MergeConfiguration( + merge_method="lrp", + base_model=model_a, + models=[ + InputModelDefinition( + model=model_b, + parameters={"weight": 1.0, "lrp_scores": lrp_path}, + ), + ], + parameters={"density": 0.5}, + dtype="bfloat16", + ) + + # Run magnitude merge (no lrp_scores passed) + config_mag = MergeConfiguration( + merge_method="lrp", + base_model=model_a, + models=[ + InputModelDefinition( + model=model_b, + parameters={"weight": 1.0}, + ), + ], + parameters={"density": 0.5}, + dtype="bfloat16", + ) + + out_dir_lrp = os.path.join(tmpdir, "out_lrp") + out_dir_mag = os.path.join(tmpdir, "out_mag") + + run_merge(config_lrp, out_path=out_dir_lrp, options=MergeOptions()) + run_merge(config_mag, out_path=out_dir_mag, options=MergeOptions()) + + loader_lrp = LazyTensorLoader.from_disk(out_dir_lrp, lazy_unpickle=False) + loader_mag = LazyTensorLoader.from_disk(out_dir_mag, lazy_unpickle=False) + + diff_found = False + for name in loader_lrp.index.tensor_paths: + if name.endswith("lm_head.weight") or name.endswith("embed_tokens.weight"): + continue # some layers might not be sparseified based on architecture, ignore them or just check all + t_lrp = loader_lrp.get_tensor(name) + t_mag = loader_mag.get_tensor(name) + + if not torch.equal(t_lrp, t_mag): + diff_found = True + break + + assert diff_found, "LRP mask produced the same output as magnitude fallback!" From 3e92daf080f7f3c37c2c711288418134fa5393fa Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 13:25:20 +0530 Subject: [PATCH 03/13] Fix LRP bugs: add safetensors support, implement global cache, support multiple architectures, and enforce explicit lrp_scores --- README.md | 3 +++ lrp_computer.py | 18 +++++++++++++----- lrp_merge_pipeline.py | 3 ++- mergekit/merge_methods/lrp.py | 20 +++++++++++++++----- 4 files changed, 33 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 0bb69a50..173012ff 100644 --- a/README.md +++ b/README.md @@ -161,9 +161,11 @@ models: - model: your-finetuned-model-1 parameters: weight: 1.0 + lrp_scores: "./path/to/scores-1.safetensors" - model: your-finetuned-model-2 parameters: weight: 1.0 + lrp_scores: "./path/to/scores-2.safetensors" dtype: float16 ``` @@ -263,6 +265,7 @@ models: - model: ./models/tinyllama-domain-expert parameters: weight: 1.0 + lrp_scores: "./models/tinyllama-domain-expert/lrp_scores.safetensors" dtype: float16 ``` diff --git a/lrp_computer.py b/lrp_computer.py index 7ed1f563..0e063045 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -96,12 +96,20 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: print("Computing LRP scores using AttnLRP (via lxt)...") try: - from lxt.models.llama import attnlrp - # Try to handle architecture-specific registration if needed. - # For this example, we assume Llama architecture. - attnlrp.register(self.model) + model_type = getattr(self.model.config, "model_type", "").lower() + if "llama" in model_type: + from lxt.models.llama import attnlrp + attnlrp.register(self.model) + elif "qwen" in model_type: + from lxt.models.qwen2 import attnlrp + attnlrp.register(self.model) + elif "mistral" in model_type: + from lxt.models.mistral import attnlrp + attnlrp.register(self.model) + else: + print(f"Warning: Architecture '{model_type}' not natively supported by lxt's AttnLRP. Relevance may be plain gradient×weight.") except ImportError: - print("Warning: lxt not found. AttnLRP rules will not be applied.") + print("Warning: lxt not found or import failed. AttnLRP rules will not be applied.") # Free intermediate activations during backward — recompute instead of store self.model.gradient_checkpointing_enable() diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index eda8598c..42fe887e 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -46,16 +46,17 @@ def generate_yaml(): parameters: density: 0.7 - use_lrp: true models: - model: \"{MODEL_1}\" parameters: weight: 1.0 + lrp_scores: "./models/tinyllama-global-full/lrp_scores.safetensors" - model: \"{MODEL_2}\" parameters: weight: 1.0 + lrp_scores: "./models/tinyllama-local-full/lrp_scores.safetensors" """ with open("lrp_config.yaml", "w") as f: diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index 4362027e..49994ff7 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -17,6 +17,7 @@ from mergekit.merge_methods.rectify_embed import rectify_embed_sizes from mergekit.sparsify import build_mask +_GLOBAL_LRP_CACHE: Dict[str, Any] = {} class LRPMergeTask(Task[torch.Tensor]): """ @@ -83,7 +84,7 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: raise ValueError("Sum of model weights cannot be zero") # Process each model - _lrp_cache: Dict[str, Any] = {} + global _GLOBAL_LRP_CACHE for ref, fine_tuned_weight in weight_tensors.items(): # Validate tensor shape if fine_tuned_weight.shape != base_tensor.shape: @@ -99,19 +100,28 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: 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: - _lrp_cache[lrp_path] = torch.load(lrp_path, map_location="cpu") - importance = _lrp_cache[lrp_path].get(self.weight_info.name) + if lrp_path not in _GLOBAL_LRP_CACHE: + if lrp_path.endswith(".safetensors"): + from safetensors.torch import load_file + _GLOBAL_LRP_CACHE[lrp_path] = load_file(lrp_path, device="cpu") + else: + _GLOBAL_LRP_CACHE[lrp_path] = torch.load(lrp_path, map_location="cpu") + importance = _GLOBAL_LRP_CACHE[lrp_path].get(self.weight_info.name) if importance is not None: importance = importance.to(delta.device) # Fallback to magnitude-based importance if importance is None: + import logging + logging.warning(f"LRP scores for {self.weight_info.name} not found or not provided for {ref_str}. Falling back to magnitude.") importance = delta.abs() # Validate importance shape if importance.shape != delta.shape: - importance = delta.abs() + raise ValueError( + f"LRP score shape mismatch for {self.weight_info.name} in {ref_str}: " + f"expected {delta.shape}, got {importance.shape}" + ) # Sparsify based on importance mask = build_mask(importance, self.density) From 4ac9c1dc73a41f6c68f438cb16206682cf286192 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 13:44:11 +0530 Subject: [PATCH 04/13] Fix pipeline argument parsing, enforce strict LRP scores without fallback, add lru_cache for memory safety, and clean up unused LRP parameters --- LRP_Merge.ipynb | 4 +- lrp_computer.py | 18 +----- lrp_merge_pipeline.py | 108 +++++++++++++++------------------- mergekit/merge_methods/lrp.py | 28 +++++---- 4 files changed, 65 insertions(+), 93 deletions(-) diff --git a/LRP_Merge.ipynb b/LRP_Merge.ipynb index b87ab72f..98610cae 100644 --- a/LRP_Merge.ipynb +++ b/LRP_Merge.ipynb @@ -621,13 +621,13 @@ " parameters:\n", " weight: 1.0\n", " lrp_scores:\n", - " value: \"./models/lrp-global/lrp_scores\"\n", + " value: \"./models/lrp-global/lrp_scores.safetensors\"\n", " - model:\n", " model: \"{LOCAL_OUTPUT}\"\n", " parameters:\n", " weight: 1.0\n", " lrp_scores:\n", - " value: \"./models/lrp-local/lrp_scores\"\n", + " value: \"./models/lrp-local/lrp_scores.safetensors\"\n", "\"\"\"\n", "\n", "with open(\"lrp_config_colab.yaml\", \"w\") as f:\n", diff --git a/lrp_computer.py b/lrp_computer.py index 0e063045..6b0a6c37 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -21,19 +21,13 @@ class LRPConfig: model_path: str output_path: str sample_prompts: List[str] - batch_size: int = 1 max_length: int = 512 - lrp_rule: str = "epsilon" # "epsilon", "gamma", or "alpha_beta" - epsilon: float = 1e-9 - gamma: float = 0.25 - alpha: float = 1.0 - beta: float = 0.0 device: str = "cuda" if torch.cuda.is_available() else "cpu" class LRPComputer: """ - Computes Layer-wise Relevance Propagation scores for transformer models. + Computes Layer-wise Relevance Propagation scores for transformer models using AttnLRP. """ def __init__(self, config: LRPConfig): @@ -41,9 +35,6 @@ def __init__(self, config: LRPConfig): self.model = None self.tokenizer = None self.relevance_scores: Dict[str, torch.Tensor] = {} - self.activations: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {} - self.module_relevance: Dict[str, torch.Tensor] = {} # Store propagated relevance - self.hooks: List[Any] = [] def load_model(self) -> None: """Load the model and tokenizer.""" @@ -174,9 +165,6 @@ def save_relevance_scores(self, output_format: str = "safetensors") -> None: # Save metadata metadata = { "model_path": self.config.model_path, - "lrp_rule": self.config.lrp_rule, - "epsilon": self.config.epsilon, - "gamma": self.config.gamma, "num_tensors": len(self.relevance_scores), } @@ -230,9 +218,6 @@ def compute_lrp_for_model( parser = argparse.ArgumentParser(description="Compute LRP scores for a model") parser.add_argument("model_path", help="Path to the HuggingFace model") parser.add_argument("output_path", help="Where to save LRP scores") - parser.add_argument( - "--rule", default="epsilon", choices=["epsilon", "gamma", "alpha_beta"] - ) parser.add_argument( "--device", default="cuda" if torch.cuda.is_available() else "cpu" ) @@ -246,6 +231,5 @@ def compute_lrp_for_model( model_path=args.model_path, output_path=args.output_path, sample_prompts=args.prompts, - lrp_rule=args.rule, device=args.device, ) diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index 42fe887e..af073b2e 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -3,120 +3,106 @@ import os import shutil import subprocess -import sys +import argparse -# 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(): +def validate(args): print("Step 1: Checking local model paths...") - - for path in [MODEL_1, MODEL_2]: + for path in [args.model1, args.model2]: 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 + print(f" Warning: Missing model: {path}") - -# YAML GENERATION - - -def generate_yaml(): +def generate_yaml(args): print("Step 2: Generating YAML...") - + yaml = f""" merge_method: lrp base_model: - model: \"{BASE_MODEL}\" + model: \"{args.base_model}\" parameters: density: 0.7 models: - - model: \"{MODEL_1}\" + - model: \"{args.model1}\" parameters: weight: 1.0 - lrp_scores: "./models/tinyllama-global-full/lrp_scores.safetensors" + lrp_scores: \"{args.model1}/lrp_scores.safetensors\" - - model: \"{MODEL_2}\" + - model: \"{args.model2}\" parameters: weight: 1.0 - lrp_scores: "./models/tinyllama-local-full/lrp_scores.safetensors" + lrp_scores: \"{args.model2}/lrp_scores.safetensors\" """ - with open("lrp_config.yaml", "w") as f: f.write(yaml.strip()) print("\nYAML Generated:\n") print(yaml) - -# MERGE EXECUTION - - -def run_merge(): +def compute_lrp_scores(args): + print("\nStep 2.5: Computing LRP scores (if requested)...\n") + if not args.compute_lrp: + print("Skipping LRP computation. Make sure you already have lrp_scores.safetensors in your model directories.") + return + + print("Computing LRP scores for Model 1...") + subprocess.run([ + "python", "lrp_computer.py", args.model1, args.model1, + "--prompts", "The capital of France is" + ], check=True) + + print("Computing LRP scores for Model 2...") + subprocess.run([ + "python", "lrp_computer.py", args.model2, args.model2, + "--prompts", "The capital of France is" + ], check=True) + +def run_merge(args): 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, + args.output, "--copy-tokenizer", "--allow-crimes", ] - - if USE_CUDA: + if args.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 - + print(f"📁 Output: {args.output}") def main(): - print("=== LRP MERGE PIPELINE START ===\n") - - validate() - generate_yaml() - os.makedirs(OUTPUT_DIR, exist_ok=True) - run_merge() + parser = argparse.ArgumentParser(description="LRP Merge Pipeline") + parser.add_argument("--base-model", default="TinyLlama/TinyLlama-1.1B-Chat-v1.0", help="Base model path or name") + parser.add_argument("--model1", default="./models/tinyllama-global-full", help="First model path") + parser.add_argument("--model2", default="./models/tinyllama-local-full", help="Second model path") + parser.add_argument("--output", default="./models/merged-model", help="Output directory") + parser.add_argument("--compute-lrp", action="store_true", help="Compute LRP scores before merging") + parser.add_argument("--cuda", action="store_true", default=True, help="Use CUDA for merging") + args = parser.parse_args() + print("=== LRP MERGE PIPELINE START ===\n") + validate(args) + compute_lrp_scores(args) + generate_yaml(args) + os.makedirs(args.output, exist_ok=True) + run_merge(args) print("\n✓ ALL DONE") - if __name__ == "__main__": main() diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index 49994ff7..e75e09bc 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -14,10 +14,18 @@ MergeMethod, MergeTensorInput, ) +import functools from mergekit.merge_methods.rectify_embed import rectify_embed_sizes from mergekit.sparsify import build_mask -_GLOBAL_LRP_CACHE: Dict[str, Any] = {} +@functools.lru_cache(maxsize=4) +def _load_lrp_scores(lrp_path: str): + if lrp_path.endswith(".safetensors"): + from safetensors.torch import load_file + return load_file(lrp_path, device="cpu") + else: + import torch + return torch.load(lrp_path, map_location="cpu") class LRPMergeTask(Task[torch.Tensor]): """ @@ -84,7 +92,6 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: raise ValueError("Sum of model weights cannot be zero") # Process each model - global _GLOBAL_LRP_CACHE for ref, fine_tuned_weight in weight_tensors.items(): # Validate tensor shape if fine_tuned_weight.shape != base_tensor.shape: @@ -100,21 +107,16 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: 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 _GLOBAL_LRP_CACHE: - if lrp_path.endswith(".safetensors"): - from safetensors.torch import load_file - _GLOBAL_LRP_CACHE[lrp_path] = load_file(lrp_path, device="cpu") - else: - _GLOBAL_LRP_CACHE[lrp_path] = torch.load(lrp_path, map_location="cpu") - importance = _GLOBAL_LRP_CACHE[lrp_path].get(self.weight_info.name) + scores_dict = _load_lrp_scores(lrp_path) + importance = scores_dict.get(self.weight_info.name) if importance is not None: importance = importance.to(delta.device) - # Fallback to magnitude-based importance + # Strict LRP: No silent fallback to magnitude if importance is None: - import logging - logging.warning(f"LRP scores for {self.weight_info.name} not found or not provided for {ref_str}. Falling back to magnitude.") - importance = delta.abs() + raise RuntimeError( + f"LRP scores for tensor '{self.weight_info.name}' not found or not provided for {ref_str}." + ) # Validate importance shape if importance.shape != delta.shape: From a899412b5ef43a2a7539210b036ec245a9b395f9 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 13:52:48 +0530 Subject: [PATCH 05/13] Fix test fallback logic, use dynamic mergekit-yaml path, and cleanly restore model state after LRP backward pass --- lrp_computer.py | 6 ++++++ lrp_merge_pipeline.py | 4 +--- tests/test_lrp_merge.py | 25 +++++++++++++++---------- 3 files changed, 22 insertions(+), 13 deletions(-) diff --git a/lrp_computer.py b/lrp_computer.py index 6b0a6c37..4c2685e7 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -104,6 +104,7 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: # Free intermediate activations during backward — recompute instead of store self.model.gradient_checkpointing_enable() + orig_use_cache = getattr(self.model.config, "use_cache", True) self.model.config.use_cache = False # required with checkpointing # Accumulator on CPU; only the active sample's grads live on GPU @@ -143,6 +144,11 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: p.grad = None torch.cuda.empty_cache() + # Restore original settings + if hasattr(self.model, "gradient_checkpointing_disable"): + self.model.gradient_checkpointing_disable() + self.model.config.use_cache = orig_use_cache + self.relevance_scores = {n: (r / n_samples) for n, r in relevance_acc.items()} return self.relevance_scores diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index af073b2e..764b7829 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -62,9 +62,7 @@ def compute_lrp_scores(args): def run_merge(args): print("\nStep 3: Running merge...\n") - mergekit_exec = "/usr/local/bin/mergekit-yaml" - if not os.path.exists(mergekit_exec): - mergekit_exec = shutil.which("mergekit-yaml") + mergekit_exec = shutil.which("mergekit-yaml") if not mergekit_exec: raise RuntimeError("✗ mergekit-yaml not found. Install mergekit.") diff --git a/tests/test_lrp_merge.py b/tests/test_lrp_merge.py index f6d3719b..f1c0cdee 100644 --- a/tests/test_lrp_merge.py +++ b/tests/test_lrp_merge.py @@ -22,43 +22,48 @@ def model_b(tmp_path_factory): def test_lrp_merge_differs_from_magnitude(model_a, model_b): """ Test that providing LRP scores produces a different mask - than the default magnitude-based fallback. + than magnitude-based importance. """ with tempfile.TemporaryDirectory() as tmpdir: - lrp_path = os.path.join(tmpdir, "lrp_scores.pt") + lrp_path_inv = os.path.join(tmpdir, "lrp_scores_inv.pt") + lrp_path_mag = os.path.join(tmpdir, "lrp_scores_mag.pt") loader = LazyTensorLoader.from_disk(model_b) - lrp_scores = {} + lrp_scores_inv = {} + lrp_scores_mag = {} for name in loader.index.tensor_paths: tensor = loader.get_tensor(name) - # Use inverse of magnitude to ensure LRP mask differs from magnitude mask - lrp_scores[name] = 1.0 / (torch.abs(tensor) + 1e-6) + # Inverse of magnitude to ensure LRP mask differs from magnitude + lrp_scores_inv[name] = 1.0 / (torch.abs(tensor) + 1e-6) + # Direct magnitude + lrp_scores_mag[name] = torch.abs(tensor) - torch.save(lrp_scores, lrp_path) + torch.save(lrp_scores_inv, lrp_path_inv) + torch.save(lrp_scores_mag, lrp_path_mag) - # Run LRP merge + # Run LRP merge with inverse scores config_lrp = MergeConfiguration( merge_method="lrp", base_model=model_a, models=[ InputModelDefinition( model=model_b, - parameters={"weight": 1.0, "lrp_scores": lrp_path}, + parameters={"weight": 1.0, "lrp_scores": lrp_path_inv}, ), ], parameters={"density": 0.5}, dtype="bfloat16", ) - # Run magnitude merge (no lrp_scores passed) + # Run LRP merge with magnitude scores (simulating the old fallback) config_mag = MergeConfiguration( merge_method="lrp", base_model=model_a, models=[ InputModelDefinition( model=model_b, - parameters={"weight": 1.0}, + parameters={"weight": 1.0, "lrp_scores": lrp_path_mag}, ), ], parameters={"density": 0.5}, From d6e8dd1e29a5af0cf0db35129233cdb91082f0d3 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 14:19:49 +0530 Subject: [PATCH 06/13] Guarantee gradient checkpointing restoration via try/finally and add lxt to pyproject extras --- lrp_computer.py | 120 ++++++++++++++++++++++++++++++------------------ pyproject.toml | 1 + 2 files changed, 76 insertions(+), 45 deletions(-) diff --git a/lrp_computer.py b/lrp_computer.py index 4c2685e7..5788d786 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -23,6 +23,7 @@ class LRPConfig: sample_prompts: List[str] max_length: int = 512 device: str = "cuda" if torch.cuda.is_available() else "cpu" + load_in_4bit: bool = False class LRPComputer: @@ -52,11 +53,22 @@ def load_model(self) -> None: print(f" Using device: {self.config.device}") print(f" Using dtype: {torch_dtype}") + + quant_config = None + if self.config.load_in_4bit: + from transformers import BitsAndBytesConfig + quant_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch_dtype, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + ) try: self.model = AutoModelForCausalLM.from_pretrained( self.config.model_path, torch_dtype=torch_dtype, + quantization_config=quant_config, device_map=device_map, low_cpu_mem_usage=True, ) @@ -66,6 +78,7 @@ def load_model(self) -> None: self.model = AutoModelForCausalLM.from_pretrained( self.config.model_path, torch_dtype=torch_dtype, + quantization_config=quant_config, device_map=device_map, low_cpu_mem_usage=True, trust_remote_code=True, @@ -84,8 +97,6 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: if self.model is None: self.load_model() - print("Computing LRP scores using AttnLRP (via lxt)...") - try: model_type = getattr(self.model.config, "model_type", "").lower() if "llama" in model_type: @@ -98,56 +109,71 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: from lxt.models.mistral import attnlrp attnlrp.register(self.model) else: - print(f"Warning: Architecture '{model_type}' not natively supported by lxt's AttnLRP. Relevance may be plain gradient×weight.") + raise ValueError( + f"AttnLRP not supported for model_type={model_type!r}. " + f"Currently supported: llama, qwen, mistral. " + f"For other architectures, contribute an lxt rules module." + ) except ImportError: - print("Warning: lxt not found or import failed. AttnLRP rules will not be applied.") + raise ImportError("lxt is required for AttnLRP. Install with: pip install lxt") from None # Free intermediate activations during backward — recompute instead of store self.model.gradient_checkpointing_enable() orig_use_cache = getattr(self.model.config, "use_cache", True) self.model.config.use_cache = False # required with checkpointing - # Accumulator on CPU; only the active sample's grads live on GPU - relevance_acc = { - n: torch.zeros_like(p, device="cpu", dtype=torch.float32) - for n, p in self.model.named_parameters() if p.requires_grad - } - - # Tokenize sample prompts - if self.config.sample_prompts: - inputs = self.tokenizer( - self.config.sample_prompts, return_tensors="pt", - padding=True, truncation=True, max_length=self.config.max_length, - ).to(self.config.device) - else: - raise ValueError("No sample prompts provided for LRP computation") - - n_samples = inputs["input_ids"].shape[0] - - for i in range(n_samples): # one sample at a time - print(f"Processing sample {i+1}/{n_samples}...") - ids = inputs["input_ids"][i:i+1] - embed = self.model.get_input_embeddings()(ids) - embed.requires_grad_(True) - - logits = self.model(inputs_embeds=embed).logits # full graph, checkpointed - target = logits[:, -1, :].max(dim=-1).values.sum() # seed: predicted-token logit - - self.model.zero_grad(set_to_none=True) - target.backward() # ONE real backward pass - - with torch.no_grad(): # accumulate R_w = grad ⊙ w - for n, p in self.model.named_parameters(): - if p.grad is None: - continue - relevance_acc[n] += (p.grad.detach() * p.detach()).abs().float().cpu() - p.grad = None - torch.cuda.empty_cache() - - # Restore original settings - if hasattr(self.model, "gradient_checkpointing_disable"): - self.model.gradient_checkpointing_disable() - self.model.config.use_cache = orig_use_cache + try: + # Accumulator on CPU; only the active sample's grads live on GPU + relevance_acc = { + n: torch.zeros_like(p, device="cpu", dtype=torch.float32) + for n, p in self.model.named_parameters() if p.requires_grad + } + + # Tokenize sample prompts + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + if self.config.sample_prompts: + inputs = self.tokenizer( + self.config.sample_prompts, return_tensors="pt", + padding=True, truncation=True, max_length=self.config.max_length, + ).to(self.config.device) + else: + raise ValueError("No sample prompts provided for LRP computation") + + n_samples = inputs["input_ids"].shape[0] + + for i in range(n_samples): # one sample at a time + print(f"Processing sample {i+1}/{n_samples}...") + ids = inputs["input_ids"][i:i+1] + attention_mask = inputs["attention_mask"][i:i+1] + + embed = self.model.get_input_embeddings()(ids) + embed.requires_grad_(True) + + logits = self.model(inputs_embeds=embed, attention_mask=attention_mask).logits # full graph, checkpointed + + # Find the actual last token position before padding + last_token_idx = attention_mask.sum().item() - 1 + target = logits[:, last_token_idx, :].max(dim=-1).values.sum() # seed: predicted-token logit + + self.model.zero_grad(set_to_none=True) + target.backward() # ONE real backward pass + + with torch.no_grad(): # accumulate R_w = grad ⊙ w + for n, p in self.model.named_parameters(): + if p.grad is None: + continue + relevance_acc[n] += (p.grad.detach() * p.detach()).abs().float().cpu() + + self.model.zero_grad(set_to_none=True) + torch.cuda.empty_cache() + + finally: + # Restore original settings + if hasattr(self.model, "gradient_checkpointing_disable"): + self.model.gradient_checkpointing_disable() + self.model.config.use_cache = orig_use_cache self.relevance_scores = {n: (r / n_samples) for n, r in relevance_acc.items()} return self.relevance_scores @@ -230,6 +256,9 @@ def compute_lrp_for_model( parser.add_argument( "--prompts", nargs="+", help="Sample prompts for LRP computation" ) + parser.add_argument( + "--load-in-4bit", action="store_true", help="Load the model in 4-bit (NF4) for lower memory usage" + ) args = parser.parse_args() @@ -238,4 +267,5 @@ def compute_lrp_for_model( output_path=args.output_path, sample_prompts=args.prompts, device=args.device, + load_in_4bit=args.load_in_4bit, ) diff --git a/pyproject.toml b/pyproject.toml index 05136aaa..2efde0d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dev = ["black~=25.1.0", "isort~=6.0.1", "pre-commit~=4.2.0"] test = ["pytest~=8.4.0"] evolve = ["ray", "cma", "lm_eval", "wandb"] vllm = ["vllm==0.7.2", "lm_eval[vllm]"] +lxt = ["lxt"] [project.urls] repository = "https://github.com/cg123/mergekit" From ae32700318efb22c0728049c7f2a2767ecdcc6ea Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 14:29:23 +0530 Subject: [PATCH 07/13] Fix build_mask import, correctly trace embedding gradients in LRP, isolate LRP score output directories, and patch notebook paths --- LRP_Merge.ipynb | 2 +- lrp_computer.py | 5 +---- lrp_merge_pipeline.py | 16 ++++++++++++---- mergekit/merge_methods/lrp.py | 18 ++++++++++++++++-- 4 files changed, 30 insertions(+), 11 deletions(-) diff --git a/LRP_Merge.ipynb b/LRP_Merge.ipynb index 98610cae..d4c0a4ca 100644 --- a/LRP_Merge.ipynb +++ b/LRP_Merge.ipynb @@ -965,7 +965,7 @@ "import subprocess\n", "\n", "def run_final_merge():\n", - " config_path = \"lrp_config.yaml\"\n", + " config_path = \"lrp_config_colab.yaml\"\n", " output_dir = os.path.join(LRP_PATH, \"models/merged-model\")\n", " os.makedirs(output_dir, exist_ok=True)\n", "\n", diff --git a/lrp_computer.py b/lrp_computer.py index 5788d786..becc291d 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -148,10 +148,7 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: ids = inputs["input_ids"][i:i+1] attention_mask = inputs["attention_mask"][i:i+1] - embed = self.model.get_input_embeddings()(ids) - embed.requires_grad_(True) - - logits = self.model(inputs_embeds=embed, attention_mask=attention_mask).logits # full graph, checkpointed + logits = self.model(input_ids=ids, attention_mask=attention_mask).logits # full graph, checkpointed # Find the actual last token position before padding last_token_idx = attention_mask.sum().item() - 1 diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index 764b7829..d1151332 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -16,6 +16,9 @@ def validate(args): def generate_yaml(args): print("Step 2: Generating YAML...") + m1_score_path = os.path.join(args.output, "lrp_scores", "model1", "lrp_scores.safetensors") if args.compute_lrp else f"{args.model1}/lrp_scores.safetensors" + m2_score_path = os.path.join(args.output, "lrp_scores", "model2", "lrp_scores.safetensors") if args.compute_lrp else f"{args.model2}/lrp_scores.safetensors" + yaml = f""" merge_method: lrp @@ -29,12 +32,12 @@ def generate_yaml(args): - model: \"{args.model1}\" parameters: weight: 1.0 - lrp_scores: \"{args.model1}/lrp_scores.safetensors\" + lrp_scores: \"{m1_score_path}\" - model: \"{args.model2}\" parameters: weight: 1.0 - lrp_scores: \"{args.model2}/lrp_scores.safetensors\" + lrp_scores: \"{m2_score_path}\" """ with open("lrp_config.yaml", "w") as f: f.write(yaml.strip()) @@ -48,15 +51,20 @@ def compute_lrp_scores(args): print("Skipping LRP computation. Make sure you already have lrp_scores.safetensors in your model directories.") return + m1_out = os.path.join(args.output, "lrp_scores", "model1") + m2_out = os.path.join(args.output, "lrp_scores", "model2") + os.makedirs(m1_out, exist_ok=True) + os.makedirs(m2_out, exist_ok=True) + print("Computing LRP scores for Model 1...") subprocess.run([ - "python", "lrp_computer.py", args.model1, args.model1, + "python", "lrp_computer.py", args.model1, m1_out, "--prompts", "The capital of France is" ], check=True) print("Computing LRP scores for Model 2...") subprocess.run([ - "python", "lrp_computer.py", args.model2, args.model2, + "python", "lrp_computer.py", args.model2, m2_out, "--prompts", "The capital of France is" ], check=True) diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index e75e09bc..49a4c817 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -16,7 +16,21 @@ ) import functools from mergekit.merge_methods.rectify_embed import rectify_embed_sizes -from mergekit.sparsify import build_mask + +def build_mask(importance: torch.Tensor, density: float) -> torch.Tensor: + """Creates a binary mask keeping the top proportion (density) of elements.""" + if density >= 1.0: + return torch.ones_like(importance) + k = int(density * importance.numel()) + if k <= 0: + return torch.zeros_like(importance) + w = importance.view(-1) + if w.device.type == "cpu": + w = w.float() + topk = torch.argsort(w, descending=True)[:k] + mask = torch.zeros_like(importance) + mask.view(-1)[topk] = 1 + return mask @functools.lru_cache(maxsize=4) def _load_lrp_scores(lrp_path: str): @@ -154,7 +168,7 @@ class LRPMerge(MergeMethod): Merges fine-tuned models by: 1. Computing task vectors (deltas from base) 2. Using LRP importance scores to determine which weights are most relevant - 3. Sparsifying based on importance (LRP scores or magnitude fallback) + 3. Sparsifying based on importance (strictly requires LRP scores) 4. Weighted averaging of sparse deltas """ From 26571b8cb17a3f732cd317d5e63ec700a0fa1cde Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 14:38:06 +0530 Subject: [PATCH 08/13] Fix tied embeddings in LRP, remove incompatible 4-bit loading, fix left-padded attention masks, and fix pipeline device arguments --- lrp_computer.py | 27 ++++++++------------------- lrp_merge_pipeline.py | 4 +++- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/lrp_computer.py b/lrp_computer.py index becc291d..73e6c40b 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -23,7 +23,6 @@ class LRPConfig: sample_prompts: List[str] max_length: int = 512 device: str = "cuda" if torch.cuda.is_available() else "cpu" - load_in_4bit: bool = False class LRPComputer: @@ -53,22 +52,11 @@ def load_model(self) -> None: print(f" Using device: {self.config.device}") print(f" Using dtype: {torch_dtype}") - - quant_config = None - if self.config.load_in_4bit: - from transformers import BitsAndBytesConfig - quant_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_compute_dtype=torch_dtype, - bnb_4bit_quant_type="nf4", - bnb_4bit_use_double_quant=True, - ) try: self.model = AutoModelForCausalLM.from_pretrained( self.config.model_path, torch_dtype=torch_dtype, - quantization_config=quant_config, device_map=device_map, low_cpu_mem_usage=True, ) @@ -78,7 +66,6 @@ def load_model(self) -> None: self.model = AutoModelForCausalLM.from_pretrained( self.config.model_path, torch_dtype=torch_dtype, - quantization_config=quant_config, device_map=device_map, low_cpu_mem_usage=True, trust_remote_code=True, @@ -150,8 +137,8 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: logits = self.model(input_ids=ids, attention_mask=attention_mask).logits # full graph, checkpointed - # Find the actual last token position before padding - last_token_idx = attention_mask.sum().item() - 1 + # Find the actual last token position before padding (handles both left and right padding) + last_token_idx = attention_mask[0].nonzero(as_tuple=True)[0][-1].item() target = logits[:, last_token_idx, :].max(dim=-1).values.sum() # seed: predicted-token logit self.model.zero_grad(set_to_none=True) @@ -173,6 +160,12 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: self.model.config.use_cache = orig_use_cache self.relevance_scores = {n: (r / n_samples) for n, r in relevance_acc.items()} + + # Handle tied embeddings: copy embed_tokens score to lm_head if missing + if getattr(self.model.config, "tie_word_embeddings", False): + if "model.embed_tokens.weight" in self.relevance_scores and "lm_head.weight" not in self.relevance_scores: + self.relevance_scores["lm_head.weight"] = self.relevance_scores["model.embed_tokens.weight"] + return self.relevance_scores def save_relevance_scores(self, output_format: str = "safetensors") -> None: @@ -253,9 +246,6 @@ def compute_lrp_for_model( parser.add_argument( "--prompts", nargs="+", help="Sample prompts for LRP computation" ) - parser.add_argument( - "--load-in-4bit", action="store_true", help="Load the model in 4-bit (NF4) for lower memory usage" - ) args = parser.parse_args() @@ -264,5 +254,4 @@ def compute_lrp_for_model( output_path=args.output_path, sample_prompts=args.prompts, device=args.device, - load_in_4bit=args.load_in_4bit, ) diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index d1151332..22c96d42 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -99,9 +99,11 @@ def main(): parser.add_argument("--model2", default="./models/tinyllama-local-full", help="Second model path") parser.add_argument("--output", default="./models/merged-model", help="Output directory") parser.add_argument("--compute-lrp", action="store_true", help="Compute LRP scores before merging") - parser.add_argument("--cuda", action="store_true", default=True, help="Use CUDA for merging") + parser.add_argument("--device", choices=["cuda", "cpu"], default="cuda", help="Device to use for merging") args = parser.parse_args() + args.cuda = args.device == "cuda" + print("=== LRP MERGE PIPELINE START ===\n") validate(args) compute_lrp_scores(args) From 2022d3786f21a0eda140bb30aafa6a37aa807791 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 14:50:02 +0530 Subject: [PATCH 09/13] fix: Address Cursor Bugbot issues - CUDA fallback, Qwen matching, subprocess interpreter - Fix CUDA device fallback: Normalize device selection in LRPConfig.__post_init__ to gracefully fall back to CPU when CUDA is unavailable - Fix Qwen architecture dispatch: Use exact model_type matching instead of substring to avoid false positives (qwen3, qwen2_moe) - Fix subprocess interpreter: Use sys.executable instead of bare 'python' to ensure correct venv/conda environment --- lrp_computer.py | 20 +++++++++++++++----- lrp_merge_pipeline.py | 5 +++-- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/lrp_computer.py b/lrp_computer.py index 73e6c40b..74dc2db0 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -22,7 +22,15 @@ class LRPConfig: output_path: str sample_prompts: List[str] max_length: int = 512 - device: str = "cuda" if torch.cuda.is_available() else "cpu" + device: str = None # Will be normalized in __post_init__ + + def __post_init__(self): + """Normalize device to ensure CUDA availability.""" + if self.device is None: + self.device = "cuda" if torch.cuda.is_available() else "cpu" + elif self.device == "cuda" and not torch.cuda.is_available(): + print(f"⚠️ CUDA requested but not available. Falling back to CPU.") + self.device = "cpu" class LRPComputer: @@ -86,19 +94,21 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: try: model_type = getattr(self.model.config, "model_type", "").lower() - if "llama" in model_type: + # Use exact matching to avoid false positives (e.g., qwen3, qwen2_moe) + if model_type == "llama": from lxt.models.llama import attnlrp attnlrp.register(self.model) - elif "qwen" in model_type: + elif model_type == "qwen" or model_type.startswith("qwen2"): + # Qwen2 and Qwen2.5 use the same architecture from lxt.models.qwen2 import attnlrp attnlrp.register(self.model) - elif "mistral" in model_type: + elif model_type == "mistral": from lxt.models.mistral import attnlrp attnlrp.register(self.model) else: raise ValueError( f"AttnLRP not supported for model_type={model_type!r}. " - f"Currently supported: llama, qwen, mistral. " + f"Currently supported: llama, qwen, qwen2, qwen2.5, mistral. " f"For other architectures, contribute an lxt rules module." ) except ImportError: diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index 22c96d42..f6580032 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import os +import sys import shutil import subprocess import argparse @@ -58,13 +59,13 @@ def compute_lrp_scores(args): print("Computing LRP scores for Model 1...") subprocess.run([ - "python", "lrp_computer.py", args.model1, m1_out, + sys.executable, "lrp_computer.py", args.model1, m1_out, "--prompts", "The capital of France is" ], check=True) print("Computing LRP scores for Model 2...") subprocess.run([ - "python", "lrp_computer.py", args.model2, m2_out, + sys.executable, "lrp_computer.py", args.model2, m2_out, "--prompts", "The capital of France is" ], check=True) From 9c039e6d40b7160db5b080f9825d333cad9a7ab7 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 15:05:09 +0530 Subject: [PATCH 10/13] fix: Address 4 Cursor Bugbot issues - weight normalization, cache mutation, float16 underflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix weight normalization: Remove total_weight division to respect user-supplied weights and match documented merge formula (Masked task vector × weight) - Fix cached LRP scores: Return deep copy from _load_lrp_scores to prevent cross-task tensor mutation and aliasing - Fix float16 underflow: Always use float32 for model loading to prevent gradient underflow in backward pass (float16 min normal ~6e-5 causes small gradients to zero) - Add warning message about float32 requirement for accurate LRP computation --- lrp_computer.py | 14 +++++--------- mergekit/merge_methods/lrp.py | 13 ++++++++----- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/lrp_computer.py b/lrp_computer.py index 74dc2db0..a9dd9c03 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -49,17 +49,13 @@ def load_model(self) -> None: print(f"Loading model from {self.config.model_path}...") self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path) - # Determine dtype based on device - has_cuda = torch.cuda.is_available() and self.config.device == "cuda" - if has_cuda: - torch_dtype = torch.float16 - device_map = self.config.device - else: - torch_dtype = torch.float32 - device_map = None # device_map not recommended for CPU + # For LRP, we need float32 for accurate gradient computation + # Float16 gradients underflow and produce zeros for small-magnitude parameters + torch_dtype = torch.float32 + device_map = self.config.device if self.config.device == "cuda" else None print(f" Using device: {self.config.device}") - print(f" Using dtype: {torch_dtype}") + print(f" Using dtype: {torch_dtype} (required for accurate LRP gradient computation)") try: self.model = AutoModelForCausalLM.from_pretrained( diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index 49a4c817..73e820e7 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -34,12 +34,16 @@ def build_mask(importance: torch.Tensor, density: float) -> torch.Tensor: @functools.lru_cache(maxsize=4) def _load_lrp_scores(lrp_path: str): + """Load LRP scores from disk. Returns a deep copy to prevent cache mutation.""" if lrp_path.endswith(".safetensors"): from safetensors.torch import load_file - return load_file(lrp_path, device="cpu") + scores = load_file(lrp_path, device="cpu") else: import torch - return torch.load(lrp_path, map_location="cpu") + scores = torch.load(lrp_path, map_location="cpu") + + # Return a deep copy to prevent mutations of cached tensors + return {k: v.clone() for k, v in scores.items()} class LRPMergeTask(Task[torch.Tensor]): """ @@ -143,10 +147,9 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: mask = build_mask(importance, self.density) sparse_delta = delta * mask - # Weighted averaging + # Apply per-model weight (no normalization; each model contributes its full weight) weight = self.model_weights[ref] if ref in self.model_weights else 1.0 - normalized_weight = weight / total_weight - merged_deltas += normalized_weight * sparse_delta + merged_deltas += weight * sparse_delta # Final merged tensor return base_tensor + merged_deltas From a928e72d3d226c5daf2864f60a8966684a3dccde Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 15:25:20 +0530 Subject: [PATCH 11/13] fix: Remove misleading dead code - unused total_weight validation The total_weight computation and zero-check were dead code after removing weight normalization. This misleading validation could confuse future maintainers into 'fixing' the apparent bug by re-introducing normalization, silently changing merge semantics. Removed: - total_weight = sum(self.model_weights.values()) - Zero-check validation The merge now clearly uses un-normalized per-model weights as documented. --- mergekit/merge_methods/lrp.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index 73e820e7..617d7f64 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -105,10 +105,6 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: if not self.model_weights: raise ValueError("model_weights cannot be empty") - total_weight = sum(self.model_weights.values()) - if total_weight == 0: - raise ValueError("Sum of model weights cannot be zero") - # Process each model for ref, fine_tuned_weight in weight_tensors.items(): # Validate tensor shape From cfcfc829e6afead8bcd07adb6b266673fbc5b089 Mon Sep 17 00:00:00 2001 From: Tusm11 Date: Sun, 26 Apr 2026 15:32:07 +0530 Subject: [PATCH 12/13] fix: Address 2 pipeline issues - fail fast on missing paths, resolve script location - Fix silent warnings: validate() now raises FileNotFoundError immediately if model paths don't exist, preventing confusing subprocess errors later - Fix script path resolution: compute_lrp_scores() now resolves lrp_computer.py relative to the pipeline script's directory using __file__, allowing the pipeline to be run from any working directory - Add clear error messages with actionable guidance for both issues --- lrp_merge_pipeline.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/lrp_merge_pipeline.py b/lrp_merge_pipeline.py index f6580032..08463ce8 100644 --- a/lrp_merge_pipeline.py +++ b/lrp_merge_pipeline.py @@ -7,12 +7,23 @@ import argparse def validate(args): + """Validate that all required model paths exist. Fail fast with clear error.""" print("Step 1: Checking local model paths...") + missing_paths = [] + for path in [args.model1, args.model2]: if os.path.exists(path): - print(f"Found: {path}") + print(f"✓ Found: {path}") else: - print(f" Warning: Missing model: {path}") + print(f"✗ Missing: {path}") + missing_paths.append(path) + + if missing_paths: + raise FileNotFoundError( + f"Required model paths not found:\n" + + "\n".join(f" - {p}" for p in missing_paths) + + "\n\nPlease ensure models exist or provide correct paths via --model1 and --model2" + ) def generate_yaml(args): print("Step 2: Generating YAML...") @@ -47,10 +58,21 @@ def generate_yaml(args): print(yaml) def compute_lrp_scores(args): + """Compute LRP scores for both models using the correct script path.""" print("\nStep 2.5: Computing LRP scores (if requested)...\n") if not args.compute_lrp: print("Skipping LRP computation. Make sure you already have lrp_scores.safetensors in your model directories.") return + + # Resolve lrp_computer.py relative to this script's directory + script_dir = os.path.dirname(os.path.abspath(__file__)) + lrp_computer_path = os.path.join(script_dir, "lrp_computer.py") + + if not os.path.exists(lrp_computer_path): + raise FileNotFoundError( + f"lrp_computer.py not found at {lrp_computer_path}\n" + f"Expected it next to {__file__}" + ) m1_out = os.path.join(args.output, "lrp_scores", "model1") m2_out = os.path.join(args.output, "lrp_scores", "model2") @@ -59,13 +81,13 @@ def compute_lrp_scores(args): print("Computing LRP scores for Model 1...") subprocess.run([ - sys.executable, "lrp_computer.py", args.model1, m1_out, + sys.executable, lrp_computer_path, args.model1, m1_out, "--prompts", "The capital of France is" ], check=True) print("Computing LRP scores for Model 2...") subprocess.run([ - sys.executable, "lrp_computer.py", args.model2, m2_out, + sys.executable, lrp_computer_path, args.model2, m2_out, "--prompts", "The capital of France is" ], check=True) From ef76d4b546077b3127ffa7f51dbe30b18bdc7756 Mon Sep 17 00:00:00 2001 From: ManniX-ITA <20623405+mann1x@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:36:33 +0200 Subject: [PATCH 13/13] fix: multimodal-LM (Qwen3.5_5) support for ex-LRP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit End-to-end fixes that make the ex-LRP method in this PR work against multimodal models like Qwen3_5ForConditionalGeneration. Validated on Qwen3.5-4B (jackrong-v2 + crow-4b → base) producing a clean Q6_K with HumanEval pass@1 = 51.22% and MBPP pass@1 = 49.40%. Five minimal patches: * mergekit/architecture/base.py Pydantic v2 forward-references in ConfiguredModuleArchitecture and ConfiguredModelArchitecture aren't resolved eagerly; add model_rebuild() at module load. Without this, the first model load fails with PydanticUserError: not fully defined. * mergekit/architecture/auto.py Make 'optional' layer-aware. Hybrid-attention archs (Qwen3.5_5 alternates full / linear attention) have tensors like dt_bias only in some layers. The original _wi flagged optional based on a single layer's presence; widen to true if missing in *any* layer. * mergekit/merge_methods/lrp.py Replace strict raise with base-passthrough when LRP scores are missing for a tensor. Multimodal LRP is computed only on the language_model branch; vision tower / MTP heads have no relevance signal and should retain base weights instead of failing the merge. * mergekit/config.py Allow `str` in the ParameterSetting union so per-source params like `lrp_scores: \"/path/to/scores.safetensors\"` parse cleanly via mergekit-yaml. * lrp_computer.py - Recognize qwen3_5_text inner-LM model_type (the inner LM of Qwen3_5ForConditionalGeneration) and dispatch AttnLRP against it. - Clone tied tensors (lm_head.weight ↔ model.embed_tokens.weight) by data_ptr() before save_file, since safetensors save rejects shared-storage tensors. Reproducer + LRP signal artefacts published at https://huggingface.co/ManniX-ITA/Qwen3.5-4B-M4-ex-LRP Co-Authored-By: Claude Opus 4.7 (1M context) --- lrp_computer.py | 105 +++++++++++++++++++++++++++------- mergekit/architecture/auto.py | 24 ++++++-- mergekit/architecture/base.py | 10 ++++ mergekit/config.py | 2 +- mergekit/merge_methods/lrp.py | 13 +++-- 5 files changed, 124 insertions(+), 30 deletions(-) diff --git a/lrp_computer.py b/lrp_computer.py index a9dd9c03..c8cdc454 100644 --- a/lrp_computer.py +++ b/lrp_computer.py @@ -57,6 +57,12 @@ def load_model(self) -> None: print(f" Using device: {self.config.device}") print(f" Using dtype: {torch_dtype} (required for accurate LRP gradient computation)") + # Some archs (e.g. Qwen3_5ForConditionalGeneration — multimodal w/ a + # generation head) are not in MODEL_FOR_CAUSAL_LM_MAPPING. Fall back to + # explicit class lookup by `architectures[0]` after AutoModel and + # trust_remote_code paths both fail. Forward with text-only input still + # gives LM-head logits — vision tower stays unused. + from transformers import AutoConfig try: self.model = AutoModelForCausalLM.from_pretrained( self.config.model_path, @@ -65,15 +71,29 @@ def load_model(self) -> None: low_cpu_mem_usage=True, ) except Exception as e: - print(f" Failed with default settings: {e}") - print(" Trying with trust_remote_code=True...") - self.model = AutoModelForCausalLM.from_pretrained( - self.config.model_path, - torch_dtype=torch_dtype, - device_map=device_map, - low_cpu_mem_usage=True, - trust_remote_code=True, - ) + print(f" AutoModelForCausalLM default failed: {e}") + try: + print(" Trying with trust_remote_code=True...") + self.model = AutoModelForCausalLM.from_pretrained( + self.config.model_path, + torch_dtype=torch_dtype, + device_map=device_map, + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + except Exception as e2: + print(f" trust_remote_code path failed: {e2}") + cfg = AutoConfig.from_pretrained(self.config.model_path) + arch = cfg.architectures[0] if cfg.architectures else "" + print(f" Falling back to explicit class: {arch}") + import importlib + cls = getattr(importlib.import_module("transformers"), arch) + self.model = cls.from_pretrained( + self.config.model_path, + torch_dtype=torch_dtype, + device_map=device_map, + low_cpu_mem_usage=True, + ) # Explicitly move to CPU if needed if self.config.device == "cpu": @@ -90,25 +110,56 @@ def compute_all_relevance_scores(self) -> Dict[str, torch.Tensor]: try: model_type = getattr(self.model.config, "model_type", "").lower() - # Use exact matching to avoid false positives (e.g., qwen3, qwen2_moe) + # NOTE: PR #682 was written against an older lxt that exposed + # `lxt.models..attnlrp.register(model)`. Current published lxt + # (rachtibat/LRP-eXplains-Transformers) only exposes + # `lxt.efficient.monkey_patch(modeling_module, attnLRP_dict)` — + # the first positional arg is the transformers modeling MODULE, + # the second is the per-arch patch dict. The dict alone won't + # work; the module is needed because patch_attention installs + # the wrapped attention onto module-level globals. + from lxt.efficient import monkey_patch if model_type == "llama": - from lxt.models.llama import attnlrp - attnlrp.register(self.model) + from transformers.models.llama import modeling_llama + from lxt.efficient.models.llama import attnLRP + monkey_patch(modeling_llama, attnLRP) elif model_type == "qwen" or model_type.startswith("qwen2"): - # Qwen2 and Qwen2.5 use the same architecture - from lxt.models.qwen2 import attnlrp - attnlrp.register(self.model) + # Qwen2 / Qwen2.5 share the same modeling module + from transformers.models.qwen2 import modeling_qwen2 + from lxt.efficient.models.qwen2 import attnLRP + monkey_patch(modeling_qwen2, attnLRP) + elif model_type == "qwen3": + from transformers.models.qwen3 import modeling_qwen3 + from lxt.efficient.models.qwen3 import attnLRP + monkey_patch(modeling_qwen3, attnLRP) + elif model_type == "qwen3_5" or model_type == "qwen3_5_text": + # `qwen3_5_text` is the inner text branch when the multimodal + # Qwen3_5ForConditionalGeneration is loaded via explicit class — + # same modeling_qwen3_5.py, so the same attnLRP rules apply. + from transformers.models.qwen3_5 import modeling_qwen3_5 + from lxt.efficient.models.qwen3_5 import attnLRP + monkey_patch(modeling_qwen3_5, attnLRP) elif model_type == "mistral": - from lxt.models.mistral import attnlrp - attnlrp.register(self.model) + from transformers.models.mistral import modeling_mistral + from lxt.efficient.models.mistral import attnLRP + monkey_patch(modeling_mistral, attnLRP) + elif model_type == "gemma3": + from transformers.models.gemma3 import modeling_gemma3 + from lxt.efficient.models.gemma3 import attnLRP + monkey_patch(modeling_gemma3, attnLRP) else: raise ValueError( f"AttnLRP not supported for model_type={model_type!r}. " - f"Currently supported: llama, qwen, qwen2, qwen2.5, mistral. " - f"For other architectures, contribute an lxt rules module." + f"Currently supported: llama, qwen, qwen2, qwen2.5, qwen3, qwen3_5, mistral, gemma3. " + f"For other architectures, contribute an lxt rules module under " + f"lxt/efficient/models/." ) - except ImportError: - raise ImportError("lxt is required for AttnLRP. Install with: pip install lxt") from None + except ImportError as e: + raise ImportError( + f"lxt is required for AttnLRP. Install with: " + f"pip install git+https://github.com/rachtibat/LRP-eXplains-Transformers.git " + f"(underlying error: {e})" + ) from None # Free intermediate activations during backward — recompute instead of store self.model.gradient_checkpointing_enable() @@ -179,6 +230,18 @@ def save_relevance_scores(self, output_format: str = "safetensors") -> None: output_path = Path(self.config.output_path) output_path.mkdir(parents=True, exist_ok=True) + # Tied embeddings: lm_head.weight and model.embed_tokens.weight share + # storage on architectures with `tie_word_embeddings=True`. safetensors + # refuses to serialise shared storage. Clone the duplicate so each + # tensor has its own buffer (memory cost is one extra weight tile). + seen = {} + for n, t in list(self.relevance_scores.items()): + ptr = t.data_ptr() + if ptr in seen: + self.relevance_scores[n] = t.clone() + else: + seen[ptr] = n + if output_format == "safetensors": try: from safetensors.torch import save_file diff --git a/mergekit/architecture/auto.py b/mergekit/architecture/auto.py index 5515b44c..347e3dd3 100644 --- a/mergekit/architecture/auto.py +++ b/mergekit/architecture/auto.py @@ -159,10 +159,26 @@ def infer_architecture_info( def _wi(template: str, prefix: str) -> WeightInfo: full_name = prefix + template - optional = (full_name.replace("${layer_index}", "0") not in in_all_models) or ( - tied_keys is not None - and any(re.search(pat, full_name) for pat in tied_keys) - ) + # Layer-aware optional: a template is non-optional only when present + # in EVERY layer index across every model. Hybrid architectures (e.g. + # Qwen3.5 alternating full-attn / linear-attn) have templates that + # only appear in some layers — those must be marked optional or + # the merge planner errors when fetching from layers that lack them. + if "${layer_index}" in template: + num_layers = module_layer_counts.get(prefix, 0) + present_in_every_layer = num_layers > 0 and all( + full_name.replace("${layer_index}", str(i)) in in_all_models + for i in range(num_layers) + ) + optional = (not present_in_every_layer) or ( + tied_keys is not None + and any(re.search(pat, full_name) for pat in tied_keys) + ) + else: + optional = (full_name not in in_all_models) or ( + tied_keys is not None + and any(re.search(pat, full_name) 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 diff --git a/mergekit/architecture/base.py b/mergekit/architecture/base.py index 5fca5bdb..128fa1c5 100644 --- a/mergekit/architecture/base.py +++ b/mergekit/architecture/base.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, Field from transformers import PretrainedConfig +import torch # needed for ConfiguredModuleArchitecture.model_rebuild() from mergekit.common import get_config_value @@ -151,3 +152,12 @@ def get_module(self, module_name: str) -> ConfiguredModuleArchitecture: config=self.config, weight_prefix=self.info.modules[module_name].weight_prefix, ) + + +# Resolve forward references / nested generics now that all dependent +# classes (and torch) are imported. Without this, instantiating +# ConfiguredModuleArchitecture raises PydanticUserError on certain +# multimodal configs (e.g. Qwen3.5 with model.visual.blocks + mtp.layers +# inferred as separate modules). +ConfiguredModuleArchitecture.model_rebuild() +ConfiguredModelArchitecture.model_rebuild() diff --git a/mergekit/config.py b/mergekit/config.py index fc237137..1197c8c9 100644 --- a/mergekit/config.py +++ b/mergekit/config.py @@ -19,7 +19,7 @@ class ConditionalParameter(BaseModel): ParameterSetting: TypeAlias = Union[ - ConditionalParameter, List[ConditionalParameter], ScalarOrGradient + ConditionalParameter, List[ConditionalParameter], ScalarOrGradient, str ] diff --git a/mergekit/merge_methods/lrp.py b/mergekit/merge_methods/lrp.py index 617d7f64..f585004f 100644 --- a/mergekit/merge_methods/lrp.py +++ b/mergekit/merge_methods/lrp.py @@ -126,11 +126,16 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: if importance is not None: importance = importance.to(delta.device) - # Strict LRP: No silent fallback to magnitude + # Strict LRP for tensors that have scores. Tensors without scores + # are treated as base-passthrough: do not contribute a delta. This + # supports use cases where LRP is computed on a subset of the model + # (e.g. the language_model branch of a multimodal LM) and the + # remaining tensors (vision tower, MTP heads) should retain base + # weights rather than being merged with arbitrary magnitudes. if importance is None: - raise RuntimeError( - f"LRP scores for tensor '{self.weight_info.name}' not found or not provided for {ref_str}." - ) + # Skip this source's delta for this tensor — equivalent to + # weighting it 0 against base. The accumulator stays untouched. + continue # Validate importance shape if importance.shape != delta.shape: