使用 🤗 Optimum Intel 和 fastRAG 最佳化 CPU Embedding
Embedding 模型可用於許多應用，例如檢索、重排、聚類和分類。近年來，研究界在 Embedding 模型方面取得了重大進展，從而顯著增強了所有基於語義表示的應用。諸如 BGE、GTE 和 E5 等模型在 MTEB 基準測試中名列前茅，在某些情況下甚至優於專有的 Embedding 服務。Hugging Face 的模型中心提供了各種大小的模型，從輕量級（1-3.5 億引數）到 7B 引數模型（例如 Salesforce/SFR-Embedding-Mistral）。這些基於編碼器架構的輕量級模型是在 CPU 後端上執行基於語義搜尋的應用（例如檢索增強生成，RAG）的理想最佳化和利用物件。
在這篇部落格中，我們將展示如何在基於 Xeon 的 CPU 上實現顯著的效能提升，並展示使用 fastRAG 將最佳化後的模型整合到現有 RAG 管道中是多麼容易。
使用 Embedding 模型進行資訊檢索
Embedding 模型將文字資料編碼為密集向量，捕捉語義和上下文含義。這透過更具上下文的方式表示詞和文件關係，從而實現準確的資訊檢索。通常，語義相似度將透過 Embedding 向量之間的餘弦相似度來衡量。
是否應該始終使用密集向量進行資訊檢索？兩種主要方法各有優缺點。
- 稀疏檢索匹配 n-gram、短語或元資料，以高效且大規模地搜尋大型集合。然而，由於查詢和文件之間的詞彙差距，它可能會錯過相關文件。
- 語義檢索將文字編碼為密集向量，比詞袋模型更好地捕捉上下文和含義。即使存在詞彙不匹配，它也能檢索到語義相關的文件。然而，與 BM25 等詞彙匹配方法相比，它的計算密集、延遲更高，並且需要複雜的編碼模型。
Embedding 模型和 RAG
Embedding 模型在 RAG 應用中具有多種關鍵作用：
- 離線過程：在索引/更新檢索文件庫（索引）時，將文件編碼為密集向量。
- 查詢編碼：在查詢時，它們將輸入查詢編碼為用於檢索的密集向量表示。
- 重排：在初次檢索後，它們可以透過將檢索到的文件編碼為密集向量並與查詢向量進行比較來對它們進行重排。這允許對最初缺少密集表示的文件進行重排。
最佳化 RAG 管道中的 Embedding 模型元件對於獲得更高效的體驗非常重要，特別是：
- 文件索引/更新：更高的吞吐量允許在初始設定或定期更新期間更快速地編碼和索引大型文件集合。
- 查詢編碼：較低的查詢編碼延遲對於響應式即時檢索至關重要。更高的吞吐量支援高效地編碼許多併發查詢，從而實現可擴充套件性。
- 重排檢索到的文件：初次檢索後，Embedding 模型需要快速編碼檢索到的候選項以進行重排。較低的延遲允許對時間敏感型應用快速重排文件。更高的吞吐量支援並行重排更大的候選集，以實現更全面的重排。
使用 Optimum Intel 和 IPEX 最佳化 Embedding 模型
Optimum Intel 是一個開源庫，可加速在英特爾硬體上使用 Hugging Face 庫構建的端到端管道。Optimum Intel 包含多種加速模型的技術，如低位元量化、模型權重剪枝、蒸餾和加速執行時。
Optimum Intel 中包含的執行時和最佳化利用了英特爾 CPU 上的英特爾® 高階向量擴充套件 512（英特爾® AVX-512）、向量神經網路指令（VNNI）和英特爾® 高階矩陣擴充套件（英特爾® AMX）來加速模型。具體而言，它在每個核心中都內建了 BFloat16 (bf16) 和 int8 GEMM 加速器，以加速深度學習訓練和推理工作負載。AMX 加速推理在 PyTorch 2.0 和 Intel Extension for PyTorch (IPEX) 中引入，此外還有針對各種常見運算元的其他最佳化。
使用 Optimum Intel 可以輕鬆最佳化預訓練模型；許多簡單的示例可以在這裡找到。
示例：最佳化 BGE Embedding 模型
在這篇部落格中，我們重點關注北京人工智慧研究院的研究人員最近釋出的 Embedding 模型，因為他們的模型在廣泛採用的 MTEB 排行榜上顯示出有競爭力的結果。
BGE 技術細節
雙編碼器模型是基於 Transformer 的編碼器，經過訓練以最小化兩個語義相似文字向量之間的相似度度量，例如餘弦相似度。例如，流行的 Embedding 模型使用 BERT 模型作為基礎預訓練模型，並對其進行微調以用於 Embedding 文件。表示編碼文字的向量由模型輸出建立；例如，它可以是 [CLS] 標記向量或所有標記向量的平均值。
與更復雜的 Embedding 架構不同，雙編碼器僅編碼單個文件，因此它們缺乏編碼實體（如查詢-文件和文件-文件）之間的上下文互動。然而，由於其簡單的架構，最先進的雙編碼器 Embedding 模型表現出有競爭力的效能並且速度極快。
我們專注於 3 個 BGE 模型：small、base 和 large，分別包含 4500 萬、1.1 億和 3.55 億引數，編碼為 384/768/1024 大小的 Embedding 向量。
我們注意到，我們下面展示的最佳化過程是通用的，可以應用於其他 Embedding 模型（包括雙編碼器、交叉編碼器等）。
分步指南：透過量化進行最佳化
我們提供了一個分步指南，用於提高 Embedding 模型的效能，重點是減少延遲（批次大小為 1）和增加吞吐量（以每秒編碼的文件數衡量）。此方法利用 optimum-intel 和 Intel Neural Compressor 對模型進行量化，並使用 IPEX 在基於英特爾的硬體上進行最佳化執行時。
第 1 步：安裝軟體包
要安裝 optimum-intel 和 intel-extension-for-transformers，請執行以下命令：
pip install -U optimum[neural-compressor] intel-extension-for-transformers
第 2 步：訓練後靜態量化
訓練後靜態量化需要一個校準集來確定權重和啟用的動態範圍。校準是透過在模型上執行一組有代表性的資料樣本，收集統計資訊，然後根據收集到的資訊對模型進行量化，以最小化精度損失。
以下程式碼片段顯示了量化的程式碼：
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,
save_directory=output_path,
batch_size=1,
)
tokenizer.save_pretrained(output_path)
在我們的校準過程中，我們使用了 qasper 資料集的一個子集。
第 3 步：載入和執行推理
載入量化模型只需執行以下命令：
from optimum.intel import IPEXModel
model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")
將句子編碼為向量可以像我們習慣使用 Transformers 庫一樣完成。
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)
# get the [CLS] token
embeddings = outputs[0][:, 0]
我們在下面的評估部分提供了關於如何配置 CPU 後端設定的其他重要細節（正確的機器設定）。
使用 MTEB 進行模型評估
將模型權重以較低精度量化會引入精度損失，因為我們從 `fp32` 權重轉換為 `int8` 時會丟失精度。因此，我們旨在透過將最佳化後的模型與原始模型在兩項 MTEB 任務上進行比較來驗證其準確性：
- 檢索 (Retrieval) - 對語料庫進行編碼，並透過搜尋索引給定查詢來建立排名列表。
- 重排 (Reranking) - 重新排列檢索結果以提高與查詢的相關性。
下表顯示了每種任務型別的平均準確率（在多個數據集上）（重排為 MAP，檢索為 NDCG@10），其中 `int8` 是我們的量化模型，`fp32` 是原始模型（結果來自官方 MTEB 排行榜）。量化模型在重排任務中的錯誤率與原始模型相比低於 1%，在檢索任務中低於 1.55%。
| 重排 | 檢索 | |||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
|
速度和延遲
我們將我們模型的效能與另外兩種常見的模型使用方法進行比較：
- 使用 PyTorch 和 Huggingface 的 Transformers 庫，精度為 `bf16`。
- 使用 Intel extension for PyTorch (IPEX) 執行時，精度為 `bf16`，並使用 torchscript 對模型進行跟蹤。
實驗設定說明：
- 硬體 (CPU)：第四代英特爾至強 8480+，2 個插槽，每個插槽 56 個核心。
- Pytorch 模型在 1 個 CPU 插槽上使用 56 個核心進行評估。
- IPEX/Optimum 設定使用 ipexrun 在 1 個 CPU 插槽上進行評估，核心數從 22 到 56 不等。
- 在所有執行中，都安裝了 TCMalloc 並將其定義為環境變數。
我們如何進行評估？
我們建立了一個指令碼，使用模型的詞彙表生成隨機示例。我們載入了原始模型和最佳化模型，並比較了在我們上面提到的兩種場景下編碼這些示例所需的時間：批次大小為 1 時的延遲，以及使用批處理示例編碼時的吞吐量。
- 基準 PyTorch 和 Hugging Face
import torch
from transformers import AutoModel
model = AutoModel.from_pretrained("BAAI/bge-small-en-v1.5")
@torch.inference_mode()
def encode_text():
outputs = model(inputs)
with torch.cpu.amp.autocast(dtype=torch.bfloat16):
encode_text()
- IPEX torchscript 和 `bf16`
import torch
from transformers import AutoModel
import intel_extension_for_pytorch as ipex
model = AutoModel.from_pretrained("BAAI/bge-small-en-v1.5")
model = ipex.optimize(model, dtype=torch.bfloat16)
vocab_size = model.config.vocab_size
batch_size = 1
seq_length = 512
d = torch.randint(vocab_size, size=[batch_size, seq_length])
model = torch.jit.trace(model, (d,), check_trace=False, strict=False)
model = torch.jit.freeze(model)
@torch.inference_mode()
def encode_text():
outputs = model(inputs)
with torch.cpu.amp.autocast(dtype=torch.bfloat16):
encode_text()
- Optimum Intel 與 IPEX 和 `int8` 模型
import torch
from optimum.intel import IPEXModel
model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")
@torch.inference_mode()
def encode_text():
outputs = model(inputs)
encode_text()
延遲效能
在本次評估中，我們旨在衡量模型的響應速度。這是一個在 RAG 管道中編碼查詢的用例示例。在本次評估中，我們將批處理大小設定為 1，並測量不同文件長度的延遲。
我們可以看到，量化模型在整體上具有最佳的延遲，對於 small 和 base 模型，延遲低於 10 毫秒，對於 large 模型，延遲低於 20 毫秒。與原始模型相比，量化模型的延遲速度提升高達 4.5 倍。
圖 1. BGE 模型的延遲。
吞吐量效能
在我們的吞吐量評估中，我們旨在尋找以每秒文件數為單位的峰值編碼效能。我們將文字長度設定為 256 個標記，因為這是 RAG 管道中平均文件長度的一個很好的估計，並使用不同的批處理大小（4、8、16、32、64、128、256）進行評估。
結果顯示，與其他模型相比，量化模型達到了更高的吞吐量值，並在批處理大小為 128 時達到峰值吞吐量。總體而言，對於所有模型大小，量化模型在各種批處理大小下，與基準 `bf16` 模型相比，效能提升高達 4 倍。
圖 2. BGE small 的吞吐量。
圖 3. BGE base 的吞吐量。
圖 4. BGE large 的吞吐量。
使用 fastRAG 最佳化 Embedding 模型
作為一個例子，我們將演示如何將最佳化後的檢索/重排模型整合到 fastRAG 中（它也可以輕鬆地整合到其他 RAG 框架中，如 Langchain 和 LlamaIndex）。
fastRAG 是由 英特爾實驗室開發的一個研究框架，用於高效和最佳化的檢索增強生成管道，結合了最先進的 LLM 和資訊檢索。fastRAG 與 Haystack 完全相容，幷包括新穎高效的 RAG 模組，以便在英特爾硬體上進行高效部署。要開始使用 fastRAG，我們邀請讀者檢視安裝說明並使用我們的指南開始使用 fastRAG。
我們將最佳化的雙編碼器 Embedding 模型整合在兩個模組中：
QuantizedBiEncoderRetriever– 用於從密集索引中索引和檢索文件。QuantizedBiEncoderRanker– 用於使用 Embedding 模型作為複雜檢索管道的一部分，對文件列表進行重排。
使用最佳化的 Retriever 進行快速索引
讓我們透過使用一個利用最佳化 Embedding 模型的密集檢索器來建立一個密集索引。
首先，建立一個文件儲存庫：
from haystack.document_store import InMemoryDocumentStore
document_store = InMemoryDocumentStore(use_gpu=False, use_bm25=False, embedding_dim=384, return_embedding=True)
然後，向其中新增一些文件：
from haystack.schema import Document
# example documents to index
examples = [
"There is a blue house on Oxford Street.",
"Paris is the capital of France.",
"The first commit in fastRAG was in 2022"
]
documents = []
for i, d in enumerate(examples):
documents.append(Document(content=d, id=i))
document_store.write_documents(documents)
載入一個帶有最佳化雙編碼器 Embedding 模型的 Retriever，並對文件儲存庫中的所有文件進行編碼：
from fastrag.retrievers import QuantizedBiEncoderRetriever
model_id = "Intel/bge-small-en-v1.5-rag-int8-static"
retriever = QuantizedBiEncoderRetriever(document_store=document_store, embedding_model=model_id)
document_store.update_embeddings(retriever=retriever)
使用最佳化的 Ranker 進行重排
下面是一個將最佳化模型載入到 ranker 節點中的示例，該節點會根據查詢對從索引中檢索到的所有文件進行編碼和重排：
from haystack import Pipeline
from fastrag.rankers import QuantizedBiEncoderRanker
ranker = QuantizedBiEncoderRanker("Intel/bge-large-en-v1.5-rag-int8-static")
p = Pipeline()
p.add_node(component=retriever, name="retriever", inputs=["Query"])
p.add_node(component=ranker, name="ranker", inputs=["retriever"])
results = p.run(query="What is the capital of France?")
# print the documents retrieved
print(results)
完成！建立的管道可用於從文件儲存中檢索文件，並使用（另一個）Embedding 模型對檢索到的文件進行重排，以重新排序文件。此 Jupyter Notebook 中提供了一個更完整的示例。
有關更多與 RAG 相關的方法、模型和示例，我們邀請讀者探索 fastRAG/examples 筆記本。
