来喽，手把手优化向量模型的吞吐量(2)实战

上篇我们聊了RAG系统里向量模型选型的事儿，这次咱们把焦点切换到另一个关键问题：怎么让向量模型跑得更快、吞吐量更大。先看一个直观的对比——下图是优化后BGE模型的吞吐量提升效果，感受一下差距。 ## 使用 Optimum Intel 和 IPEX 优化嵌入模型 **Optimum Intel** 是一个开源库，专门针对Intel硬件加速Hugging Face生态里的模型流水线。它集成了低比特量化、权重剪枝、知识蒸馏和运行时优化等多种技术，底层充分利用了Intel® A VX-512、VNNI（矢量神经网络指令）以及Intel® AMX（高级矩阵扩展）等硬件特性。每个CPU核里都内置了BFloat16和int8的GEMM翻跟斗，专门用来跑深度学习推理。PyTorch 2.0和Intel Extension for PyTorch（IPEX）更是直接调用了AMX来加速。用Optimum Intel优化预训练模型推理很简单，下面就拿MTEB排行榜上表现亮眼的BGE嵌入模型来演示。 ## BGE技术细节 BGE属于双编码器模型，基于Transformer编码器架构，训练目标就是让两个语义相似的文本的嵌入向量余弦相似度最大化。举个典型例子：拿BERT做基座，微调成嵌入模型，然后给文档生成向量。构造向量有很多方法，比如直接取[CLS] token的嵌入，或者对全部输入token的嵌入取平均。双编码器架构简单，只对单个文档上下文编码，没法处理查询-文档或文档-文档之间的交叉上下文。但现在的先进双编码器已经很有竞争力了，加上速度快，成了标配。 这里我们关注三个BGE模型：small（45M参数，384维）、base（110M参数，768维）、large（355M参数，1024维）。注意，下面讲的优化方法可以应用到任何嵌入模型，包括双编码器和交叉编码器。 ## 模型量化分步指南 目标是提升CPU上的嵌入性能：降低延迟（batch size=1时）和提高吞吐量（每秒编码文档数）。我们使用optimum-intel和Intel Neural Compressor（INC）做量化，再用IPEX优化Intel硬件上的运行。 ### 第1步：安装包 ```bash pip install -U optimum[neural-compressor] intel-extension-for-transformers ``` ### 第2步：训后静态量化 静态量化需要一个校准集来确定权重和激活的动态范围。模型跑一组有代表性的样本，收集统计数据，然后按这些信息量化，尽量降低精度损失。下面是一个量化代码片段： ```python def quantize(model_name: str, output_path: str, calibration_set: "datasets.Dataset"): model = AutoModel.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) def preprocess_function(examples): return tokenizer(examples["text"], padding="max_length", max_length=512, truncation=True) vectorized_ds = calibration_set.map(preprocess_function, num_proc=10) vectorized_ds = vectorized_ds.remove_columns(["text"]) quantizer = INCQuantizer.from_pretrained(model) quantization_config = PostTrainingQuantConfig(approach="static", backend="ipex", domain="nlp") quantizer.quantize( quantization_config=quantization_config, calibration_dataset=vectorized_ds, sa ve_directory=output_path, batch_size=1, ) tokenizer.sa ve_pretrained(output_path) ``` 这里用了qasper数据集的一个子集做校准集（数据集地址：https://huggingface.co/datasets/allenai/qasper）。 ### 第3步：加载模型，运行推理 加载量化模型只需一行： ```python from optimum.intel import IPEXModel model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static") ``` 然后用transformers API把句子编码成向量： ```python from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static") inputs = tokenizer(sentences, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) # 取[CLS] token embeddings = outputs[0][:, 0] ``` ## 使用MTEB进行模型评估 权重从fp32降到int8肯定有精度损失。我们在两个MTEB任务上对比量化模型和原始模型： - **检索**：对语料库编码建索引，然后搜索给定查询，找出相似文本并排序。 - **重排**：对检索结果重新排序，细化相关性排名。 下表是每个任务在多个数据集上的平均准确度（重排用MAP，检索用NDCG@10）。int8是量化模型，fp32是原始模型（来自官方MTEB排行榜）。量化模型在重排上精度损失低于1%，检索低于1.55%。 ## 延迟性能 这里测的是响应速度，直接关系到RAG流水线中查询编码的快慢。batch size设为1，测不同文档长度下的延迟。总体看，量化模型延迟最小，small和base模型低于10毫秒，large模型低于20毫秒。相比原始模型，提升达到4.5倍。 *图 1: 各尺寸BGE模型的延迟* ## 吞吐性能 测峰值编码性能，单位是每秒处理文档数。文本长度设为256 token（接近RAG平均文档长度），batch size分别取4、8、16、32、64、128、256。结果显示，量化模型吞吐在各batch size上都比基线bf16模型高约4倍，且在batch size=128时达到峰值。 *图 2: BGE small模型的吞吐* *图 3: BGE base模型的吞吐* *图 4: BGE large模型的吞吐* 从上面的结果看，量化后延迟和吞吐量都有了质的飞跃。是不是已经蠢蠢欲动想试试了？下篇咱们接着聊一个辅助管理RAG流程的实用工具，敬请期待。