Explain this code like i am a layperson: ""# Extracting LoRA adapters from model diff using SVD

Low-rank adaptation (LoRA) is a popular technique for parameter efficient training (PEFT) of neural networks.

It is usually used together with gradient descent in order to fine-tune models to perform specific task. However it can be also derived from already trained weights using a technique called *Singular value decomposition*.

Inspired by an [already existing technique  for Stable Diffusion](https://github.com/bmaltais/kohya_ss/blob/master/networks/extract_lora_from_models.py), we implement LoRA extraction for `transformers` language models.

First, let's choose a base model and a target model we want to extract a LoRA adapter from. I think it's probably better to use the model that was used as the base model before our target model was fine-tuned (so in this instance, as our target model `valine/OpenPirate` was trained on top of `teknium/OpenHermes-2.5-Mistral-7B` I choose OpenHermes instead of the foundation Mistral 7B model as base model)

base_model_id = 'teknium/OpenHermes-2.5-Mistral-7B'
target_model_id = 'valine/OpenSnark'

setup enviroment

!pip install torch transformers huggingface_hub bitsandbytes accelerate peft sentencepiece protobuf pyyaml safetensors tqdm

## Model download

Define a utility function to download transformer models from the Hugging Face Hub, prefering `.safetensors` files over PyTorch `.bin` files.

import os
from huggingface_hub import list_repo_files, snapshot_download

def download_transformers_model(repo_id, cache_dir=None):
    # Check for .safetensors files in the repository
    repo_files = list_repo_files(repo_id)
    has_safetensors = any(file.endswith('.safetensors') for file in repo_files)

    # Define ignore_patterns based on the presence of .safetensors files
    ignore_patterns = ["*.bin"] if has_safetensors else None

    # Download the repository, ignoring PyTorch .bin files if .safetensors files are present
    local_path = snapshot_download(repo_id=repo_id,
                                    cache_dir=cache_dir,
                                    ignore_patterns=ignore_patterns,
                                    )

    print(f"Model downloaded to: {local_path}")
    if has_safetensors:
        print("Note: PyTorch .bin files were ignored due to the presence of .safetensors files.")
    return os.path.abspath(local_path), has_safetensors


### Downloading the base model

cache_dir = './models'
base_model_download_path, base_model_has_safetensors = download_transformers_model(base_model_id, cache_dir)

models = {
    'base' : {
        'download_path' : base_model_download_path,
        'has_safetensors' : base_model_has_safetensors
    },
    'target' : None
}

### Identifying relevant model layers

Define functions to identify linear and embedding layers within transformer models. These layers are targets for LoRA adapters extraction.

# This code has been modified from its original version on the Axolotl project.
# Copyright 2023 Axolotl contributors.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import torch
import torch
import bitsandbytes as bnb
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
from peft.tuners.lora import QuantLinear


def get_linear_embedding_layers(model_type):
    """
    returns the linear embedding layers needed for loras, dependent on the model arch
    """
    if model_type == "gpt_neox":
        return ["embed_in", "embed_out"]
    if model_type == "falcon":
        return ["word_embeddings", "lm_head"]
    return ["embed_tokens", "lm_head"]


def find_all_linear_names(model):
    cls = (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt, torch.nn.Linear, QuantLinear)

    names = []
    for name, module in model.named_modules():
        if (
            isinstance(module, cls)
            or "Linear" in module.__class__.__name__
            and module.__class__.__name__ not in ("LlamaLinearScalingRotaryEmbedding",)
        ):
            names.append(name)


    return names

def get_linear_module_names(model_id):
    model = AutoModelForCausalLM.from_pretrained(model_id, state_dict={}, device_map="meta") #avoid loading weights as we won't need them
    return find_all_linear_names(model)

linear_module_names = get_linear_module_names(models['base']['download_path'])

### Downloading the target model

target_model_download_path, target_model_has_safetensors = download_transformers_model(target_model_id, cache_dir)

models['target'] = {
    'download_path' : target_model_download_path,
    'has_safetensors' : target_model_has_safetensors
}

### Loading tensors from .bin files

Define functions to load PyTorch tensors from `.bin` files or `.safetensors` file.

import torch
import glob

def load_pytorch_tensors(directory, device='cpu'):
    """
    Loads tensors from .bin files in the specified directory into a dictionary.

    Args:
    - directory (str): Path to the directory containing .bin files.
    - device (str): The device to load the tensors on ('cpu', 'cuda', etc.). Default is 'cpu'.

    Returns:
    - dict: A dictionary containing all tensors from the .bin files.
    """
    tensors_dict = {}
    # Use glob to find all .bin files in the directory
    file_paths = glob.glob(f"{directory}/*.bin")

    # Loop through each file and load its tensors into the dictionary
    for file_path in sorted(file_paths):
        loaded_tensors = torch.load(file_path, map_location=torch.device(device))
        for k, v in loaded_tensors.items():
            tensors_dict[k] = v

    return tensors_dict


import glob
from safetensors import safe_open

def load_safetensors(directory, framework="pt", device='cpu'):
    """
    Loads tensors from .safetensors files in the specified directory into a dictionary.

    Args:
    - directory (str): Path to the directory containing .safetensors files.
    - framework (str): The framework to use ('pt' for PyTorch, 'tf' for TensorFlow, etc.). Default is 'pt'.
    - device (str): The device to load the tensors on ('cpu', 'cuda', etc.). Default is 'cpu'.

    Returns:
    - dict: A dictionary containing all tensors from the .safetensors files.
    """
    tensors_dict = {}
    # Use glob to find all .safetensors files in the directory
    file_paths = glob.glob(f"{directory}/*.safetensors")

    # Loop through each file and load its tensors into the dictionary
    for file_path in sorted(file_paths):
        with safe_open(file_path, framework=framework, device=device) as f:
            for k in f.keys():
                tensors_dict[k] = f.get_tensor(k)

    return tensors_dict

### Loading model weights

Load weights for both base and target models

base_model_weights = load_safetensors(models['base']['download_path']) if models['base']['has_safetensors'] else load_pytorch_tensors(models['base']['download_path'])
target_model_weights = load_safetensors(models['target']['download_path']) if models['target']['has_safetensors'] else load_pytorch_tensors(models['target']['download_path'])

### Weight matrix decomposition

The crux of what we're doing here. We define a function to decompose weight matrices into low-rank matrices using SVD

import torch

def _low_rank_decomposition(weight, reduced_rank=16):
    """
    Decompose a 2D matrix into low-rank matrices A and B using SVD.a

    :param weight: The matrix to decompose, of shape (H, W)
    :param reduced_rank: The final rank of the decomposition
    :return: A tuple of tensors (A, B)
    """
    if weight.dim() != 2:
        raise ValueError(f"Only support 2D matrix, but your input has {weight.dim()} dimensions.")

    # SVD Decomposition
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)

    # Truncated matrices
    A = Vh[:reduced_rank, :]
    B = U[:, :reduced_rank] @ torch.diag(S[:reduced_rank])

    return A, B

