代码
所有代码均可在 Google Colab 免费 GPU 上运行,或本地 RTX 4090(24G VRAM)运行。
Notebook 列表
| Notebook | 说明 | 运行时间 |
|---|---|---|
paper02_finetune | QLoRA 微调主脚本(510 条数据,5 epochs) | ~40 min(A100) |
paper02_eval | Q(f) / H / 四模型评估 | ~15 min |
paper02_latency | Interface A 延迟测试(P95 = 1323 ms) | ~10 min |
QLoRA 微调脚本关键代码
模型加载与量化配置
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
# 4-bit NF4 量化
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-3B-Instruct",
quantization_config=bnb_config,
device_map="auto",
)
# LoRA 配置
lora_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
lora_dropout=0.1,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 10,649,600 || all params: 3,096,559,616 || trainable%: 0.3440
训练参数
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir="./brain_llm_lora",
num_train_epochs=5,
per_device_train_batch_size=2,
gradient_accumulation_steps=4, # 等效 batch_size=8
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.05,
bf16=True,
logging_steps=20,
save_strategy="epoch",
report_to="none",
)
d_Sem 评估函数
from FlagEmbedding import FlagModel
import numpy as np
def compute_d_sem(model, tokenizer, vocab_v167l, test_instructions):
"""
计算子流形语义距离 d_Sem。
d_Sem = 1 - (1/|V|) * sum_v max_u cos(phi(u), psi(v))
其中 phi(u) 为 LLM 指令嵌入,psi(v) 为 BGE-large 参考嵌入。
"""
bge = FlagModel("BAAI/bge-large-zh-v1.5", use_fp16=True)
ref_embeddings = bge.encode(vocab_v167l) # psi(v)
# 获取 LLM 指令嵌入 phi(u)
llm_embeddings = get_llm_embeddings(model, tokenizer, test_instructions)
scores = []
for v_emb in ref_embeddings:
max_cos = max(cosine_sim(u_emb, v_emb) for u_emb in llm_embeddings)
scores.append(max_cos)
d_sem = 1.0 - np.mean(scores)
return d_sem
# 微调后 d_Sem 显著降低,验证谱泛化界 Q(f_gen)-Q(f_dom) ≤ L·d_Sem
路由评估(Q(f) 和 H 计算)
def evaluate_routing(model, tokenizer, test_set, vocab_v167l, danger_pairs):
"""
评估路由质量:全局命中率 Q(f) 和危险对幻觉率 H。
"""
correct = 0
total = 0
hazard_events = 0
hazard_total = 0
for item in test_set:
instruction = item["instruction"]
true_route = item["route"]
# 生成路由预测
predicted = generate_route(model, tokenizer, instruction, vocab_v167l)
if predicted == true_route:
correct += 1
total += 1
# 危险对幻觉检测
if true_route in [p[1] for p in danger_pairs if p[0] == predicted]:
hazard_events += 1
hazard_total += 1
Q = correct / total # 目标: ≥ 0.90
H = hazard_events / hazard_total # 目标: < 0.05
delta_align = 1 - Q
return {"Q": Q, "H": H, "delta_align": delta_align}
环境配置
# Python 环境(建议 Python 3.10+)
pip install transformers==4.44.0
pip install peft==0.12.0
pip install bitsandbytes==0.43.3
pip install accelerate==0.33.0
pip install FlagEmbedding==1.2.11
pip install trl==0.10.1
# VRAM 要求
# 训练:6.8 GB(RTX 4090 24G)
# 推理:~3.2 GB(bfloat16)
说明
训练数据构造(510 条)的代码在 paper02_finetune.ipynb 的 Data Generation 部分,包含:
- 模板扩充:
generate_paraphrases(v167l_entry, n=8),基于 Qwen2.5-72B API 生成 8 种改写; - 危险对强化:
generate_danger_pair_samples(danger_pairs)构造混淆-正确样本对; - 欠定/矛盾样本:
generate_ambiguous_samples(n=15)标注为 reject。
完整数据集见数据集页。