File size: 4,633 Bytes
302920f |
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 158 159 160 161 162 163 164 165 166 |
# Copyright 2024-present the HuggingFace Inc. team.
#
# 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 argparse
import os
import numpy as np
import torch
from datautils import get_calib_data
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model
from peft.tuners.lora.config import CordaConfig, LoraConfig
from peft.tuners.lora.corda import preprocess_corda
@torch.no_grad()
def run_model(model, calib_loader):
model.eval()
for batch in tqdm(calib_loader):
batch = {k: v.to(model.device) for k, v in batch.items()}
model(**batch)
def main(args):
# Setting random seed of numpy and torch
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(args.seed)
elif torch.xpu.is_available():
torch.xpu.manual_seed_all(args.seed)
torch.use_deterministic_algorithms(True)
# Load model
model_id = args.model_id
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_id, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True
)
# Collect data
calib_loader = get_calib_data(args.calib_dataset, tokenizer, model_id, args.calib_loader_size, seed=args.seed)
# Evaluate the original model
print("\n---- model before svd ---\n")
print(model)
# Perform decomposition
corda_config = CordaConfig(
corda_method="ipm" if args.first_eigen else "kpm",
)
lora_config = LoraConfig(
init_lora_weights="corda",
target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
r=args.r,
lora_alpha=args.r,
corda_config=corda_config,
)
preprocess_corda(
model,
lora_config,
run_model=lambda: run_model(model, calib_loader),
)
model = get_peft_model(model, lora_config)
# Evaluate again to check if the model is consistent
# Using `model.model` here because `get_peft_model` wraps a layer to the model
print("\n---- model after svd ---\n")
print(model)
# Save as hugging face model
if args.save_model:
assert args.save_path is not None
save_path = args.save_path
# Save CorDA modules
model.peft_config["default"].init_lora_weights = True
model.save_pretrained(os.path.join(save_path, "corda_init"))
# Save residual model
model = model.unload()
model.save_pretrained(save_path)
# Save tokenizer
tokenizer.save_pretrained(save_path)
print(f"Done building CorDA huggingface model in {save_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_id",
type=str,
default="meta-llama/Llama-2-7b-hf",
help="Pretrained model ID",
)
parser.add_argument(
"--calib_loader_size",
type=int,
default=256,
help="number of samples used for covariance matrices",
)
parser.add_argument(
"--calib_dataset",
type=str,
default="wikitext2",
choices=[
"wikitext2",
"c4",
"ptb",
"traivia_qa",
"nqopen",
"MetaMATH",
"codefeedback",
"WizLMinstruct",
"alpaca",
],
help="calibration dataset",
)
parser.add_argument(
"--eval_mmlu",
action="store_true",
help="evaluate mmlu",
)
parser.add_argument(
"--seed",
type=int,
default=233,
help="random seed",
)
parser.add_argument(
"--r",
type=int,
default=None,
)
parser.add_argument(
"--first_eigen",
action="store_true",
)
parser.add_argument(
"--save_model",
action="store_true",
)
parser.add_argument(
"--save_path",
type=str,
default=None,
)
args = parser.parse_args()
main(args)
|