def decompose_delta_weight(new_weight, base_weight, alpha, reduced_rank, device=None):
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

    new_weight = new_weight.to(device)
    base_weight = base_weight.to(device)

    """
    Decompose the delta weight into low-rank matrices A and B, considering the alpha scaling factor.

    :param new_weight: The updated weight matrix after applying LoRA.
    :param base_weight: The original weight matrix before LoRA.
    :param alpha: The alpha scaling factor used in LoRA.
    :param reduced_rank: The rank for the low-rank decomposition.
    :return: A tuple of tensors (A, B)
    """
    delta_weight = new_weight - base_weight

    # Check if alpha is applied uniformly
    # Adjust the implementation if alpha is applied differently
    adjusted_delta_weight = delta_weight / alpha

    A, B = _low_rank_decomposition(adjusted_delta_weight, reduced_rank=reduced_rank)

    return A, B

## Extract the LoRAs

from tqdm.notebook import tqdm

loras = {

}

# lower rank captures less of the original model, a rank of 32 is probably reasonable for small delta (task specific finetunes and such)
alpha = 1
rank = 32

for module in tqdm(linear_module_names):
  target_tensor = target_model_weights[module+".weight"]
  base_tensor = base_model_weights[module+".weight"]

  lora_A, lora_B = decompose_delta_weight(target_tensor, base_tensor, alpha, rank)
  loras[f"base_model.model.{module}.lora_A.weight"] = lora_A.to('cpu')
  loras[f"base_model.model.{module}.lora_B.weight"] = lora_B.to('cpu')


