File size: 3,355 Bytes
55c92b3 | 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 | from transformers import AutoModelForCausalLM, AutoTokenizer, QuantoConfig
from alphalora.expert_number import calculate_expert
import json
from tqdm import tqdm
import torch
from lsaq_quant_assign import quantize_llama_like
import os
from datasets import load_dataset
def compute_bi(x_in, x_out):
x_in = x_in.view(-1, x_in.size(-1))
x_out = x_out.view(-1, x_out.size(-1))
cos_sim = torch.nn.functional.cosine_similarity(x_in, x_out, dim=-1)
return (1 - cos_sim.mean()).item()
import argparse
import numpy as np
parser = argparse.ArgumentParser(description="parser")
parser.add_argument("--model_id", type=str, required=True)
parser.add_argument("--cuda_id", type=int, default=1)
args = parser.parse_args()
# model_names = ['Qwen3-8B', ''Llama-2-7b-hf'', 'Mistral-7B-Instruct-v0.3','Llama-3.2-3B-Instruct']
import os
# 强制将 HF_HOME 设置为您有权限的目录
# 注意：这必须在 import datasets 之前执行
os.environ["HF_HOME"] = "/mnt/bn/life-mllm/users/cxr/.cache/huggingface"
os.environ["XDG_CACHE_HOME"] = "/mnt/bn/life-mllm/users/cxr/.cache"
my_cache_dir = "./data_cache"
dataset = load_dataset(
"wikitext",
"wikitext-2-raw-v1",
split="train[:1%]",
cache_dir=my_cache_dir # <--- 强制指定缓存路径
)
# print(dataset[:8])
# model_names = ['Qwen3-8B', 'Qwen3-4B', 'Mistral-7B-Instruct-v0.3','Llama-3.2-3B-Instruct']
model_ids = [args.model_id]
for model_id in model_ids:
# Llama-3.2-3B-Instruct
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map=f"cuda:{args.cuda_id}",
torch_dtype=torch.float16
)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_id)
num_layers = len(model.model.layers)
bi_sums = torch.zeros(num_layers, device=model.device)
count = 0
batch_size = 8
max_len = 128
for i in tqdm(range(0, len(dataset), batch_size), desc="Processing batches"):
texts = dataset[i: i + batch_size]["text"]
# import pdb; pdb.set_trace()
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=max_len).to(model.device)
hidden_states_in, hidden_states_out = {}, {}
def make_hook(layer_id):
def hook(module, input, output):
hidden_states_in[layer_id] = input[0].detach()
hidden_states_out[layer_id] = output.detach()
return hook
hooks = [block.register_forward_hook(make_hook(i)) for i, block in enumerate(model.model.layers)]
with torch.no_grad():
_ = model(**inputs)
for h in hooks: h.remove() # 移除 hook，避免重复注册
# 累积每层 BI
for l in range(num_layers):
bi_sums[l] += compute_bi(hidden_states_in[l], hidden_states_out[l])
count += 1
# break
bi_scores = (bi_sums / count).cpu().tolist()
sorted_bi = sorted(enumerate(bi_scores), key=lambda x: x[1])
print(sorted_bi)
model_name = os.path.basename(model_id.rstrip("/"))
os.makedirs(f"metrics/{model_name}", exist_ok=True)
with open(f"metrics/{model_name}/BI.json", "w") as f:
json.dump(bi_scores, f, indent=4) |