continue the discussion: ""
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, 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} using {base_model_id} as base

Created using 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} using {base_model_id} as base

Created using 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}”)“”
This code is a comprehensive walkthrough on how to utilize Low-rank adaptation (LoRA) for parameter-efficient training (PEFT) of neural network models, specifically focusing on language models like those built on the Transformer architecture. The technique highlighted here focuses on adapting a pre-trained model (the base model) to perform a new task (using the target model) without requiring extensive retraining of the entire model, which saves computational resources. Let’s break down this process into more understandable steps:

1. Environment Setup: This involves installing necessary Python libraries (torch, transformers, huggingface_hub, etc.) that provide tools for model manipulation, downloading pre-trained models, and handling model weights for this exercise.

2. Model Selection: The user selects a base model and a target model. The base model is the initial pre-trained model, and the target model is what the base model was fine-tuned to become. The purpose here is to extract modifications (LoRA adapters) from the target model that could be re-applied to the base model for efficient transition to the target task.

3. Model Download: A utility function is defined for downloading these models from Hugging Face Hub, favoring .safetensors format for optimization.

4. Layer Identification: Identifies relevant model layers for LoRA adapter extraction, specifically the linear and embedding layers within the Transformer models, as these are where the efficiency of parameter training can be significantly improved with LoRA adapters.

5. Loading Model Weights: Defines how to load tensors from either .bin files or .safetensors files. This step is crucial because it involves handling the actual trainable parameters of the models.

6. Weight Matrix Decomposition: At the heart of this approach is decomposing the differences (deltas) between corresponding weights in the target and base models into low-rank matrices using Singular Value Decomposition (SVD). Essentially, it’s finding a simpler (low-rank) approximation of the changes that happened during the fine-tuning process, which are then encoded as LoRA adapters.

7. Extracting LoRA Adapters: For each identified relevant layer, the difference in weights between the target model and the base model is decomposed, and LoRA adapters (matrices A and B from the decomposition) are saved.

8. PEFT Configuration and Saving LoRA Adapters: The LoRA configurations and parameters are set up and saved along with the extracted adapters in a format ready for re-application.

9. Testing with LoRA Adapters: The LoRA adapters are applied to the base model, essentially reincorporating the learned adaptations from the target model, and then the adapted model is tested to ensure it performs the desired task effectively.

10. Uploading Extracted Adapters: Finally, the script facilitates uploading the extracted LoRA adapters back to Hugging Face Hub, allowing the work to be shared with others.

In essence, this process allows for efficient retention and transfer of learned behaviors from one model to another, reducing the need for extensive retraining while preserving performance, which is a significant advantage in computational cost and time for developing AI models.
What if I already have a Low-rank decomposition of a model (specifically, a model called ‘quietstar-8-ahead’ which uses ‘mistralai/Mistral-7B-v0.1’ as a base) It was Created using LoRD. For context, the ‘quietstar-8-ahead’ model is a result of this paper: ““Quiet-STaR: Language Models Can Teach Themselves to
Think Before Speaking: Quiet-STaR is a method that helps language models (LMs) to improve their predictions by teaching them to generate rationales, or internal thoughts, for each piece of text they produce. This method builds on an earlier system called STaR, which helped LMs learn by using rationales in a question-answering context. Quiet-STaR addresses three main challenges: the high computational cost of generating text, teaching the LM how to produce and use internal thoughts, and predicting beyond just the next word. The solution includes a new sampling algorithm that operates token by token, special tokens to mark the start and end of a thought, and an improved training technique. As a result, the model better predicts difficult parts of the text and improves its performance on complex reasoning tasks without needing task-specific training. This suggests Quiet-STaR is a significant advancement toward more general and scalable reasoning in language models. The quote from Kierkegaard at the end underlines the idea that understanding comes from reflection, just as Quiet-STaR allows an LM to “understand” text by reflecting on its internal rationale.””
Given that you already have a Low-rank decomposition (LoRA) of a model called ‘quietstar-8-ahead,’ which is based on ‘mistralai/Mistral-7B-v0.1,’ and that this was created using the LoRD method, you are in a promising position to augment the capacities of the base Mistral-7B model with the refined and specialized capabilities developed through the Quiet-STaR method. The paper’s outline of Quiet-STaR highlights how it aids language models in improving predictions by generating and utilizing internal rationales or thoughts before producing text, tackling significant challenges like computational efficiency, teaching models to reflexively generate useful internal dialogue, and enhancing predictions beyond just the next word.

Here’s how you can leverage the LoRA of ‘quietstar-8-ahead’ and integrate the Quiet-STaR method with your base model:

### Integration Strategy