### Extracting correct module names for PEFT

PEFT config uses partial module names, let's extract them correctly

def get_module_peft_name(module_name):
    return module_name.split('.')[-1]

### Configuring LoRA model with PEFT

Set up a PEFT LoRA configuration for the model. Load the base model and apply this configuration, saving the configuration on disk. The LoRA weights will be saved later from our SVD decomposition.

from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
from peft import get_peft_model, LoraConfig

LORA_OUT_DIR = "./lora"

lora_config = LoraConfig(
        lora_alpha=32, # Setting the alpha to the to decomposition rank value (instead of alpha value used) seems to give better performance. Further testing would be needed to understand what is the optimal alpha value to use
        lora_dropout=0,
        r=32,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules= list(set([get_module_peft_name(e) for e in linear_module_names])),
)

model = AutoModelForCausalLM.from_pretrained(models['base']['download_path'], load_in_4bit=True)

peft_model = get_peft_model(model, lora_config)

# Save to disk
peft_model.save_pretrained(LORA_OUT_DIR)

del peft_model

### Saving LoRA adapters as SafeTensors

Save the decomposed LoRA weights along our PEFT adapter config

import torch
from safetensors.torch import save_file


for key in loras.keys():
    loras[key] = loras[key].to('cpu').contiguous()

save_file(loras, os.path.join(LORA_OUT_DIR, 'adapter_model.safetensors'))

### Testing the result

Load the LoRA adapters from its saved location. If everything went well, this model base model + extracted adapter will behave like the target module. Let's try and generate responses with it.

import os
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM

config = PeftConfig.from_pretrained(os.path.abspath(LORA_OUT_DIR))
model = AutoModelForCausalLM.from_pretrained(models['base']['download_path'], load_in_4bit=True)
model = PeftModel.from_pretrained(model, os.path.abspath(LORA_OUT_DIR))

tokenizer = AutoTokenizer.from_pretrained(models['base']['download_path'], use_fast=True)

# Test input
messages = [
    {
        "role": "system",
        "content": "You are a human just going about your day.",
    },
    {"role": "user", "content": "Hey what's up?"},
]

# Tokenize and format the chat for the model
tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")

# Generate a response
outputs = model.generate(tokenized_chat, max_new_tokens=128)  # Adjust max_new_tokens if needed
response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(response_text)

Seems snarky to me! It works!

## Uploading the extracted adapter to Hugging Face Hub

First, let's replace the `base_model_name_or_path` value of the adapter config with the base model id instead of the local path

import os
import json

adapter_config_path = os.path.join(LORA_OUT_DIR, 'adapter_config.json')

# Load the configuration from the file
with open(adapter_config_path, 'r') as file:
    config = json.load(file)

# Update the base_model_name_or_path in the configuration
config['base_model_name_or_path'] = base_model_id

# Save the updated configuration back to the file
with open(adapter_config_path, 'w') as file:
    json.dump(config, file, indent=2)

print("Configuration updated successfully.")

Now let's create a readme

import yaml

# Define your metadata as a Python dictionary
metadata = {
    'library_name': 'peft',
    'base_model': base_model_id
}

# Convert the dictionary to YAML format
yaml_frontmatter = yaml.dump(metadata, sort_keys=False)

