引言
随着语言大模型的兴起,蚂蚁内部也出现了很多服务于大模型相关的场景。作为蚂蚁模型推理的重要技术底盘,Ant Ray Serving 是和 Ray 社区的 Ray Serve 合作并在此基础上做了大量扩展的AI服务框架,例如 LDC、高可用、Java/C++ 支持、负载均衡优化和流式通信等处理。接下来会有一系列文章介绍 LLM 与 Ray Serve 相结合的应用场景。
在这篇文章中,我们将介绍 LangChain 和 Ray Serve,以及如何使用 LLM embedding 和向量数据库构建搜索引擎。在之后部分,我们将展示如何加速 embedding 并结合向量数据库和 LLM 创建基于事实的问答服务。此外,我们将优化代码并测量性能:成本、延迟和吞吐量。
这篇文章的涵盖内容如下:
-
LangChain 的介绍,以及为什么它表现非常好;
-
解释 Ray 如何与 LangChain 相辅相成:
-
展示如何在几个小的更改下,部分处理的速度提高4倍或更多;
-
使用 Ray Serve 在 Cloud 中提供 LangChain 的功能;
-
使用自托管模型在同一 Ray 集群中运行 Ray Serve、LangChain 和 Model,且无需关心机器的运维。
使用LangChain和Ray用100行代码构建开源LLM搜索引擎
从搜索开始
Ray 是一款非常强大的 ML 协调框架,与Ray随之而诞生的是大量的文档,文档大小为 120MB。那么我们如何使文档更易于访问?
答案是:使其可搜索!以前创建高质量的自定义搜索结果很困难,但是使用 LangChain,我们可以在大约 100 行代码的情况下完成。这就是 LangChain 的作用。LangChain 为 LLM 周围的所有内容提供了一个惊人的工具套件。它有点像 HuggingFace,但其专为 LLM 设计。它包含有关提示、索引、生成和摘要文本的工具(链)。与 Ray 结合可以使 LangChain 更加强大,Ray 可以简单而快速地帮助您部署 LangChain 服务。相比于依赖远程 API 调用,在同一个 Ray 集群上运行 pipeline 和 LLM 具体降低成本、缩短延迟和控制数据的优点。
创建索引
首先,我们将通过以下步骤来构建索引:
-
下载要索引的内容。
-
读取内容并将其分成小块(每个句子)。这是因为将查询与页面的一部分进行匹配,比与整个页面进行匹配更容易。
-
使用 HuggingFace 的 Sentence Transformer 库生成每个句子的向量表示。
-
将这些向量放回到向量数据库中(这里以 FAISS 为例,但也可以使用喜欢的任何库)。
这段代码的优点在于它的简单性。正如看到的那样,LangChain 已经为我们完成了所有的繁重工作。假设我们下载了 Ray 文档,并读取了所有文档:
loader = ReadTheDocsLoader("docs.ray.io/en/master/")
docs = loader.load()
接下来的步骤是将每个文档分解为小块。LangChain 使用拆分器来完成此操作:
chunks = text_splitter.create_documents(
[doc.page_content for doc in docs],
metadatas=[doc.metadata for doc in docs])
我们希望保留原始 URL 的元数据,因此必须保留这些文档的元数据。现在有了 chunk,我们可以将 chunk 中的数据 embed 成为向量。LLM 提供商提供了远程API来完成这项工作(这是大多数人使用 LangChain 的方式)。这里我们选择从 HuggingFace下载 Sentence Transformer,并在本地运行(只需要很少代码,受了 LangChain对 llama.cpp 支持的启发)。以下是胶水代码。通过这样做,我们可以降低延迟,使用开源技术,并避免需要 HuggingFace 密钥或支付 API 使用费用。
最后,我们有了向量的 embedding,现在可以使用向量数据库(这里使用 FAISS)来存储 embedding。向量数据库经过优化,可在高维空间中快速进行搜索。
from langchain.vectorstores import FAISS
db = FAISS.from_documents(chunks, embeddings)
db.save_local(FAISS_INDEX_PATH)
之后,我们可以构建存储库。
python build_vector_store.py
这需要大约 8 分钟才能执行完毕。其中大部分时间都花在了进行 embedding 上。当然,这种情况并不是什么大问题,但是想象一下如果要索引数百GB而不是数百MB会发生什么。
使用Ray加速索引
【注:这是一个稍微高级一点的话题,初次阅读时可以跳过。这只是展示我们如何将速度提高4倍到8倍】
我们通过并行化embedding来加速索引:
-
将chunk列表切分为 8 个分片。
-
分别对 8 个分片进行 embed 得到向量表示。
-
合并片段。
需要认识到的关键一点是,embedding 是 GPU 加速的,因此如果我们想要这么做,就需要 8 个 GPU。使用 Ray,8 个 GPU 不必在同一台机器上。但即使在单个机器上,Ray 也有显着的优势。而且,Ray 屏蔽了设置集群的复杂度,你所需要做的就是 pip install ray[default],然后 import ray。
这需要对代码进行一些微小的改动。首先,创建一个任务来创建 embedding,然后使用它来索引一个分片。Ray annotation(@ray.remote) 告诉我们每个任务需要一个完整的 GPU:
@ray.remote(num_gpus=1)
def process_shard(shard):
embeddings = LocalHuggingFaceEmbeddings('multi-qa-mpnet-base-dot-v1')
result = FAISS.from_documents(shard, embeddings)
return result
接下来,按照分片切分 chunk 列表。代码如下:
shards = np.array_split(chunks, db_shards)
然后,为每个分片创建一个任务并等待结果。
futures = [process_shard.remote(shards[i]) for i in range(db_shards)]
results = ray.get(futures)
最后,让我们合并片段。我们使用简单的线性合并方式来完成这个操作。
db = results[0]
for i in range(1,db_shards):
db.merge_from(results[i])
你可能会想知道,这真的起作用吗?我们在一个具有 8 个 GPU 的 g4dn.metal 实例上进行了一些测试。原始代码需要 313 秒来创建 embedding,新代码需要 70 秒,提升了 4.5 倍。而且,这里存在创建任务、设置 GPU 等一次性开销。随着数据的增加,这种开销会减少。例如,我们进行了一个简单的测试,使用的是4倍的数据,它约为理论最大性能的 80%(即快 6.5 倍)。
我们可以使用 Ray Dashboard 查看这些 GPU 的状态。很明显,它们都接近 100% 地运行我们刚刚编写的 process_shard 方法。
结果证明,合并向量数据库非常快,仅需 0.3 秒即可合并所有 8 个片段。
服务化
服务化是另一个领域,LangChain 和 Ray Serve 的组合表现出了它的优势。在本系列的下一篇文章中,我们将探索独立的 auto scaling 和 request batching 等功能。
部署成服务所需的步骤有:
-
加载我们创建的 FAISS 数据库,然后将embedding实例化。
-
开始使用 FAISS 进行相似度搜索。
Ray Serve 使这变得非常容易。Ray 使用“deployment”来封装一个简单的 Python 类。__init__ 方法用于加载,而 __call__ 则完成实际的工作。Ray 负责启动服务、部署http等等。这是简化版代码:
@serve.deployment
class VectorSearchDeployment:
def __init__(self):
self.embeddings = …
self.db = FAISS.load_local(FAISS_INDEX_PATH, self.embeddings)
def search(self,query):
results = self.db.max_marginal_relevance_search(query)
retval = <some string processing of the results>
return retval
async def __call__(self, request: Request) -> List[str]:
return self.search(request.query_params["query"])
deployment = VectorSearchDeployment.bind()
用命令行启动这个服务(当然 Serve 还有更多的部署选项):
% serve run serve_vector_store:deployment
现在,我们可以编写一个简单的 Python 脚本来查询服务以获取相关向量(它只是在端口 8000 上运行的 Web 服务)。
import requests
import sys
query = sys.argv[1]
response = requests.post(f'http://localhost:8000/?query={query}')
print(response.content.decode())
最后是一个完整运行流程:
$ python query.py 'Does Ray Serve support batching?'
From http://docs.ray.io/en/master/serve/performance.html
You can check out our microbenchmark instructions
to benchmark Ray Serve on your hardware.
Request Batching#
====
From http://docs.ray.io/en/master/serve/performance.html
You can enable batching by using the ray.serve.batch decorator. Let’s take a look at a simple example by modifying the MyModel class to accept a batch.
from ray import serve
import ray
@serve.deployment
class Model:
def __call__(self, single_sample: int) -> int:
return single_sample * 2
====
From http://docs.ray.io/en/master/ray-air/api/doc/ray.train.lightgbm.LightGBMPredictor.preferred_batch_format.html
native batch format.
DeveloperAPI: This API may change across minor Ray releases.
====
From http://docs.ray.io/en/master/serve/performance.html
Machine Learning (ML) frameworks such as Tensorflow, PyTorch, and Scikit-Learn support evaluating multiple samples at the same time.
Ray Serve allows you to take advantage of this feature via dynamic request batching.
====
结论
我们在上面的代码中展示了如何通过结合 LangChain 和 Ray Serve 的强大之处来构建基于 LLM 的搜索引擎的关键组件,并转化为服务。
请继续关注第二部分,我们将展示如何将其转化为一个类似于 chatgpt 的问答系统。我们将使用开源 LLM(例如 Dolly 2.0)来完成这个任务。
最后,我们将分享第三部分,讲述扩展性和成本。每秒几百个查询的场景还好,但如果需要扩展到更多呢?延迟是否可以接受?
下一步
在以下 Github 仓库中查看本文章中使用的代码和数据。
如果您想了解更多关于 Ray 的信息,请访问 Ray.io和 Docs.Ray.io。
Ray Summit 2023:如果您有兴趣了解 Ray 如何用于构建高性能和可扩展的 LLM 应用,并在 Ray 上进行LLM的调优/训练/服务,请于 9 月 18 日至 20 日加入 Ray 峰会!我们有一系列优秀的主题演讲人,包括来自 OpenAI 的 John Schulman 和来自 Cohere 的 Aidan Gomez,还有关于 Ray 的社区和技术演讲以及 LLM 的实用培训。
如果您正在寻找一种可靠、高效的AI服务框架,那么 Ant Ray Serving 是您应该高优考虑的选择。欢迎使用或者加入Ray Serving团队,与我们一起打造高效、可靠的AI在线服务框架。