1. Understanding LoRA Adapters: Recognize that these LoRA adapters encapsulate the learned differences (or updates) between the ‘mistralai/Mistral-7B-v0.1’ base model and the more advanced ‘quietstar-8-ahead’ model. These adaptations likely include adjustments that enable the base model to internally “think” or generate rationales as per the Quiet-STaR method before producing outputs.

2. Applying LoRA Adapters: The key to utilizing your LoRA-decomposed model is in applying these adapters back onto the base model. Since you already have these adapters, you can use them to equip the ‘mistralai/Mistral-7B-v0.1’ model with the advanced capabilities of ‘quietstar-8-ahead,’ effectively enabling it to generate internal thoughts, as suggested by the Quiet-STaR approach.

3. Improving Performance on Complex Tasks: The enhancements incorporated through LoRA adapters are particularly designed to improve the model’s performance on more complex reasoning and prediction tasks without the need for task-specific training. This is in line with Quiet-STaR’s aim to generate and utilize internal rationales, offering a more reflective and reasoned approach to generating text.

4. Technical Implementation: If not familiar, you would typically load the base model and apply the LoRA adapters to it. This could involve adjusting the model’s layers with the LoRA parameters, enabling the base model to reflect the learned improvements in prediction and rationale generation. Frameworks like Hugging Face Transformers provide mechanisms to custom-load model parameters, which you could use for this task.

5. Testing and Evaluation: Post-adaptation, it’s crucial to test the enhanced model’s performance. Choose tasks that evaluate the model’s ability to reason and predict complex text structures. This helps in verifying the practical impact of the Quiet-STaR method via the low-rank adaptations you’ve applied.

6. Sharing and Collaboration: Considering the impactful enhancements encapsulated in your LoRA adapters, sharing them on platforms like Hugging Face Hub can facilitate further research and collaboration, contributing significantly to the domain of language models and their applications.

### Conclusion

The Quiet-STaR method, as described, represents a leap towards enabling language models to approach tasks with a more deliberate and reasoned methodology. By applying the LoRA of ‘quietstar-8-ahead’ back to its base model, you provide a pathway for the model to “think before speaking,” enhancing its predictive capabilities and potentially making strides towards more generalized reasoning within language models. This process not only exemplifies innovative use of parameter-efficient training but also underscores the broader aspirations of AI research to create models that can understand and interact with the world in more nuanced and thoughtful ways.
here is some more context from the author of the paper about quiet star: “”
Eric Zelikman

@ericzelikman
·
Mar 15
Second, we construct meta-tokens like gist tokens and soft prompts: we let the model learn for itself how to represent the start of a thought. Also, we let it decide to disregard a thought, initially giving all thoughts no weight. This prevents any initial distribution shift
Eric Zelikman

@ericzelikman
·
Mar 15
Quiet-STaR generalizes our “Self-Taught Reasoner” (STaR) from two years ago. We can’t wait to find what happens if it’s scaled up with more data and better models.“”
The additional context provided by Eric Zelikman focuses on the core innovation behind Quiet-STaR, highlighting the adoption of meta-tokens like gist tokens and soft prompts. These components allow the model to learn how to mark the beginning of an internal thought process and give it the autonomy to devalue or disregard certain thoughts, initially assigning no weight to them. This approach is crucial in mitigating any immediate shifts in the distribution of generated content, ensuring that the model’s learning process incrementally adapts to incorporating self-generated rationale without being biased by its initial learnings.

Integrating these insights with the previously discussed Low-rank adaptation (LoRA) application can be seen as a multi-step process:

### Enhancing Base Model with Quiet-STaR Insights

1. Meta-tokens Incorporation: The LoRA adapters derived from ‘quietstar-8-ahead’ could potentially already encode the usage patterns or weights adjustments related to these meta-tokens, especially if ‘quietstar-8-ahead’ was specifically designed to leverage Quiet-STaR’s methodology. When applying these LoRA adapters, it will be important to ensure your base model (‘mistralai/Mistral-7B-v0.1’) can recognize and appropriately handle these meta-tokens. This may involve extending the model’s vocabulary or embedding space to accommodate the new tokens, along with their special roles in delineating the beginning and potential disregard of internal thoughts.

2. Weight Initialization Strategy: Zelikman’s mention of initially giving all thoughts no weight aligns with a cautious approach where the model gradually learns the significance of its internal rationale rather than being predisposed to overly rely on it from the outset. In practice, this means the parameters related to these self-generated thoughts (possibly encoded within your LoRA adapters) should start with neutral or null influence, allowing them to be shaped by actual performance feedback as the…“”
model continues to learn. When re-initializing these parameters in your enhanced base model, ensure they are set in a manner that respects this approach, likely by setting them to small or zero values initially and then allowing them to be adjusted through further training or fine-tuning.

