代码

所有代码均可在 Google Colab 免费 GPU 上运行,或本地 RTX 4090(24G VRAM)运行。


Notebook 列表

Notebook说明运行时间
paper02_finetuneQLoRA 微调主脚本(510 条数据,5 epochs)~40 min(A100)
paper02_evalQ(f) / H / Δalign\Delta_\mathrm{align} 四模型评估~15 min
paper02_latencyInterface A 延迟测试(P95 = 1323 ms)~10 min
Open in Colabpaper02_finetune.ipynb — QLoRA 微调Open in Colabpaper02_eval.ipynb — 四模型对比评估

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。

完整数据集见数据集页