# Define your Markdown content
markdown_content = f"""
# Low-rank decomposition of [{target_model_id}](https://huggingface.co/{target_model_id}) using [{base_model_id}](https://huggingface.co/{base_model_id}) as base

Created using [LoRD](https://github.com/thomasgauthier/LoRD)
"""

# Combine the YAML frontmatter and Markdown content
full_content = f"---\n{yaml_frontmatter}---\n{markdown_content}"

adapter_readme_path = os.path.join(LORA_OUT_DIR, 'README.md')

# Write to a Markdown file
with open(adapter_readme_path, 'w') as md_file:
    md_file.write(full_content)

print("Markdown file successfully created.")

from huggingface_hub import notebook_login

notebook_login()

from huggingface_hub import HfApi

# Initialize the API
hf_api = HfApi()

# Get the current user's username
user_info = hf_api.whoami()
username = user_info['name']

# Define your model name and the local directory path
model_name = input("Enter your desired model name: ")

repo_id = f"{username}/{model_name}"

# Create the repository on the Hugging Face Hub, setting it to private
try:
    # This will return the URL of the newly created repo if successful
    repo_url = hf_api.create_repo(repo_id=repo_id, exist_ok=True)
    print(f"Repository '{repo_id}' created or already exists on the Hugging Face Hub.")
except Exception as e:
    print(f"Error creating repository: {e}")
    exit(1)  # Exit if there's an issue creating the repo

# Upload all the content from the local folder to your remote repository
try:
    hf_api.upload_folder(folder_path=LORA_OUT_DIR,
                         repo_id=repo_id,
                         repo_type="model")
    print(f"All files in '{LORA_OUT_DIR}' have been uploaded to {repo_url.url} successfully.")
except Exception as e:
    prLoading...
Skip to main content
LoRD.ipynb
LoRD.ipynb_
Extracting LoRA adapters from model diff using SVD
Low-rank adaptation (LoRA) is a popular technique for parameter efficient training (PEFT) of neural networks.

It is usually used together with gradient descent in order to fine-tune models to perform specific task. However it can be also derived from already trained weights using a technique called Singular value decomposition.

Inspired by an already existing technique for Stable Diffusion, we implement LoRA extraction for transformers language models.

First, let's choose a base model and a target model we want to extract a LoRA adapter from. I think it's probably better to use the model that was used as the base model before our target model was fine-tuned (so in this instance, as our target model valine/OpenPirate was trained on top of teknium/OpenHermes-2.5-Mistral-7B I choose OpenHermes instead of the foundation Mistral 7B model as base model)

[ ]
base_model_id = 'teknium/OpenHermes-2.5-Mistral-7B'
target_model_id = 'valine/OpenSnark'
setup enviroment

[ ]
!pip install torch transformers huggingface_hub bitsandbytes accelerate peft sentencepiece protobuf pyyaml safetensors tqdm
Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.1.0+cu118)
Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.37.2)
Requirement already satisfied: huggingface_hub in /usr/local/lib/python3.10/dist-packages (0.20.3)
Requirement already satisfied: bitsandbytes in /usr/local/lib/python3.10/dist-packages (0.42.0)
Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (0.27.2)
Requirement already satisfied: peft in /usr/local/lib/python3.10/dist-packages (0.8.2)
Requirement already satisfied: sentencepiece in /usr/local/lib/python3.10/dist-packages (0.1.99)
Requirement already satisfied: protobuf in /usr/local/lib/python3.10/dist-packages (4.25.2)
Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (6.0.1)
Requirement already satisfied: safetensors in /usr/local/lib/python3.10/dist-packages (0.4.2)
Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (4.66.2)
Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch) (3.9.0)
Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch) (4.4.0)
Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch) (1.12)
Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.0)
Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.2)
Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch) (2024.2.0)
Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch) (2.1.0)
Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.24.1)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.2)
Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2023.12.25)
Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.31.0)
Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.15.2)
Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from bitsandbytes) (1.12.0)
Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate) (5.9.6)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch) (2.1.2)
Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.1.1)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)
Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.13)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2022.12.7)
Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch) (1.3.0)
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv

