Building a RAG Chatbot with Friendli, MongoDB Atlas, and LangChain

pip install langchain langchain-community langchain-mongodb friendli-client pypdf pymongo langchain-openai tiktoken

from pymongo import MongoClient

MONGODB_ATLAS_CLUSTER_URI = "YOUR CLUSTER URI"

client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)

# Fill in your information here
DB_NAME = "my_rag"
COLLECTION_NAME = "test"
ATLAS_VECTOR_SEARCH_INDEX_NAME = "index_name"

MONGODB_COLLECTION = client[DB_NAME][COLLECTION_NAME]

client.server_info()

{
  "fields": [
    {
      "numDimensions": 1536,
      "path": "embedding",
      "similarity": "cosine",
      "type": "vector"
    }
  ]
}

from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

loader = PyPDFLoader("https://openreview.net/pdf?id=HVKmLi1iR4")
data = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
docs = text_splitter.split_documents(data)

vector_store = MongoDBAtlasVectorSearch.from_documents(
    documents=docs,
    embedding=OpenAIEmbeddings(disallowed_special=()),
    collection=MONGODB_COLLECTION,
    index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,
)
retriever = vector_store.as_retriever()

from langchain_community.chat_models.friendli import ChatFriendli

llm = ChatFriendli(model="meta-llama-3.1-70b-instruct")

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough

retriever = vector_store.as_retriever()

template = """Use the following pieces of context to answer the question at the end.
If you don’t know the answer, just say that you don’t know, don’t try to make up an answer.
Use three sentences maximum and keep the answer as concise as possible.
Always say “thanks for asking!” at the end of the answer.

{context}

Question: {question}

Helpful Answer:"""

prompt = PromptTemplate.from_template(template)

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

rag_chain.invoke("What is the memory imbalance problem that BPipe solves?")

The memory imbalance problem refers to a situation where some pipeline stages in a pipeline parallelism setup consume more memory than others, leading to inefficient utilization of GPU memory. This can occur when training large language models (LLMs) using pipeline parallelism, where the model is split into multiple stages, and each stage is processed sequentially on a different GPU.

In a pipeline parallelism setup, each stage typically requires a certain amount of memory to store the model’s parameters, optimizer states, and intermediate activations. However, the memory requirements can vary significantly between stages, depending on the model’s architecture and the specific computations performed at each stage.

The memory imbalance problem arises when some stages require significantly more memory than others, leading to:

1. **Memory underutilization**: GPUs with less memory-intensive stages may not be fully utilized, wasting computational resources.
2. **Memory overflow**: GPUs with more memory-intensive stages may run out of memory, causing the training process to slow down or even fail.

The memory imbalance problem can be exacerbated by the following factors:

* **Model size**: Larger models require more memory, making it more challenging to balance memory usage across pipeline stages.
* **Batch size**: Increasing the batch size can amplify the memory imbalance problem, as more data needs to be stored in memory.
* **Pipeline schedule**: The order in which pipeline stages are executed can affect the memory imbalance, with some schedules leading to more pronounced imbalances than others.

The memory imbalance problem can have significant consequences, including:

* **Slower training times**: Inefficient memory utilization can lead to slower training times, making it more challenging to train large language models.
* **Increased costs**: Underutilized GPUs can result in wasted computational resources, increasing the overall cost of training.

The BPIPE approach, described in the original document, addresses the memory imbalance problem by transferring activations between pipeline stages to balance memory usage and ensure efficient GPU memory utilization.