Diffusers
English
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)