[notice] A new release of pip is available: 23.3.1 -> 24.0
[notice] To update, run: python -m pip install --upgrade pip
Model download
Define a utility function to download transformer models from the Hugging Face Hub, prefering .safetensors files over PyTorch .bin files.

[ ]
import os
from huggingface_hub import list_repo_files, snapshot_download

def download_transformers_model(repo_id, cache_dir=None):
    # Check for .safetensors files in the repository
    repo_files = list_repo_files(repo_id)
    has_safetensors = any(file.endswith('.safetensors') for file in repo_files)

    # Define ignore_patterns based on the presence of .safetensors files
    ignore_patterns = ["*.bin"] if has_safetensors else None

    # Download the repository, ignoring PyTorch .bin files if .safetensors files are present
    local_path = snapshot_download(repo_id=repo_id,
                                    cache_dir=cache_dir,
                                    ignore_patterns=ignore_patterns,
                                    )

    print(f"Model downloaded to: {local_path}")
    if has_safetensors:
        print("Note: PyTorch .bin files were ignored due to the presence of .safetensors files.")
    return os.path.abspath(local_path), has_safetensors

Downloading the base model
[ ]
cache_dir = './models'
base_model_download_path, base_model_has_safetensors = download_transformers_model(base_model_id, cache_dir)

models = {
    'base' : {
        'download_path' : base_model_download_path,
        'has_safetensors' : base_model_has_safetensors
    },
    'target' : None
}

Identifying relevant model layers
Define functions to identify linear and embedding layers within transformer models. These layers are targets for LoRA adapters extraction.

[ ]
# This code has been modified from its original version on the Axolotl project.
# Copyright 2023 Axolotl contributors.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import torch
import torch
import bitsandbytes as bnb
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
from peft.tuners.lora import QuantLinear


def get_linear_embedding_layers(model_type):
    """
    returns the linear embedding layers needed for loras, dependent on the model arch
    """
    if model_type == "gpt_neox":
        return ["embed_in", "embed_out"]
    if model_type == "falcon":
        return ["word_embeddings", "lm_head"]
    return ["embed_tokens", "lm_head"]


def find_all_linear_names(model):
    cls = (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt, torch.nn.Linear, QuantLinear)

    names = []
    for name, module in model.named_modules():
        if (
            isinstance(module, cls)
            or "Linear" in module.__class__.__name__
            and module.__class__.__name__ not in ("LlamaLinearScalingRotaryEmbedding",)
        ):
            names.append(name)


    return names

def get_linear_module_names(model_id):
    model = AutoModelForCausalLM.from_pretrained(model_id, state_dict={}, device_map="meta") #avoid loading weights as we won't need them
    return find_all_linear_names(model)

linear_module_names = get_linear_module_names(models['base']['download_path'])

Downloading the target model
[ ]
target_model_download_path, target_model_has_safetensors = download_transformers_model(target_model_id, cache_dir)

models['target'] = {
    'download_path' : target_model_download_path,
    'has_safetensors' : target_model_has_safetensors
}

Loading tensors from .bin files
Define functions to load PyTorch tensors from .bin files or .safetensors file.

[ ]
import torch
import glob

def load_pytorch_tensors(directory, device='cpu'):
    """
    Loads tensors from .bin files in the specified directory into a dictionary.

    Args:
    - directory (str): Path to the directory containing .bin files.
    - device (str): The device to load the tensors on ('cpu', 'cuda', etc.). Default is 'cpu'.

    Returns:
    - dict: A dictionary containing all tensors from the .bin files.
    """
    tensors_dict = {}
    # Use glob to find all .bin files in the directory
    file_paths = glob.glob(f"{directory}/*.bin")

    # Loop through each file and load its tensors into the dictionary
    for file_path in sorted(file_paths):
        loaded_tensors = torch.load(file_path, map_location=torch.device(device))
        for k, v in loaded_tensors.items():
            tensors_dict[k] = v

    return tensors_dict

[ ]
import glob
from safetensors import safe_open

