利用 🤗 Optimum Intel 和 fastRAG 在 CPU 上最佳化文字嵌入

HuggingFace發表於2024-04-02

嵌入模型在很多場合都有廣泛應用,如檢索、重排、聚類以及分類。近年來,研究界在嵌入模型領域取得了很大的進展,這些進展大大提高了基於語義的應用的競爭力。BGEGTE 以及 E5 等模型在 MTEB 基準上長期霸榜,在某些情況下甚至優於私有的嵌入服務。 Hugging Face 模型 hub 提供了多種尺寸的嵌入模型,從輕量級 (100-350M 引數) 到 7B (如 Salesforce/SFR-Embedding-Mistral ) 一應俱全。不少基於語義搜尋的應用會選用基於編碼器架構的輕量級模型作為其嵌入模型,此時,CPU 就成為執行這些輕量級模型的有力候選,一個典型的場景就是 檢索增強生成 (Retrieval Augmented Generation,RAG)

使用嵌入模型進行資訊檢索

嵌入模型把文字資料編碼為稠密向量,這些稠密向量中濃縮了文字的語義及上下文資訊。這種上下文相關的單詞和文件表徵方式使得我們有可能實現更準確的資訊檢索。通常,我們可以用嵌入向量之間的餘弦相似度來度量文字間的語義相似度。

在資訊檢索中是否僅依賴稠密向量就可以了?這需要一定的權衡:

  • 稀疏檢索透過把文字集建模成 n- 元組、短語或後設資料的集合,並透過在集合上進行高效、大規模的搜尋來實現資訊檢索。然而,由於查詢和文件在用詞上可能存在差異,這種方法有可能會漏掉一些相關的文件。
  • 語義檢索將文字編碼為稠密向量,相比於詞袋,其能更好地捕獲上下文及詞義。此時,即使用詞上不能精確匹配,這種方法仍然可以檢索出語義相關的文件。然而,與 BM25 等詞匹配方法相比,語義檢索的計算量更大,延遲更高,並且需要用到複雜的編碼模型。

嵌入模型與 RAG

嵌入模型在 RAG 應用的多個環節中均起到了關鍵的作用:

  • 離線處理: 在生成或更新文件資料庫的索引時,要用嵌入模型將文件編碼為稠密向量。
  • 查詢編碼: 在查詢時,要用嵌入模型將輸入查詢編碼為稠密向量以供後續檢索。
  • 重排: 首輪檢索出初始候選文件列表後,要用嵌入模型將檢索到的文件編碼為稠密向量並與查詢向量進行比較,以完成重排。

可見,為了讓整個應用更高效,最佳化 RAG 流水線中的嵌入模型這一環節非常必要,具體來說:

  • 文件索引/更新: 追求高吞吐,這樣就能更快地對大型文件集進行編碼和索引,從而大大縮短建庫和更新耗時。
  • 查詢編碼: 較低的查詢編碼延遲對於檢索的實時性至關重要。更高的吞吐可以支援更高查詢併發度,從而實現高擴充套件度。
  • 對檢索到的文件進行重排: 首輪檢索後,嵌入模型需要快速對檢索到的候選文件進行編碼以支援重排。較低的編碼延遲意味著重排的速度會更快,從而更能滿足時間敏感型應用的要求。同時,更高的吞吐意味著可以並行對更大的候選集進行重排,從而使得更全面的重排成為可能。

使用 Optimum Intel 和 IPEX 最佳化嵌入模型

Optimum Intel 是一個開源庫,其針對英特爾硬體對使用 Hugging Face 庫構建的端到端流水線進行加速和最佳化。 Optimum Intel 實現了多種模型加速技術,如低位元量化、模型權重修剪、蒸餾以及執行時最佳化。

Optimum Intel 在最佳化時充分利用了英特爾® 先進向量擴充套件 512 (英特爾® AVX-512) 、向量神經網路指令 (Vector Neural Network Instructions,VNNI) 以及英特爾® 高階矩陣擴充套件 (英特爾® AMX) 等特性以加速模型的執行。具體來說,每個 CPU 核中都內建了 BFloat16 ( bf16 ) 和 int8 GEMM 加速器,以加速深度學習訓練和推理工作負載。除了針對各種常見運算的最佳化之外,PyTorch 2.0 和 Intel Extension for PyTorch (IPEX) 中還充分利用了 AMX 以加速推理。

