BGE模型微调实战：从原理到生产环境的最佳实践

1次阅读

共计 3156 个字符，预计需要花费 8 分钟才能阅读完成。

BGE（Big Generative Embeddings）模型是近年来在 NLP 领域广泛使用的大型生成式嵌入模型。它通过预训练学习丰富的语言表示，能够捕捉文本中的深层语义信息。然而，预训练模型通常是通用的，直接应用于特定任务时效果可能不佳，因此微调（Fine-tuning）成为必要步骤。

微调的目的是让模型适应特定领域或任务，通过在小规模标注数据上继续训练，调整模型参数以提高性能。BGE 模型的微调尤其重要，因为其庞大的参数量需要精细调整才能发挥最佳效果。

传统微调方法在 BGE 模型上存在几个主要问题：

计算效率低下 ：BGE 模型通常包含数十亿参数，全参数微调需要巨大的计算资源。
内存占用高 ：训练过程中需要存储大量中间变量，容易导致内存溢出（OOM）。
过拟合风险 ：在小规模数据上微调时，模型容易过拟合训练数据，泛化能力下降。

针对上述问题，目前主要有以下几种微调策略：

全参数微调（Full Fine-tuning）：调整所有模型参数。优点是灵活性高，缺点是资源消耗大。
Adapter：在模型层间插入小型神经网络模块，仅训练这些模块。优点是参数量小，缺点是可能引入额外延迟。
LoRA（Low-Rank Adaptation）：通过低秩矩阵分解减少可训练参数。优点是在性能和资源消耗之间取得平衡。

以下是它们的对比：

策略	参数量	计算效率	内存占用	灵活性
全参数微调	高	低	高	高
Adapter	低	高	低	中
LoRA	中	高	中	高

以下是使用 PyTorch Lightning 实现 LoRA 微调的完整示例：

import torch
import pytorch_lightning as pl
from transformers import AutoModel, AutoTokenizer

# 数据预处理
class BGEDataModule(pl.LightningDataModule):
    def __init__(self, tokenizer, batch_size=32):
        super().__init__()
        self.tokenizer = tokenizer
        self.batch_size = batch_size

    def setup(self, stage=None):
        # 加载和预处理数据
        texts = [...]  # 你的文本数据
        labels = [...]  # 对应的标签

        # 分词和编码
        encodings = self.tokenizer(texts, truncation=True, padding=True, return_tensors="pt")
        self.dataset = torch.utils.data.TensorDataset(encodings["input_ids"],
            encodings["attention_mask"],
            torch.tensor(labels)
        )

    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset, batch_size=self.batch_size)

# 定义 LoRA 层
class LoRALayer(torch.nn.Module):
    def __init__(self, in_dim, out_dim, rank=4):
        super().__init__()
        self.A = torch.nn.Parameter(torch.randn(in_dim, rank))
        self.B = torch.nn.Parameter(torch.zeros(rank, out_dim))

    def forward(self, x):
        return x @ (self.A @ self.B)

# 模型定义
class BGEModel(pl.LightningModule):
    def __init__(self, model_name="bert-base-uncased", lr=1e-4):
        super().__init__()
        self.model = AutoModel.from_pretrained(model_name)
        self.classifier = torch.nn.Linear(self.model.config.hidden_size, 2)

        # 添加 LoRA 层
        for name, module in self.model.named_modules():
            if "query" in name or "key" in name:
                module.weight = torch.nn.Parameter(module.weight + LoRALayer(module.in_features, module.out_features)())

        self.lr = lr
        self.loss_fn = torch.nn.CrossEntropyLoss()

    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.last_hidden_state[:, 0, :]
        return self.classifier(pooled_output)

    def training_step(self, batch, batch_idx):
        input_ids, attention_mask, labels = batch
        logits = self(input_ids, attention_mask)
        loss = self.loss_fn(logits, labels)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=self.lr)

# 训练流程
def train():
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    data_module = BGEDataModule(tokenizer)
    model = BGEModel()

    trainer = pl.Trainer(
        max_epochs=3,
        accelerator="gpu",
        devices=1,
        precision=16  # 混合精度训练
    )

    trainer.fit(model, data_module)

if __name__ == "__main__":
    train()