def load_safetensors(directory, framework="pt", device='cpu'):
    """
    Loads tensors from .safetensors files in the specified directory into a dictionary.

    Args:
    - directory (str): Path to the directory containing .safetensors files.
    - framework (str): The framework to use ('pt' for PyTorch, 'tf' for TensorFlow, etc.). Default is 'pt'.
    - device (str): The device to load the tensors on ('cpu', 'cuda', etc.). Default is 'cpu'.

    Returns:
    - dict: A dictionary containing all tensors from the .safetensors files.
    """
    tensors_dict = {}
    # Use glob to find all .safetensors files in the directory
    file_paths = glob.glob(f"{directory}/*.safetensors")

    # Loop through each file and load its tensors into the dictionary
    for file_path in sorted(file_paths):
        with safe_open(file_path, framework=framework, device=device) as f:
            for k in f.keys():
                tensors_dict[k] = f.get_tensor(k)

    return tensors_dict
Loading model weights
Load weights for both base and target models

[ ]
base_model_weights = load_safetensors(models['base']['download_path']) if models['base']['has_safetensors'] else load_pytorch_tensors(models['base']['download_path'])
target_model_weights = load_safetensors(models['target']['download_path']) if models['target']['has_safetensors'] else load_pytorch_tensors(models['target']['download_path'])
Weight matrix decomposition
The crux of what we're doing here. We define a function to decompose weight matrices into low-rank matrices using SVD

[ ]
import torch

def _low_rank_decomposition(weight, reduced_rank=16):
    """
    Decompose a 2D matrix into low-rank matrices A and B using SVD.a

    :param weight: The matrix to decompose, of shape (H, W)
    :param reduced_rank: The final rank of the decomposition
    :return: A tuple of tensors (A, B)
    """
    if weight.dim() != 2:
        raise ValueError(f"Only support 2D matrix, but your input has {weight.dim()} dimensions.")

    # SVD Decomposition
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)

    # Truncated matrices
    A = Vh[:reduced_rank, :]
    B = U[:, :reduced_rank] @ torch.diag(S[:reduced_rank])

    return A, B

def decompose_delta_weight(new_weight, base_weight, alpha, reduced_rank, device=None):
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

    new_weight = new_weight.to(device)
    base_weight = base_weight.to(device)

    """
    Decompose the delta weight into low-rank matrices A and B, considering the alpha scaling factor.

    :param new_weight: The updated weight matrix after applying LoRA.
    :param base_weight: The original weight matrix before LoRA.
    :param alpha: The alpha scaling factor used in LoRA.
    :param reduced_rank: The rank for the low-rank decomposition.
    :return: A tuple of tensors (A, B)
    """
    delta_weight = new_weight - base_weight

    # Check if alpha is applied uniformly
    # Adjust the implementation if alpha is applied differently
    adjusted_delta_weight = delta_weight / alpha

    A, B = _low_rank_decomposition(adjusted_delta_weight, reduced_rank=reduced_rank)

    return A, B
Extract the LoRAs
[ ]
from tqdm.notebook import tqdm

loras = {

}

# lower rank captures less of the original model, a rank of 32 is probably reasonable for small delta (task specific finetunes and such)
alpha = 1
rank = 32

for module in tqdm(linear_module_names):
  target_tensor = target_model_weights[module+".weight"]
  base_tensor = base_model_weights[module+".weight"]

  lora_A, lora_B = decompose_delta_weight(target_tensor, base_tensor, alpha, rank)
  loras[f"base_model.model.{module}.lora_A.weight"] = lora_A.to('cpu')
  loras[f"base_model.model.{module}.lora_B.weight"] = lora_B.to('cpu')


Extracting correct module names for PEFT
PEFT config uses partial module names, let's extract them correctly

[ ]
def get_module_peft_name(module_name):
    return module_name.split('.')[-1]
Configuring LoRA model with PEFT
Set up a PEFT LoRA configuration for the model. Load the base model and apply this configuration, saving the configuration on disk. The LoRA weights will be saved later from our SVD decomposition.

