Upload LoRD-for-windows.py

LoRD-for-windows.py

@@ -0,0 +1,391 @@
+#My dirty fix for offline LORA conversion on cpu on Windows
+### Pick locations here:
+
+base_model_id = 'C:/path/to/base-model' #eg C:/models/llama2-hf
+target_model_id = 'C:/path/to/target-model' #eg C:/models/euryale
+cache_dir = './models'
+LORA_OUT_DIR = "./lora"
+
+import os
+from huggingface_hub import list_repo_files, snapshot_download
+
+def init_transformers_model(local_path,cache_dir):
+    from os import listdir
+    from os.path import isfile, join
+    onlyfiles = [f for f in listdir(local_path) if isfile(join(local_path, f))]
+    has_safetensors = any(file.endswith('.safetensors') for file in onlyfiles)
+
+    print(f"Model will be loaded from: {local_path}")
+    if has_safetensors:
+        print("Note: .safetensors found. You better don't have .bin in there.")
+    return os.path.abspath(local_path), has_safetensors
+
+
+# ### Downloading the base model
+
+# In[4]:
+
+base_model_download_path, base_model_has_safetensors = init_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.
+
+# In[5]:
+
+
+# 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
+
+# In[6]:
+
+
+target_model_download_path, target_model_has_safetensors = init_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.
+
+# In[7]:
+
+
+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
+
+
+# In[8]:
+
+
+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
+
+# In[9]:
+
+
+base_model_weights = load_safetensors(models['base']['download_path']) if models['base']['has_safetensors'] else load_pytorch_tensors(models['base']['download_path'])
+print("Base model weights loaded.")
+target_model_weights = load_safetensors(models['target']['download_path']) if models['target']['has_safetensors'] else load_pytorch_tensors(models['target']['download_path'])
+print("Target model weights loaded.")
+
+# ### 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
+
+# In[10]:
+
+
+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.")
+
+    weight = weight.to(torch.float32)
+    # 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 = '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
+
+# In[11]:
+
+
+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
+print("Decomposing LORA...(This may take a few hours for larger models)")
+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')
+
+del target_model_weights
+del base_model_weights
+print("LORA decomposed.")
+# ### Extracting correct module names for PEFT
+# 
+# PEFT config uses partial module names, let's extract them correctly
+
+# In[12]:
+
+
+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.
+
+# In[13]:
+
+
+from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
+from peft import get_peft_model, LoraConfig
+
+
+
+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])),
+)
+
+print("Saving LORA to disk...")
+model = AutoModelForCausalLM.from_pretrained(models['base']['download_path'], load_in_4bit=False)
+
+peft_model = get_peft_model(model, lora_config)
+
+# Save to disk
+peft_model.save_pretrained(LORA_OUT_DIR)
+
+del peft_model
+del model
+print("LORA saved to disk.")
+
+# ### Saving LoRA adapters as SafeTensors
+# 
+# Save the decomposed LoRA weights along our PEFT adapter config
+
+# In[14]:
+
+
+import torch
+from safetensors.torch import save_file
+
+print("Saving LoRA adapters as SafeTensors...")
+for key in loras.keys():
+    loras[key] = loras[key].to('cpu').contiguous()
+
+save_file(loras, os.path.join(LORA_OUT_DIR, 'adapter_model.safetensors'))
+
+print("Saved LoRA adapters as SafeTensors.")
+# 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
+
+# In[22]:
+
+
+import os
+import json
+
+print("Adding metadata...")
+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
+
+# In[23]:
+
+
+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.")