使用 Optimum Intel 可以輕鬆最佳化預訓練模型的推理任務。你可在 此處 找到很多簡單的例子。

示例: 最佳化 BGE 嵌入模型

本文,我們主要關注 北京人工智慧研究院 的研究人員最近釋出的嵌入模型,它們在廣為人知的 MTEB 排行榜上取得了亮眼的排名。

BGE 技術細節

雙編碼器模型基於 Transformer 編碼器架構,其訓練目標是最大化兩個語義相似的文字的嵌入向量之間的相似度,常見的指標是餘弦相似度。舉個常見的例子,我們可以使用 BERT 模型作為基礎預訓練模型,並對其進行微調以生成嵌入模型從而為文件生成嵌入向量。有多種方法可用於根據模型輸出構造出文字的嵌入向量,例如,可以直接取 [CLS] 詞元的嵌入向量,也可以對所有輸入詞元的嵌入向量取平均值。

雙編碼器模型是個相對比較簡單的嵌入編碼架構,其僅針對單個文件上下文進行編碼,因此它們無法對諸如 查詢 - 文件文件 - 文件 這樣的交叉上下文進行編碼。然而,最先進的雙編碼器嵌入模型已能表現出相當有競爭力的效能,再加上因其架構簡單帶來的極快的速度,因此該架構的模型成為了當紅炸子雞。

這裡,我們主要關注 3 個 BGE 模型: smallbase 以及 large,它們的引數量分別為 45M、110M 以及 355M,嵌入向量維度分別為 384、768 以及 1024。

請注意,下文展示的最佳化過程是通用的,你可以將它們應用於任何其他嵌入模型 (包括雙編碼器模型、交叉編碼器模型等)。

模型量化分步指南

下面,我們展示如何提高嵌入模型在 CPU 上的效能,我們的最佳化重點是降低延遲 (batch size 為 1) 以及提高吞吐量 (以每秒編碼的文件數來衡量)。我們用 optimum-intelINC (Intel Neural Compressor) 對模型進行量化,並用 IPEX 來最佳化模型在 Intel 的硬體上的執行時間。

第 1 步: 安裝軟體包

請執行以下命令安裝 optimum-intelintel-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 資料集的一個子集作為校準集。

第 2 步: 載入模型,執行推理

僅需執行以下命令,即可載入量化模型:

from optimum.intel import IPEXModel

model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")

隨後,我們使用 transformers 的 API 將句子編碼為向量:

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 任務上對量化模型與原始模型進行比較以驗證量化模型的準確度到底如何:

  • 檢索 - 對語料庫進行編碼,並生成索引庫,然後在索引庫中搜尋給定查詢,以找出與給定查詢相似的文字並排序。
  • 重排 - 對檢索結果進行重排,以細化與給定查詢的相關性排名。

下表展示了每個任務在多個資料集上的平均準確度 (其中,MAP 用於重排,NDCG@10 用於檢索),表中 int8 表示量化模型, fp32 表示原始模型 (原始模型結果取自官方 MTEB 排行榜)。與原始模型相比,量化模型在重排任務上的準確度損失低於 1%,在檢索任務中的準確度損失低於 1.55%。

重排 檢索
BGE-small
BGE-base
BGE-large
int8 fp32 準確度損失
0.5826 0.5836 -0.17%
0.5886 0.5886 0%
0.5985 0.6003 -0.3%
int8 fp32 準確度損失
0.5138 0.5168 -0.58%
0.5242 0.5325 -1.55%
0.5346 0.5429 -1.53%

速度與延遲

我們用量化模型進行推理,並將其與如下兩種常見的模型推理方法進行效能比較:

  1. 使用 PyTorch 和 Hugging Face 的 transformers 庫以 bf16 精度執行模型。
  2. 使用 IPEXbf16 精度執行模型,並使用 torchscript 對模型進行圖化。

