Advanced LoRA Merging Techniques for Multi-Task LLM Serving

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration ---
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
ADAPTER_A_PATH = "./path/to/sql-lora-adapter" # Fictional path to a SQL generation adapter
ADAPTER_B_PATH = "./path/to/python-lora-adapter" # Fictional path to a Python code gen adapter
MERGE_OUTPUT_PATH = "./models/mistral-7b-sql-python-linear-merge"

# --- Load Base Model and Tokenizer ---
def load_base_model(model_name):
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="cpu", # Load on CPU to merge, then move to GPU for inference
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

model, tokenizer = load_base_model(MODEL_NAME)

# --- Load First Adapter and Merge ---
print("Loading and merging the first adapter (SQL)...")
model = PeftModel.from_pretrained(model, ADAPTER_A_PATH)
# This merges the adapter weights into the base model layers
model = model.merge_and_unload()

# --- Load Second Adapter on the *already merged* model ---
# This is a key step. We are now adapting a model that already contains SQL skills.
print("Loading the second adapter (Python) on top of the merged model...")
model = PeftModel.from_pretrained(model, ADAPTER_B_PATH)

# --- Merge the Second Adapter with a Scaling Factor ---
# The merge_and_unload function has a `weight` parameter in newer PEFT versions.
# For older versions, or for more control, we can do it manually.
# Here, we'll use the built-in functionality for simplicity.
print("Merging the second adapter with a 0.5 scaling factor...")
model = model.merge_and_unload(adapter_names=["default"], weight=0.5)

# Now, the model's weights are effectively W_base + 1.0 * ΔW_sql + 0.5 * ΔW_python

print(f"Model merged successfully. Saving to {MERGE_OUTPUT_PATH}")
model.save_pretrained(MERGE_OUTPUT_PATH)
tokenizer.save_pretrained(MERGE_OUTPUT_PATH)

print("Done.")

import torch
import copy
from collections import defaultdict
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration ---
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
ADAPTER_PATHS = {
    "sql": "./path/to/sql-lora-adapter",
    "python": "./path/to/python-lora-adapter",
    "json": "./path/to/json-lora-adapter"
}
MERGE_OUTPUT_PATH = "./models/mistral-7b-ties-merge"
DENSITY = 0.1 # The fraction of weights to keep after trimming (pruning)

# --- Load Base Model --- 
def load_base_model(model_name):
    # Note: We load the base model in its original precision (float16/bfloat16)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="cpu",
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

base_model, tokenizer = load_base_model(MODEL_NAME)

# --- Extract Task Vectors (LoRA Deltas) ---
def get_task_vector(base_model, adapter_path):
    """Loads an adapter and computes the delta weights (task vector)."""
    print(f"Loading adapter from {adapter_path}...")
    model = PeftModel.from_pretrained(copy.deepcopy(base_model), adapter_path)
    model = model.merge_and_unload()
    
    task_vector = {}
    base_sd = base_model.state_dict()
    merged_sd = model.state_dict()

    for key in base_sd:
        if base_sd[key].dtype in [torch.float32, torch.bfloat16, torch.float16]:
            # Calculate the difference and store it
            task_vector[key] = merged_sd[key] - base_sd[key]
            
    del model
    torch.cuda.empty_cache()
    return task_vector

task_vectors = {name: get_task_vector(base_model, path) for name, path in ADAPTER_PATHS.items()}