[ ]
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
from peft import get_peft_model, LoraConfig

LORA_OUT_DIR = "./lora"

lora_config = LoraConfig(
        lora_alpha=32, # Setting the alpha to the to decomposition rank value (instead of alpha value used) seems to give better performance. Further testing would be needed to understand what is the optimal alpha value to use
        lora_dropout=0,
        r=32,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules= list(set([get_module_peft_name(e) for e in linear_module_names])),
)

model = AutoModelForCausalLM.from_pretrained(models['base']['download_path'], load_in_4bit=True)

peft_model = get_peft_model(model, lora_config)

# Save to disk
peft_model.save_pretrained(LORA_OUT_DIR)

del peft_model

Saving LoRA adapters as SafeTensors
Save the decomposed LoRA weights along our PEFT adapter config

[ ]
import torch
from safetensors.torch import save_file


for key in loras.keys():
    loras[key] = loras[key].to('cpu').contiguous()

save_file(loras, os.path.join(LORA_OUT_DIR, 'adapter_model.safetensors'))
Testing the result
Load the LoRA adapters from its saved location. If everything went well, this model base model + extracted adapter will behave like the target module. Let's try and generate responses with it.

[ ]
import os
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM

config = PeftConfig.from_pretrained(os.path.abspath(LORA_OUT_DIR))
model = AutoModelForCausalLM.from_pretrained(models['base']['download_path'], load_in_4bit=True)
model = PeftModel.from_pretrained(model, os.path.abspath(LORA_OUT_DIR))

tokenizer = AutoTokenizer.from_pretrained(models['base']['download_path'], use_fast=True)

# Test input
messages = [
    {
        "role": "system",
        "content": "You are a human just going about your day.",
    },
    {"role": "user", "content": "Hey what's up?"},
]

# Tokenize and format the chat for the model
tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")

# Generate a response
outputs = model.generate(tokenized_chat, max_new_tokens=128)  # Adjust max_new_tokens if needed
response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(response_text)

Seems snarky to me! It works!

Uploading the extracted adapter to Hugging Face Hub
First, let's replace the base_model_name_or_path value of the adapter config with the base model id instead of the local path

[ ]
import os
import json

adapter_config_path = os.path.join(LORA_OUT_DIR, 'adapter_config.json')

# Load the configuration from the file
with open(adapter_config_path, 'r') as file:
    config = json.load(file)

# Update the base_model_name_or_path in the configuration
config['base_model_name_or_path'] = base_model_id

# Save the updated configuration back to the file
with open(adapter_config_path, 'w') as file:
    json.dump(config, file, indent=2)

print("Configuration updated successfully.")
Configuration updated successfully.
Now let's create a readme

[ ]
import yaml

# Define your metadata as a Python dictionary
metadata = {
    'library_name': 'peft',
    'base_model': base_model_id
}

# Convert the dictionary to YAML format
yaml_frontmatter = yaml.dump(metadata, sort_keys=False)

# Define your Markdown content
markdown_content = f"""
# Low-rank decomposition of [{target_model_id}](https://huggingface.co/{target_model_id}) using [{base_model_id}](https://huggingface.co/{base_model_id}) as base

Created using [LoRD](https://github.com/thomasgauthier/LoRD)
"""

# Combine the YAML frontmatter and Markdown content
full_content = f"---\n{yaml_frontmatter}---\n{markdown_content}"

adapter_readme_path = os.path.join(LORA_OUT_DIR, 'README.md')

# Write to a Markdown file
with open(adapter_readme_path, 'w') as md_file:
    md_file.write(full_content)

print("Markdown file successfully created.")
Markdown file successfully created.
[ ]
from huggingface_hub import notebook_login

notebook_login()

[ ]
from huggingface_hub import HfApi

# Initialize the API
hf_api = HfApi()

# Get the current user's username
user_info = hf_api.whoami()
username = user_info['name']

# Define your model name and the local directory path


Colab paid products - Cancel contracts here
int(f"Error uploading files to the repository: {e}")""