實驗環境配置:

  • 硬體 (CPU): 第四代 Intel 至強 8480+,整機有 2 路 CPU,每路 56 個核。
  • 對 PyTorch 模型進行評估時僅使用單路 CPU 上的 56 個核。
  • IPEX/Optimum 測例使用 ipexrun、單路 CPU、使用的核數在 22-56 之間。
  • 所有測例 TCMalloc,我們安裝並妥善設定了相應的環境變數以保證用到它。

如何執行評估?

我們寫了一個基於模型的詞彙表生成隨機樣本的指令碼。然後分別載入原始模型和量化模型,並比較了它們在上述兩種場景中的編碼時間: 使用單 batch size 度量編碼延遲,使用大 batch size 度量編碼吞吐。

  1. 基線 - 用 PyTorch 及 Hugging Face 執行 bf16 模型:
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()
  1. 用 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()
  1. 用基於 IPEX 後端的 Optimum Intel 執行 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 流水線中對查詢進行編碼的速度。此時,我們將 batch size 設為 1,並測量在各種文件長度下的延遲。

我們可以看到,總的來講,量化模型延遲最小,其中 small 模型和 base 模型的延遲低於 10 毫秒, large 模型的延遲低於 20 毫秒。與原始模型相比,量化模型的延遲提高了 4.5 倍。

圖 1: 各尺寸 BGE 模型的延遲

吞吐效能

在評估吞吐時,我們的目標是尋找峰值編碼效能,其單位為每秒處理文件數。我們將文字長度設定為 256 個詞元,這個長度能較好地代表 RAG 流水線中的平均文件長度,同時我們在不同的 batch size (4、8、16、32、64、128、256) 上進行評估。

結果表明,與其他模型相比,量化模型吞吐更高,且在 batch size 為 128 時達到峰值。總體而言,對於所有尺寸的模型,量化模型的吞吐在各 batch size 上均比基線 bf16 模型高 4 倍左右。

圖 2: BGE small 模型的吞吐

圖 3: BGE base 模型的吞吐

圖 3: BGE large 模型的吞吐

在 fastRAG 中使用量化嵌入模型

我們透過一個例子來演示如何將最佳化後的檢索/重排模型整合進 fastRAG 中 (你也可以很輕鬆地將其整合到其他 RAG 框架中,如 Langchain 及 LlamaIndex) 。

fastRAG 是一個高效且最佳化的檢索增強生成流水線研究框架,其可與最先進的 LLM 和資訊檢索演算法結合使用。fastRAG 與 Haystack 完全相容,並實現了多種新的、高效的 RAG 模組,可高效部署在英特爾硬體上。

大家可以參考 此說明 安裝 fastRAG,並閱讀我們的 指南 以開始 fastRAG 之旅。

我們需要將最佳化的雙編碼器嵌入模型用於下述兩個模組中:

  1. QuantizedBiEncoderRetriever – 用於建立稠密向量索引庫,以及從建好的向量庫中檢索文件
  2. QuantizedBiEncoderRanker – 在對文件列表進行重排的流水線中需要用到嵌入模型。

使用最佳化的檢索器實現快速索引

我們用基於量化嵌入模型的稠密檢索器來建立稠密索引。

首先,建立一個文件庫:

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)

使用最佳化的雙編碼器嵌入模型初始化檢索器,並對文件庫中的所有文件進行編碼:

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)

使用最佳化的排名器進行重排

下面的程式碼片段展示瞭如何將量化模型載入到排序器中,該結點會對檢索器檢索到的所有文件進行編碼和重排:

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)

搞定!我們建立的這個流水線首先從文件庫中檢索文件,並使用 (另一個) 嵌入模型對檢索到的文件進行重排。你也可從這個 Notebook 中獲取更完整的例子。

如欲瞭解更多 RAG 相關的方法、模型和示例,我們邀請大家透過 fastRAG/examples 盡情探索。


英文原文: https://huggingface.co/blog/intel-fast-embedding
原文作者: Peter Izsak,Moshe Berchansky,Daniel Fleischer,Ella Charlaix,Morgan Funtowicz,Moshe Wasserblat
譯者: Matrix Yao (姚偉峰),英特爾深度學習工程師,工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

相關文章