# --- TIES-Merging Algorithm Implementation ---
def ties_merge(task_vectors, density):
    """Performs TIES-merging on a dictionary of task vectors."""
    
    # 1. Trim: Prune away low-magnitude values
    print(f"\nStep 1: Trimming task vectors to density {density}")
    trimmed_vectors = {}
    for name, vector in task_vectors.items():
        trimmed_vectors[name] = {}
        for key, tensor in vector.items():
            if tensor.dtype not in [torch.float32, torch.bfloat16, torch.float16]:
                continue
            
            # Flatten the tensor to find the threshold
            flat_tensor = tensor.flatten()
            k = int(density * len(flat_tensor))
            if k == 0: continue

            # Find the k-th largest value's magnitude
            threshold = torch.kthvalue(torch.abs(flat_tensor), len(flat_tensor) - k).values
            
            # Create a mask to zero out values below the threshold
            mask = torch.abs(tensor) >= threshold
            trimmed_vectors[name][key] = tensor * mask

    # 2. Elect Sign: Create a disagreement mask and resolve conflicts
    print("Step 2: Electing signs to resolve conflicts")
    sign_vectors = {}
    for name, vector in trimmed_vectors.items():
        sign_vectors[name] = {key: torch.sign(tensor) for key, tensor in vector.items()}

    # Create a unified sign vector
    final_signs = {}
    param_names = list(trimmed_vectors.values())[0].keys()

    for key in param_names:
        # Sum the signs across all models for each parameter
        sign_sum = torch.zeros_like(sign_vectors["sql"][key])
        for name in sign_vectors:
            if key in sign_vectors[name]:
                sign_sum += sign_vectors[name][key]
        
        # Final sign is the sign of the sum
        final_signs[key] = torch.sign(sign_sum)

    # 3. Merge: Average the magnitudes of sign-aligned vectors
    print("Step 3: Merging the sign-aligned, trimmed vectors")
    merged_task_vector = {}
    for key in param_names:
        # Get tensors from all models for the current parameter
        tensors = [trimmed_vectors[name].get(key, torch.zeros_like(final_signs[key])) 
                   for name in trimmed_vectors]
        
        # Align signs and sum magnitudes
        sum_of_magnitudes = torch.zeros_like(final_signs[key])
        non_zero_counts = torch.zeros_like(final_signs[key])

        for tensor in tensors:
            # Align sign before taking abs
            aligned_tensor = tensor * final_signs[key]
            sum_of_magnitudes += torch.abs(aligned_tensor)
            non_zero_counts += (tensor != 0).float()

        # Average the magnitudes (avoid division by zero)
        avg_magnitude = sum_of_magnitudes / torch.clamp(non_zero_counts, min=1)
        
        # Apply the final sign to the averaged magnitude
        merged_task_vector[key] = avg_magnitude * final_signs[key]
        
    return merged_task_vector

# --- Perform Merging and Apply to Base Model ---
merged_delta = ties_merge(task_vectors, density=DENSITY)

final_model_state_dict = base_model.state_dict()
for key, delta_tensor in merged_delta.items():
    final_model_state_dict[key] += delta_tensor

base_model.load_state_dict(final_model_state_dict)

# --- Save Final Model ---
print(f"\nSaving TIES-merged model to {MERGE_OUTPUT_PATH}")
base_model.save_pretrained(MERGE_OUTPUT_PATH)
tokenizer.save_pretrained(MERGE_OUTPUT_PATH)

print("TIES-merging complete.")

import torch
import copy
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration ---
# ... (Use the same MODEL_NAME, ADAPTER_PATHS, MERGE_OUTPUT_PATH as TIES example)
MERGE_OUTPUT_PATH_DARE = "./models/mistral-7b-dare-merge"
DROP_RATE = 0.5 # The fraction of weights to reset to zero

# --- Load Base Model and Task Vectors ---
# (Assume `base_model`, `tokenizer`, and `task_vectors` are already loaded from the TIES example)

# --- DARE Algorithm Implementation ---
def dare_merge(task_vectors, drop_rate):
    """Performs DARE merging on a dictionary of task vectors."""
    if not (0 < drop_rate < 1):
        raise ValueError("drop_rate must be between 0 and 1.")

    scaling_factor = 1.0 / (1.0 - drop_rate)
    final_task_vector = {}
    param_names = list(task_vectors.values())[0].keys()

    # Initialize the final vector with zeros
    for key in param_names:
        final_task_vector[key] = torch.zeros_like(list(task_vectors.values())[0][key])

    print(f"Applying DARE with drop_rate={drop_rate} and scaling_factor={scaling_factor:.2f}")
    for name, vector in task_vectors.items():
        print(f"Processing task: {name}")
        for key, tensor in vector.items():
            if tensor.dtype not in [torch.float32, torch.bfloat16, torch.float16]:
                continue

            # 1. Drop: Create a random mask and apply it
            mask = torch.rand_like(tensor) > drop_rate
            dropped_tensor = tensor * mask
            
            # 2. Rescale: Scale the remaining weights
            rescaled_tensor = dropped_tensor * scaling_factor
            
            # 3. Add to the final task vector (linear average after DARE)
            final_task_vector[key] += rescaled_tensor
    
    # Average the summed vectors
    num_tasks = len(task_vectors)
    for key in final_task_vector:
        final_task_vector[key] /= num_tasks
        
    return final_task_vector

# --- Perform DARE Merging and Apply to a Fresh Base Model ---
dare_model = copy.deepcopy(base_model) # Use a fresh copy of the base model
merged_delta_dare = dare_merge(task_vectors, drop_rate=DROP_RATE)

final_model_state_dict_dare = dare_model.state_dict()
for key, delta_tensor in merged_delta_dare.items():
    final_model_state_dict_dare[key] += delta_tensor

dare_model.load_state_dict(final_model_state_dict_dare)

# --- Save Final Model ---
print(f"\nSaving DARE-merged model to {MERGE_OUTPUT_PATH_DARE}")
dare_model.save_pretrained(MERGE_OUTPUT_PATH_DARE)
tokenizer.save_pretrained(MERGE_OUTPUT_PATH_DARE)

print("DARE-merging complete.")