Aufbau von RAG mit Milvus, vLLM und Llama 3.1
Die University of California - Berkeley hat der LF AI & Data Foundation im Juli 2024 vLLM, eine schnelle und einfach zu bedienende Bibliothek für LLM Inferenz und Serving, als Projekt in der Inkubationsphase gespendet. Als Mitgliedsprojekt heißen wir vLLM in der LF AI & Data Familie herzlich willkommen! 🎉
Große Sprachmodelle(Large Language Models, LLMs) und Vektordatenbanken werden in der Regel kombiniert, um Retrieval Augmented Generation(RAG) zu erstellen, eine beliebte KI-Anwendungsarchitektur zur Bewältigung von KI-Halluzinationen. Dieser Blog wird Ihnen zeigen, wie Sie eine RAG mit Milvus, vLLM und Llama 3.1 erstellen und ausführen. Genauer gesagt zeige ich Ihnen, wie Sie Textinformationen als Vektoreinbettungen in Milvus einbetten und speichern und diesen Vektorspeicher als Wissensbasis nutzen, um effizient Textabschnitte abzurufen, die für Benutzerfragen relevant sind. Schließlich werden wir vLLM nutzen, um Metas Llama 3.1-8B Modell zu verwenden, um Antworten zu generieren, die durch den abgerufenen Text ergänzt werden. Tauchen wir ein!
Einführung in Milvus, vLLM und Metas Llama 3.1
Milvus Vektor-Datenbank
Milvus ist eine verteilte Open-Source-Vektordatenbank, die speziell für die Speicherung, Indizierung und Suche von Vektoren für generative KI (GenAI) entwickelt wurde. Seine Fähigkeit, eine hybride Suche, Metadatenfilterung und ein Reranking durchzuführen und Billionen von Vektoren effizient zu verarbeiten, macht Milvus zur ersten Wahl für KI- und Machine-Learning-Workloads. Milvus kann lokal, in einem Cluster oder in der vollständig verwalteten Zilliz Cloud betrieben werden.
vLLM
vLLM ist ein Open-Source-Projekt, das am UC Berkeley SkyLab gestartet wurde und sich auf die Optimierung der LLM-Serving-Leistung konzentriert. Es verwendet eine effiziente Speicherverwaltung mit PagedAttention, kontinuierliche Stapelverarbeitung und optimierte CUDA-Kernel. Im Vergleich zu herkömmlichen Methoden verbessert vLLM die Serving-Leistung um das bis zu 24-fache und halbiert gleichzeitig den Speicherbedarf der GPU.
Laut dem Papier "Efficient Memory Management for Large Language Model Serving with PagedAttention" belegt der KV-Cache etwa 30 % des GPU-Speichers, was zu potenziellen Speicherproblemen führt. Der KV-Cache wird in einem zusammenhängenden Speicher gespeichert, aber eine Änderung der Größe kann zu einer Fragmentierung des Speichers führen, was für Berechnungen ineffizient ist.
Abbildung 1. KV-Cache-Speicherverwaltung in bestehenden Systemen (2023 Paged Attention paper)
Durch die Verwendung von virtuellem Speicher für den KV-Cache weist vLLM den physischen GPU-Speicher nur bei Bedarf zu, wodurch eine Speicherfragmentierung vermieden und eine Vorabzuweisung vermieden wird. In Tests übertraf vLLM HuggingFace Transformers (HF) und Text Generation Inference (TGI) und erreichte einen bis zu 24-mal höheren Durchsatz als HF und einen 3,5-mal höheren als TGI auf NVIDIA A10G und A100 GPUs.
Abbildung 2. Serving-Durchsatz, wenn für jede Anfrage drei parallele Output-Vervollständigungen angefordert werden. vLLM erreicht einen 8,5- bis 15-mal höheren Durchsatz als HF und einen 3,3- bis 3,5-mal höheren Durchsatz als TGI (2023 vLLM Blog).
Metas Llama 3.1
Metas Llama 3.1 wurde am 23. Juli 2024 angekündigt. Das 405B-Modell bietet modernste Leistung bei mehreren öffentlichen Benchmarks und hat ein Kontextfenster von 128.000 Eingabe-Token, wobei verschiedene kommerzielle Verwendungen zulässig sind. Neben dem 405-Milliarden-Parameter-Modell hat Meta auch eine aktualisierte Version von Llama3 70B (70 Milliarden Parameter) und 8B (8 Milliarden Parameter) veröffentlicht. Die Modellgewichte stehen auf der Website von Meta zum Download bereit.
Eine wichtige Erkenntnis war, dass eine Feinabstimmung der generierten Daten die Leistung steigern kann, Beispiele von schlechter Qualität sie jedoch verschlechtern können. Das Llama-Team hat intensiv daran gearbeitet, diese schlechten Beispiele mithilfe des Modells selbst, von Hilfsmodellen und anderen Tools zu identifizieren und zu entfernen.
Aufbau und Durchführung des RAG-Retrievals mit Milvus
Bereiten Sie Ihren Datensatz vor.
Ich habe die offizielle Milvus-Dokumentation als Datensatz für diese Demo verwendet, die ich heruntergeladen und lokal gespeichert habe.
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'}
Laden Sie ein Einbettungsmodell herunter.
Laden Sie als Nächstes ein kostenloses, quelloffenes Einbettungsmodell von HuggingFace herunter.
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
Zerlegen und kodieren Sie Ihre eigenen Daten als Vektoren.
Ich werde eine feste Länge von 512 Zeichen mit 10 % Überlappung verwenden.
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.
Speichern Sie die Vektoren in Milvus.
Nehmen Sie die kodierte Vektoreinbettung in die Milvus-Vektordatenbank auf.
# 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
Führen Sie eine Vektorsuche durch.
Stellen Sie eine Frage und suchen Sie nach den nächstgelegenen Chunks aus Ihrer Wissensdatenbank in 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")
Das Ergebnis sieht wie unten dargestellt aus.
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
Erstellen und Durchführen der RAG-Generierung mit vLLM und Llama 3.1-8B
Installieren Sie vLLM und die Modelle von HuggingFace
vLLM lädt standardmäßig große Sprachmodelle von HuggingFace herunter. Wenn Sie ein neues Modell von HuggingFace verwenden wollen, sollten Sie pip install --upgrade oder -U ausführen. Außerdem benötigen Sie eine GPU, um die Inferenz von Metas Llama 3.1-Modellen mit vLLM durchzuführen.
Eine vollständige Liste aller von vLLM unterstützten Modelle finden Sie auf dieser Dokumentationsseite.
# (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
Weitere Informationen über die Installation von vLLM finden Sie auf der Installationsseite.
Holen Sie sich ein HuggingFace-Token.
Bei einigen Modellen auf HuggingFace, wie Meta Llama 3.1, muss der Benutzer die Lizenz akzeptieren, bevor er die Gewichte herunterladen kann. Daher müssen Sie ein HuggingFace-Konto erstellen, die Lizenz des Modells akzeptieren und ein Token generieren.
Wenn du diese Llama3.1-Seite auf HuggingFace besuchst, erhältst du eine Nachricht, in der du aufgefordert wirst, den Bedingungen zuzustimmen. Klicken Sie auf "Lizenz akzeptieren", um die Meta-Bedingungen zu akzeptieren, bevor Sie die Modellgewichte herunterladen. Die Genehmigung dauert normalerweise weniger als einen Tag.
Nachdem Sie die Genehmigung erhalten haben, müssen Sie ein neues HuggingFace-Token erstellen. Ihre alten Token funktionieren nicht mehr mit den neuen Berechtigungen.
Bevor Sie vLLM installieren, melden Sie sich bei HuggingFace mit Ihrem neuen Token an. Im Folgenden habe ich Colab-Geheimnisse verwendet, um das Token zu speichern.
# 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)
Ausführen der RAG-Generierung
In der Demo führen wir das Modell Llama-3.1-8B aus, das einen Grafikprozessor und einen großen Arbeitsspeicher benötigt, um zu laufen. Das folgende Beispiel wurde auf Google Colab Pro ($10/Monat) mit einer A100 GPU ausgeführt. Weitere Informationen über die Ausführung von vLLM finden Sie in der Quickstart-Dokumentation.
# 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)
Schreiben Sie eine Eingabeaufforderung unter Verwendung von Kontexten und Quellen, die von Milvus abgerufen wurden.
# 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]
Generieren Sie nun eine Antwort unter Verwendung der abgerufenen Chunks und der ursprünglichen Frage, die in die Eingabeaufforderung eingefügt wurde.
# 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.'
Die obige Antwort sieht für mich perfekt aus!
Wenn Sie an dieser Demo interessiert sind, können Sie sie gerne selbst ausprobieren und uns Ihre Meinung mitteilen. Sie sind auch herzlich eingeladen, unserer Milvus-Community auf Discord beizutreten, um sich direkt mit den GenAI-Entwicklern auszutauschen.
Referenzen
vLLM offizielle Dokumentation und Modellseite.
2023 vLLM-Präsentation auf dem Ray Summit
vLLM-Blog: vLLM: Einfaches, schnelles und günstiges LLM Serving mit PagedAttention
Hilfreicher Blog über den Betrieb des vLLM-Servers: Bereitstellung von vLLM: eine schrittweise Anleitung