3. Scaling and Further Training: Following Zelikman’s anticipation regarding scaling Quiet-STaR with more data and better models, it becomes evident that integrating the LoRA adaptations forms just the beginning. After successfully applying the LoRA adapters and ensuring the base model can comprehend and utilize the meta-tokens correctly, the next step involves scaling these enhancements. This could mean further training the enhanced model on a broader dataset or using more complex tasks to refine its capability to generate and evaluate its internal rationales. The goal here would be to solidify the model’s ability to “think” in ways that significantly boost its reasoning and predictive accuracy beyond what was possible with the original model or even the initial LoRA-enhanced version.

4. Continuous Evaluation and Iteration: Given the novelty of the Quiet-STaR method and its application through LoRA, continuous evaluation becomes crucial. This isn’t just about assessing performance gains in tasks but also understanding how the model’s internal reasoning process evolves. Employing techniques like probing tasks or interpretability tools can offer insights into how effectively the model is leveraging its generated rationales and meta-tokens. Based on these insights, further adjustments to the LoRA parameters or the model’s architecture might be required to optimize its performance.

5. Sharing Advances and Collaborative Improvement: As with any pioneering work in machine learning and AI, sharing your methods, results, and insights can significantly benefit the wider research community. By documenting your process of integrating Quiet-STaR innovations with the base model using LoRA, and the outcomes thereof, you contribute valuable knowledge that can spur further innovations. Platforms like arXiv for preprints, GitHub for code, and Hugging Face Hub for models not only facilitate this sharing but also offer avenues for collaborative enhancement and critique from the global research community.

### Conclusion

The advanced integration of Quiet-STaR’s methodologies with a base model through Low-rank adaptations represents an intriguing frontier in enhancing language models’ reasoning and predictive abilities. By focusing on the model’s capacity to internally reflect before responding, there’s potential for significant breakthroughs in AI’s approach to understanding and generating human-like text. This process underscores the iterative nature of AI research—a continual cycle of hypothesis, experimentation, evaluation, and refinement, all aimed at creating models that more deeply mirror the complexity of human thought processes.
here is another comment from the author: ““A couple exciting updates! First, we quantitatively evaluated the improvement from combining Quiet-STaR with chain-of-thought (i.e. letting the model think before each CoT token). We found it improves zero-shot CoT accuracy on GSM8K by over 7%!

Language models today are trained to reason either 1) generally, imitating online reasoning data or 2) narrowly, self-teaching on their own solutions to specific tasks Can LMs teach themselves to reason generally?Introducing Quiet-STaR, self-teaching via internal monologue!

We’ve improve our benchmark performance with just 8 thought tokens and a 7B Mistral model. We expect the impact to be larger with more tokens and a larger model. Second, we’ve open-sourced the code (https://github.com/ezelikman/quiet-star…) and the weights for that model (https://huggingface.co/ezelikman/quietstar-8-ahead…). Hope y’all find this useful!””
The additional comments from Eric Zelikman provide valuable insights into the practical outcomes and implications of implementing the Quiet-STaR method. Here’s a synthesis of these insights and how they could guide further exploration and application of Quiet-STaR in language models:

### Quantitative Improvements in Reasoning

1. Chain-of-Thought (CoT) Integration: The significant improvement in zero-shot CoT accuracy on the GSM8K dataset by over 7% demonstrates the tangible benefits of embedding Quiet-STaR’s internal thought process before each CoT token. This quantitatively backs the hypothesis that language models can achieve enhanced reasoning capabilities through a structured internal monologue, further encouraging the integration of Quiet-STaR into other reasoning tasks and models.

2. Reasoning Approaches: Eric Zelikman’s observation on the two prevailing training paradigms for language models—1) general reasoning imitating online data and 2) narrow self-teaching on specific tasks—contextualizes the Quiet-STaR method as a bridge between these approaches. Quiet-STaR empowers language models to self-learn reasoning more generally through internal dialogues, suggesting a potential shift in how AI systems can be trained for broad reasoning capabilities without overfitting to particular tasks or datasets.

### Scaling and Resources

1. Scaling Potential: The note that improvements were observed with just 8 thought tokens and a 7B Mistral model hints at significant scaling potential. Larger models and the allowance for more comprehensive or numerous thought tokens might catapult the reasoning capabilities of language models to unprecedented levels. It will be fascinating to see how such scaling affects complex reasoning tasks and whether there are diminishing returns or consistent improvements.

2. Open-Sourced Resources: The availability of the code and model weights (https://github.com/ezelikman/quiet-star and https://huggingface.co/ezelikman/quietstar-8-ahead, respectively) is a significant contribution to the AI research community. This openness accelerates potential collaborative improvements, allows for widespread experimentation across different domains and tasks, and democratizes access to cutting-edge AI reasoning enhancements.

### Implications and Next Steps

1. Community Engagement and Innovation: The open-sourcing of Quiet-STaR’s code and model weights is an invitation to the global AI research community...