File size: 5,836 Bytes
68307ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
"""
Usage:
python low_rank_lora.py --repo_id=glif/how2draw --filename="How2Draw-V2_000002800.safetensors" \
--new_rank=4 --new_lora_path="How2Draw-V2_000002800_rank_4.safetensors"
"""
import torch
from huggingface_hub import hf_hub_download
import safetensors.torch
import fire
def sparse_random_projection_matrix(original_rank, new_rank, density=0.1):
"""
Generates a sparse random projection matrix.
Args:
original_rank (int): Original rank (number of rows).
new_rank (int): Reduced rank (number of columns).
density (float): Fraction of non-zero elements.
Returns:
R (torch.Tensor): Sparse random projection matrix.
"""
R = torch.zeros(new_rank, original_rank)
num_nonzero = int(density * original_rank)
for i in range(new_rank):
indices = torch.randperm(original_rank)[:num_nonzero]
values = torch.randn(num_nonzero)
R[i, indices] = values
return R / torch.sqrt(torch.tensor(new_rank, dtype=torch.float32))
def reduce_lora_rank_random_projection(lora_A, lora_B, new_rank=4, use_sparse=False):
"""
Reduces the rank of LoRA matrices lora_A and lora_B using random projections.
Args:
lora_A (torch.Tensor): Original lora_A matrix of shape [original_rank, in_features].
lora_B (torch.Tensor): Original lora_B matrix of shape [out_features, original_rank].
new_rank (int): Desired lower rank.
use_sparse (bool): Use sparse projection matrix.
Returns:
lora_A_new (torch.Tensor): Reduced lora_A matrix of shape [new_rank, in_features].
lora_B_new (torch.Tensor): Reduced lora_B matrix of shape [out_features, new_rank].
"""
original_rank = lora_A.shape[0] # Assuming lora_A.shape = [original_rank, in_features]
# Generate random projection matrix
if use_sparse:
R = sparse_random_projection_matrix(original_rank=original_rank, new_rank=new_rank)
else:
R = torch.randn(new_rank, original_rank, dtype=torch.float32) / torch.sqrt(
torch.tensor(new_rank, dtype=torch.float32)
)
R = R.to(lora_A.device, lora_A.dtype)
# Project lora_A and lora_B
lora_A_new = (R @ lora_A.to(R.dtype)).to(lora_A.dtype) # Shape: [new_rank, in_features]
lora_B_new = (lora_B.to(R.dtype) @ R.T).to(lora_B.dtype) # Shape: [out_features, new_rank]
return lora_A_new, lora_B_new
def reduce_lora_rank_state_dict_random_projection(state_dict, new_rank=4, use_sparse=False):
"""
Reduces the rank of all LoRA matrices in the given state dict using random projections.
Args:
state_dict (dict): The state dict containing LoRA matrices.
new_rank (int): Desired lower rank.
use_sparse (bool): Use sparse projection matrix.
Returns:
new_state_dict (dict): State dict with reduced-rank LoRA matrices.
"""
new_state_dict = state_dict.copy()
keys = list(state_dict.keys())
for key in keys:
if "lora_A.weight" in key:
# Find the corresponding lora_B
lora_A_key = key
lora_B_key = key.replace("lora_A.weight", "lora_B.weight")
if lora_B_key in state_dict:
lora_A = state_dict[lora_A_key]
lora_B = state_dict[lora_B_key]
# Ensure tensors are on CPU for random projection
lora_A = lora_A.to("cuda")
lora_B = lora_B.to("cuda")
# Apply the rank reduction using random projection
lora_A_new, lora_B_new = reduce_lora_rank_random_projection(
lora_A, lora_B, new_rank=new_rank, use_sparse=use_sparse
)
# Update the state dict
new_state_dict[lora_A_key] = lora_A_new
new_state_dict[lora_B_key] = lora_B_new
print(f"Reduced rank of {lora_A_key} and {lora_B_key} to {new_rank}")
return new_state_dict
def compare_approximation_error(orig_state_dict, new_state_dict):
for key in orig_state_dict:
if "lora_A.weight" in key:
lora_A_key = key
lora_B_key = key.replace("lora_A.weight", "lora_B.weight")
lora_A_old = orig_state_dict[lora_A_key]
lora_B_old = orig_state_dict[lora_B_key]
lora_A_new = new_state_dict[lora_A_key]
lora_B_new = new_state_dict[lora_B_key]
# Original delta_W
delta_W_old = (lora_B_old @ lora_A_old).to("cuda")
# Approximated delta_W
delta_W_new = lora_B_new @ lora_A_new
# Compute the approximation error
error = torch.norm(delta_W_old - delta_W_new, p="fro") / torch.norm(delta_W_old, p="fro")
print(f"Relative error for {lora_A_key}: {error.item():.6f}")
def main(
repo_id: str,
filename: str,
new_rank: int,
use_sparse: bool = False,
check_error: bool = False,
new_lora_path: str = None,
):
# ckpt_path = hf_hub_download(repo_id="glif/how2draw", filename="How2Draw-V2_000002800.safetensors")
if new_lora_path is None:
raise ValueError("Please provide a path to serialize the converted state dict.")
ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
original_state_dict = safetensors.torch.load_file(ckpt_path)
new_state_dict = reduce_lora_rank_state_dict_random_projection(
original_state_dict, new_rank=new_rank, use_sparse=use_sparse
)
if check_error:
compare_approximation_error(original_state_dict, new_state_dict)
new_state_dict = {k: v.to("cpu") for k, v in new_state_dict.items()}
# safetensors.torch.save_file(new_state_dict, "How2Draw-V2_000002800_reduced_sparse.safetensors")
safetensors.torch.save(new_state_dict, new_lora_path)
if __name__ == "__main__":
fire.Fire(main)
|