Milvus、vLLM、Llama 3.1によるRAGの構築
カリフォルニア大学バークレー校は、2024年7月、LLM推論とサービングのための高速で使いやすいライブラリであるvLLMを、インキュベーション段階のプロジェクトとしてLF AI & Data Foundationに寄贈しました。私たちはvLLMがLF AI & Dataファミリーに加わることを歓迎します！🎉
大規模言語モデル（LLM）とベクトル・データベースは通常、AI幻覚に対処するための一般的なAIアプリケーション・アーキテクチャであるRAG（Retrieval Augmented Generation）を構築するために組み合わされます。このブログでは、Milvus、vLLM、Llama 3.1を使ってRAGを構築し、実行する方法を紹介します。具体的には、Milvusにテキスト情報をベクトル埋め込みとして埋め込み、保存し、このベクトルストアを知識ベースとして使用して、ユーザーの質問に関連するテキストチャンクを効率的に検索する方法を紹介します。最後に、vLLMを活用してMetaのLlama 3.1-8Bモデルを提供し、検索されたテキストによって拡張された回答を生成します。さあ、飛び込もう！
Milvus、vLLM、MetaのLlama 3.1の紹介
Milvusベクトルデータベース
Milvusは、Generative AI(GenAI)ワークロード用のベクトルを保存、インデックス付け、検索するための、オープンソースの目的別分散ベクトルデータベースである。ハイブリッド検索、 メタデータフィルタリング、リランキングを実行し、何兆ものベクトルを効率的に処理する能力により、MilvusはAIや機械学習のワークロードに最適な選択肢となっています。Milvusは、ローカル、クラスタ、またはフルマネージドZilliz Cloudでホストすることができます。
vLLM
vLLMは、UC Berkeley SkyLabで開始されたオープンソースプロジェクトで、LLMサービングパフォーマンスの最適化に焦点を当てています。PagedAttentionによる効率的なメモリ管理、継続的なバッチ処理、最適化されたCUDAカーネルを使用しています。従来の方法と比較して、vLLMはGPUメモリ使用量を半分に削減しながら、配信性能を最大24倍向上させた。
論文「Efficient Memory Management for Large Language Model Serving with PagedAttention」によると、KVキャッシュはGPUメモリの約30%を使用しており、潜在的なメモリ問題につながっている。KVキャッシュは連続したメモリに格納されますが、サイズが変わるとメモリの断片化が起こり、計算効率が悪くなります。
画像1.既存システムにおけるKVキャッシュのメモリ管理（2023 Paged Attention論文）
KVキャッシュに仮想メモリを使用することで、vLLMは必要に応じて物理GPUメモリのみを割り当て、メモリの断片化を排除し、事前割り当てを回避します。テストでは、vLLMはHuggingFace Transformers（HF）およびText Generation Inference（TGI）を上回り、NVIDIA A10GおよびA100 GPU上でHFより最大24倍、TGIより最大3.5倍高いスループットを達成しました。
画像2.vLLMは、HFより8.5倍～15倍、TGIより3.3倍～3.5倍高いスループットを達成している（2023vLLMブログ）。
メタのラマ3.1
Meta's Llama 3.1が2024年7月23日に発表された。405Bモデルはいくつかの公開ベンチマークで最先端の性能を発揮し、128,000の入力トークンのコンテキストウィンドウを持ち、様々な商用利用が許可されている。4,050億パラメータモデルと並行して、Meta社はLlama3 70B（700億パラメータ）と8B（80億パラメータ）の更新版をリリースした。モデルウェイトはMetaのウェブサイトからダウンロードできる。
重要な洞察は、生成されたデータを微調整することでパフォーマンスを向上させることができるが、質の低い例はパフォーマンスを低下させる可能性があるということだった。Llamaチームは、モデル自体、補助モデル、およびその他のツールを使用して、これらの悪い例を特定し、除去するために広範囲に働きました。
MilvusによるRAG検索の構築と実行
データセットを準備する。
私はこのデモのデータセットとしてMilvusの公式ドキュメントを使用した。
from langchain.document_loaders import DirectoryLoader
# Load HTML files already saved in a local directory
path = "../../RAG/rtdocs_new/"
global_pattern = '*.html'
loader = DirectoryLoader(path=path, glob=global_pattern)
docs = loader.load()
# Print num documents and a preview.
print(f"loaded {len(docs)} documents")
print(docs[0].page_content)
pprint.pprint(docs[0].metadata)
loaded 22 documents
Why Milvus Docs Tutorials Tools Blog Community Stars0 Try Managed Milvus FREE Search Home v2.4.x About ...
{'source': 'https://milvus.io/docs/quickstart.md'}
埋め込みモデルをダウンロードする。
次に、HuggingFaceから無料のオープンソースの埋め込みモデルをダウンロードする。
import torch
from sentence_transformers import SentenceTransformer
# Initialize torch settings for device-agnostic code.
N_GPU = torch.cuda.device_count()
DEVICE = torch.device('cuda:N_GPU' if torch.cuda.is_available() else 'cpu')
# Download the model from huggingface model hub.
model_name = "BAAI/bge-large-en-v1.5"
encoder = SentenceTransformer(model_name, device=DEVICE)
# Get the model parameters and save for later.
EMBEDDING_DIM = encoder.get_sentence_embedding_dimension()
MAX_SEQ_LENGTH_IN_TOKENS = encoder.get_max_seq_length()
# Inspect model parameters.
print(f"model_name: {model_name}")
print(f"EMBEDDING_DIM: {EMBEDDING_DIM}")
print(f"MAX_SEQ_LENGTH: {MAX_SEQ_LENGTH}")
model_name: BAAI/bge-large-en-v1.5
EMBEDDING_DIM: 1024
MAX_SEQ_LENGTH: 512
カスタム・データをチャンクして、ベクトルとしてエンコードする。
ここでは、512文字の固定長で、10％の重なりを使用する。
from langchain.text_splitter import RecursiveCharacterTextSplitter
CHUNK_SIZE = 512
chunk_overlap = np.round(CHUNK_SIZE * 0.10, 0)
print(f"chunk_size: {CHUNK_SIZE}, chunk_overlap: {chunk_overlap}")
# Define the splitter.
child_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=chunk_overlap)
# Chunk the docs.
chunks = child_splitter.split_documents(docs)
print(f"{len(docs)} docs split into {len(chunks)} child documents.")
# Encoder input is doc.page_content as strings.
list_of_strings = [doc.page_content for doc in chunks if hasattr(doc, 'page_content')]
# Embedding inference using HuggingFace encoder.
embeddings = torch.tensor(encoder.encode(list_of_strings))
# Normalize the embeddings.
embeddings = np.array(embeddings / np.linalg.norm(embeddings))
# Milvus expects a list of `numpy.ndarray` of `numpy.float32` numbers.
converted_values = list(map(np.float32, embeddings))
# Create dict_list for Milvus insertion.
dict_list = []
for chunk, vector in zip(chunks, converted_values):
# Assemble embedding vector, original text chunk, metadata.
chunk_dict = {
'chunk': chunk.page_content,
'source': chunk.metadata.get('source', ""),
'vector': vector,
}
dict_list.append(chunk_dict)
chunk_size: 512, chunk_overlap: 51.0
22 docs split into 355 child documents.
ベクターをMilvusに保存する。
エンコードされたベクトル埋め込みをMilvusベクトルデータベースに取り込む。
# Connect a client to the Milvus Lite server.
from pymilvus import MilvusClient
mc = MilvusClient("milvus_demo.db")
# Create a collection with flexible schema and AUTOINDEX.
COLLECTION_NAME = "MilvusDocs"
mc.create_collection(COLLECTION_NAME,
EMBEDDING_DIM,
consistency_level="Eventually",
auto_id=True,
overwrite=True)
# Insert data into the Milvus collection.
print("Start inserting entities")
start_time = time.time()
mc.insert(
COLLECTION_NAME,
data=dict_list,
progress_bar=True)
end_time = time.time()
print(f"Milvus insert time for {len(dict_list)} vectors: ", end="")
print(f"{round(end_time - start_time, 2)} seconds")
Start inserting entities
Milvus insert time for 355 vectors: 0.2 seconds
ベクトル検索を行う。
質問をし、Milvusの知識ベースから最近傍のチャンクを検索する。
SAMPLE_QUESTION = "What do the parameters for HNSW mean?"
# Embed the question using the same encoder.
query_embeddings = torch.tensor(encoder.encode(SAMPLE_QUESTION))
# Normalize embeddings to unit length.
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
# Convert the embeddings to list of list of np.float32.
query_embeddings = list(map(np.float32, query_embeddings))
# Define metadata fields you can filter on.
OUTPUT_FIELDS = list(dict_list[0].keys())
OUTPUT_FIELDS.remove('vector')
# Define how many top-k results you want to retrieve.
TOP_K = 2
# Run semantic vector search using your query and the vector database.
results = mc.search(
COLLECTION_NAME,
data=query_embeddings,
output_fields=OUTPUT_FIELDS,
limit=TOP_K,
consistency_level="Eventually")
検索結果は以下のようになります。
Retrieved result #1
distance = 0.7001987099647522
('Chunk text: layer, finds the node closest to the target in this layer, and'
...
'outgoing')
source: https://milvus.io/docs/index.md
Retrieved result #2
distance = 0.6953287124633789
('Chunk text: this value can improve recall rate at the cost of increased'
...
'to the target')
source: https://milvus.io/docs/index.md
vLLMとLlama 3.1-8BによるRAG生成のビルドと実行
vLLMとHuggingFaceのモデルをインストールする。
vLLMはデフォルトでHuggingFaceから大規模な言語モデルをダウンロードします。一般的に、HuggingFaceで新しいモデルを使いたい場合は、いつでもpip install --upgradeまたは-Uを実行してください。また、vLLMでMetaのLlama 3.1モデルの推論を実行するにはGPUが必要です。
vLLMがサポートする全モデルの一覧は、こちらのドキュメントページをご覧ください。
# (Recommended) Create a new conda environment.
conda create -n myenv python=3.11 -y
conda activate myenv
# Install vLLM with CUDA 12.1.
pip install -U vllm transformers torch
import vllm, torch
from vllm import LLM, SamplingParams
# Clear the GPU memory cache.
torch.cuda.empty_cache()
# Check the GPU.
!nvidia-smi
vLLMのインストール方法については、インストールページを参照してください。
HuggingFaceトークンを取得する。
Meta Llama 3.1のようなHuggingFace上のいくつかのモデルは、ユーザーがウェイトをダウンロードできるようにする前に、ライセンスを受け入れることを要求します。したがって、HuggingFaceアカウントを作成し、モデルのライセンスを承認し、トークンを生成する必要があります。
HuggingFaceのこのLlama3.1のページにアクセスすると、条件に同意するよう求めるメッセージが表示されます。Accept License"をクリックし、モデルの重みをダウンロードする前にメタ条件を承認してください。承認には通常1日もかかりません。
承認を受けたら、新しいHuggingFaceトークンを生成しなければなりません。古いトークンは新しい権限では使えません。
vLLMをインストールする前に、新しいトークンでHuggingFaceにログインしてください。以下では、トークンを保存するためにColab secretsを使用しています。
# Login to HuggingFace using your new token.
from huggingface_hub import login
from google.colab import userdata
hf_token = userdata.get('HF_TOKEN')
login(token = hf_token, add_to_git_credential=True)
RAGジェネレーションの実行
デモでは、Llama-3.1-8B モデルを実行する。これは、スピンアップするためにGPUとかなりのメモリを必要とする。以下の例は、A100 GPUを搭載したGoogle Colab Pro（月額10ドル）で実行した。vLLMの実行方法の詳細については、クイックスタート・ドキュメントをご覧ください。
# 1. Choose a model
MODELTORUN = "meta-llama/Meta-Llama-3.1-8B-Instruct"
# 2. Clear the GPU memory cache, you're going to need it all!
torch.cuda.empty_cache()
# 3. Instantiate a vLLM model instance.
llm = LLM(model=MODELTORUN,
enforce_eager=True,
dtype=torch.bfloat16,
gpu_memory_utilization=0.5,
max_model_len=1000,
seed=415,
max_num_batched_tokens=3000)
Milvusから取得したコンテキストとソースを使用してプロンプトを記述します。
# Separate all the context together by space.
contexts_combined = ' '.join(contexts)
# Lance Martin, LangChain, says put the best contexts at the end.
contexts_combined = ' '.join(reversed(contexts))
# Separate all the unique sources together by comma.
source_combined = ' '.join(reversed(list(dict.fromkeys(sources))))
SYSTEM_PROMPT = f"""First, check if the provided Context is relevant to
the user's question. Second, only if the provided Context is strongly relevant, answer the question using the Context. Otherwise, if the Context is not strongly relevant, answer the question without using the Context.
Be clear, concise, relevant. Answer clearly, in fewer than 2 sentences.
Grounding sources: {source_combined}
Context: {contexts_combined}
User's question: {SAMPLE_QUESTION}
"""
prompts = [SYSTEM_PROMPT]
次に、取得したチャンクとプロンプトに詰め込まれた元の質問を使って答えを生成する。
# Sampling parameters
sampling_params = SamplingParams(temperature=0.2, top_p=0.95)
# Invoke the vLLM model.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
# !r calls repr(), which prints a string inside quotes.
print()
print(f"Question: {SAMPLE_QUESTION!r}")
pprint.pprint(f"Generated text: {generated_text!r}")
Question: 'What do the parameters for HNSW MEAN!?'
Generated text: 'Answer: The parameters for HNSW (Hiera(rchical Navigable Small World Graph) are: '
'* M: The maximum degree of nodes on each layer oof the graph, which can improve '
'recall rate at the cost of increased search time. * efConstruction and ef: '
'These parameters specify a search range when building or searching an index.'
上の答えは私には完璧に見えます！
このデモに興味を持たれた方は、ご自由にお試しいただき、感想をお聞かせください。また、DiscordのMilvusコミュニティに参加して、GenAI開発者全員と直接会話することもできます。
参考文献
Ray Summitでの2023 vLLMプレゼンテーション
vLLMサーバの運用に役立つブログvLLMのデプロイ：ステップ・バイ・ステップ・ガイド
