RAG/搜索召回之二——Embedding模型bge-m3微调
今天介绍Embedding模型bge-m3的微调。
在做向量召回的时候,我们会使用bge-m3这样的模型把query和doc转化成embedding。但在一些专业领域,尤其是术语比较多的领域,bge-m3的效果可能就不好。这时候比较有效的方法,就是使用少量的高质量样本进行lora微调。
LoRA(Low-Rank Adaptation)是目前大模型微调最主流的方法之一,它的核心思想非常简单:冻结原始模型参数,只训练一小部分新增的低秩矩阵。
假设我们有一个 Embedding 模型,例如bge-m3,ransformer 中一个线性层通常是y = W x y=Wxy=Wx,其中W 是模型参数,x 是输入,y 是输出。例如W的维度是[4096,4096],参数量为4096*4096约1677万参数。整个模型有几百个这样的矩阵,如果全量微调,需要更新所有的参数,需要保存梯度、保存 Adam Optimizer 状态、更新全部参数,显存需求通常会达到模型参数量的数倍。
LoRA不更新W,只在原来的参数上学习一个变化矩阵ΔW,这样W ′ = W + Δ W W'=W+ΔWW′=W+ΔW。而Δ W = B A ΔW=BAΔW=BA,q其中A为的维度[r,d],B的维度为[d,r]。当r比较小的时候,通常为8、16,参数量就能大幅减少。这个r就是所谓的秩。
下面介绍使用LoRA对bge-m3微调过程。
1 准备高质量的微调样本
我这边针对的领域是“新能源电池”。样本不用特别多,几百条到几千条就行。格式如下:
{”query”: ”800V对电池性能有何影响?”, ”pos”: [”在新能源电池的快充(系统)场景中,800V 架构降低线束电流,有利于高功率快充,但绝缘监控与电弧防护要求更高。该点常作为材料选型、电芯设计或系统集成时的判断依据。”], ”neg”: [”在新能源电池的性能指标(低温)场景中,低温下电解液粘度升高、电荷转移阻抗增大,表现为可用容量与功率下降。该点常作为材料选型、电芯设计或系统集成时的判断依据。”, ”在新能源电池的正极材料(结构)场景中,三元正极 NCM811 中镍含量约 80%,可提高能量密度,但对空气和水更敏感,加工需严格控湿。该点常作为材料选型、电芯设计或系统集成时的判断依据。”], ”pos_ids”: [”6497b91f95f0cc53”], ”neg_ids”: [”246887d6a13e5f72”, ”349a5a7d6496d616”], ”seed_id”: ”3aabe3694a2d”, ”topic”: ”快充”, ”aspect”: ”系统”}必须包含的字段:query是查询词,pos为正样本,只要一条就行,neg为难负样本,通常包含多条。注意这里的负样本一定要是难负样本,如果是一些随机的样本效果应该就不太行。
我这里的训练集合是780条,测试集合是100条。
2 环境
我的机器是:MacBook Pro M3 16G。python环境:python3.10,安装所需的库:
torch>=2.1.0 transformers>=4.40.0 sentence-transformers>=3.0.0 peft>=0.11.0 accelerate>=0.30.0 datasets>=2.19.0 pyyaml>=6.0 numpy>=1.24.0 tqdm>=4.66.0 scikit-learn>=1.3.03 微调网络
我们采用PEFT库的get_peft_model函数把普通的Transformer模型改造成LoRA模型。然后在LoRA模型用infoNce进行微调。
首先定义一个数据集:
class TripletJsonlDataset(Dataset): def __init__(self, path: Path, num_negatives: int = 2): self.rows = read_jsonl(path) self.num_negatives = num_negatives self._pos_pool = [r[”pos”][0] for r in self.rows if r.get(”pos”)] def __len__(self) -> int: return len(self.rows) def _sample_random_negs(self, need: int, forbid_text: str) -> List[str]: ”””Sample other rows' positives as random negatives (never the current pos).””” if need <= 0 or not self._pos_pool: return [] out: List[str] = [] tries = 0 while len(out) < need and tries < need * 20: tries += 1 cand = random.choice(self._pos_pool) if cand == forbid_text or cand in out: continue out.append(cand) return out def __getitem__(self, idx: int) -> Dict[str, Any]: row = self.rows[idx] pos = row[”pos”][0] negs = list(row.get(”neg”) or [])[: self.num_negatives] if len(negs) < self.num_negatives: negs = negs + self._sample_random_negs( self.num_negatives - len(negs), forbid_text=pos ) return { ”query”: row[”query”], ”pos”: pos, ”negs”: negs, }这块就是从数据集中读取query、pos和negs字段。如果negs的数量比需要的少,还会随机从别的样本中随机采样进行补充。
从bge-m3模型构建LoRA微调模型:
def build_model(model_name: str, lora_cfg: Dict[str, Any], device: str): tokenizer = AutoTokenizer.from_pretrained(model_name) base = AutoModel.from_pretrained(model_name) peft_config = LoraConfig( task_type=TaskType.FEATURE_EXTRACTION, r=int(lora_cfg[”lora_r”]), lora_alpha=int(lora_cfg[”lora_alpha”]), lora_dropout=float(lora_cfg[”lora_dropout”]), target_modules=list(lora_cfg[”target_modules”]), bias=”none”, ) model = get_peft_model(base, peft_config) model.print_trainable_parameters() model.to(device) return tokenizer, model相关配置文件:
model_name: BAAI/bge-m3 train_file: data/train.jsonl eval_file: data/eval.jsonl output_dir: outputs/lora # LoRA lora_r: 16 lora_alpha: 32 lora_dropout: 0.05 target_modules: - query - value # Training (Mac M3 16G friendly) max_seq_length: 512 train_batch_size: 2 eval_batch_size: 4 gradient_accumulation_steps: 8 num_epochs: 2 learning_rate: 1.0e-4 warmup_ratio: 0.05 weight_decay: 0.01 seed: 42 num_negatives: 2 # Device: auto | mps | cpu device: auto这里只微调query和value,key和dense都没有微调,如果显存足够可以考虑也开启。
Loss定义:
def info_nce_loss( q: torch.Tensor, p: torch.Tensor, negs: torch.Tensor, temperature: float = 0.02, ) -> torch.Tensor: ””” q: [B, D], p: [B, D], negs: [B, N, D] ””” pos_logit = torch.sum(q * p, dim=-1, keepdim=True) / temperature # [B, 1] neg_logit = torch.einsum(”bd,bnd->bn”, q, negs) / temperature # [B, N] logits = torch.cat([pos_logit, neg_logit], dim=1) # [B, 1+N] labels = torch.zeros(q.size(0), dtype=torch.long, device=q.device) return F.cross_entropy(logits, labels)Encoder,把文本调用bge-m3得到embedding:
def encode( model: torch.nn.Module, tokenizer, texts: List[str], device: str, max_length: int, ) -> torch.Tensor: batch = tokenizer( texts, padding=True, truncation=True, max_length=max_length, return_tensors=”pt”, ) batch = {k: v.to(device) for k, v in batch.items()} outputs = model(**batch) # BGE-style: CLS token; keep compatible if pooler absent if hasattr(outputs, ”last_hidden_state”): cls = outputs.last_hidden_state[:, 0] else: cls = outputs[0][:, 0] return F.normalize(cls, p=2, dim=-1)训练代码:
def train(cfg: Dict[str, Any]) -> None: set_seed(int(cfg.get(”seed”, 42))) device = pick_device(str(cfg.get(”device”, ”auto”))) print(f”device={device}”) train_path = ROOT / cfg[”train_file”] eval_path = ROOT / cfg[”eval_file”] out_dir = ROOT / cfg[”output_dir”] out_dir.mkdir(parents=True, exist_ok=True) tokenizer, model = build_model(cfg[”model_name”], cfg, device) train_ds = TripletJsonlDataset(train_path, num_negatives=int(cfg.get(”num_negatives”, 2))) eval_rows = read_jsonl(eval_path) loader = DataLoader( train_ds, batch_size=int(cfg[”train_batch_size”]), shuffle=True, collate_fn=collate_keep, num_workers=0, ) optimizer = torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr=float(cfg[”learning_rate”]), weight_decay=float(cfg.get(”weight_decay”, 0.01)), ) steps_per_epoch = math.ceil(len(loader) / int(cfg[”gradient_accumulation_steps”])) total_steps = steps_per_epoch * int(cfg[”num_epochs”]) warmup_steps = int(total_steps * float(cfg.get(”warmup_ratio”, 0.05))) scheduler = get_linear_schedule_with_warmup(optimizer, warmup_steps, total_steps) max_length = int(cfg[”max_seq_length”]) accum = int(cfg[”gradient_accumulation_steps”]) best_acc = -1.0 global_step = 0 for epoch in range(int(cfg[”num_epochs”])): model.train() running = 0.0 optimizer.zero_grad(set_to_none=True) pbar = tqdm(loader, desc=f”epoch {epoch+1}/{cfg['num_epochs']}”) for step, batch in enumerate(pbar, start=1): q = encode(model, tokenizer, batch[”query”], device, max_length) p = encode(model, tokenizer, batch[”pos”], device, max_length) # negs: list of lists -> encode flat then reshape flat_negs = [n for group in batch[”negs”] for n in group] n_emb = encode(model, tokenizer, flat_negs, device, max_length) bsz = len(batch[”query”]) n_per = len(batch[”negs”][0]) n_emb = n_emb.view(bsz, n_per, -1) loss = info_nce_loss(q, p, n_emb) / accum loss.backward() running += loss.item() * accum if step % accum == 0 or step == len(loader): torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) global_step += 1 pbar.set_postfix(loss=f”{running / step:.4f}”) acc = eval_pair_accuracy(model, tokenizer, eval_rows, device, max_length) print(f”epoch={epoch+1} eval_pair_acc={acc:.4f}”) ckpt = out_dir / f”epoch-{epoch+1}” ckpt.mkdir(parents=True, exist_ok=True) model.save_pretrained(ckpt) tokenizer.save_pretrained(ckpt) meta = {”epoch”: epoch + 1, ”eval_pair_acc”: acc, ”global_step”: global_step} (ckpt / ”train_meta.json”).write_text(json.dumps(meta, indent=2), encoding=”utf-8”) if acc >= best_acc: best_acc = acc best_dir = out_dir / ”best” best_dir.mkdir(parents=True, exist_ok=True) model.save_pretrained(best_dir) tokenizer.save_pretrained(best_dir) (best_dir / ”train_meta.json”).write_text( json.dumps(meta, indent=2), encoding=”utf-8” ) print(f”saved best adapter -> {best_dir}”) print(f”done. best_eval_pair_acc={best_acc:.4f}”)训练日志:
trainable params: 1,572,864 || all params: 569,327,616 || trainable%: 0.2763 epoch 1/2: 100%|█████████████████████████████████████████| 390/390 [02:26<00:00, 2.66it/s, loss=0.4972] epoch=1 eval_pair_acc=0.9600 saved best adapter -> ./outputs/lora/best epoch 2/2: 100%|█████████████████████████████████████████| 390/390 [02:19<00:00, 2.79it/s, loss=0.2856] epoch=2 eval_pair_acc=0.9700 saved best adapter -> ./outputs/lora/best done. best_eval_pair_acc=0.97004 评估微调效果
=== baseline ==={ “MRR@10”: 0.6589007936507936, “nDCG@10”: 0.7273867753390105, “Recall@10”: 0.94}
=== lora ==={ “MRR@10”: 0.7681666666666669, “nDCG@10”: 0.8272662245965128, “Recall@10”: 1.0}
从上面可以看出,微调后效果有明显的提升。
学AI大模型的正确顺序,千万不要搞错了
🤔2026年AI风口已来!各行各业的AI渗透肉眼可见,超多公司要么转型做AI相关产品,要么高薪挖AI技术人才,机遇直接摆在眼前!
有往AI方向发展,或者本身有后端编程基础的朋友,直接冲AI大模型应用开发转岗超合适!
就算暂时不打算转岗,了解大模型、RAG、Prompt、Agent这些热门概念,能上手做简单项目,也绝对是求职加分王🔋
📝给大家整理了超全最新的AI大模型应用开发学习清单和资料,手把手帮你快速入门!👇👇
学习路线:
✅大模型基础认知—大模型核心原理、发展历程、主流模型(GPT、文心一言等)特点解析
✅核心技术模块—RAG检索增强生成、Prompt工程实战、Agent智能体开发逻辑
✅开发基础能力—Python进阶、API接口调用、大模型开发框架(LangChain等)实操
✅应用场景开发—智能问答系统、企业知识库、AIGC内容生成工具、行业定制化大模型应用
✅项目落地流程—需求拆解、技术选型、模型调优、测试上线、运维迭代
✅面试求职冲刺—岗位JD解析、简历AI项目包装、高频面试题汇总、模拟面经
以上6大模块,看似清晰好上手,实则每个部分都有扎实的核心内容需要吃透!
我把大模型的学习全流程已经整理📚好了!抓住AI时代风口,轻松解锁职业新可能,希望大家都能把握机遇,实现薪资/职业跃迁~
