Optimizing Generative AI Applications with Fireworks AI and MongoDB for Peak Performance

The Power of MongoDB and Fireworks AI

MongoDB is renowned for its flexible schema, efficient indexing, and distributed architecture, allowing organizations to scale their data infrastructure on demand. When paired with Fireworks.AI’s cutting-edge model tuning capabilities, enterprises are equipped to achieve remarkable efficiency and performance gains without financial strain.

Introducing FireOptimizer and FireAttention: Core Components

Central to this optimized approach is the FireOptimizer framework, integrating MongoDB’s flexible, high-performance database with Fireworks.AI’s model tuning prowess.

This framework accelerates batch inferencing through solutions like FireAttention, a method that enhances request processing using existing hardware, optimizing resource utilization. Furthermore, Parameter-Efficient Fine-Tuning (PEFT), including methods like LoRA and QLoRA, fine-tunes models to complete specialized tasks efficiently using trace data or label data, thereby enhancing portability and reducing computational footprint.

Figure 1. FireOptimizer Architecture for Adaptive Optimization and High-Quality Inference

The key benefits that can be achieved are:

• Faster Inference: With FireOptimizer’s adaptive speculative execution, production workloads across various models saw up to 3x latency improvements, ensuring applications are always highly responsive.

• Hassle-Free Optimization: FireOptimizer automates the complex optimization process, so you can focus on building your application, while we handle the rest.

Let’s dive deeper into the FireOptimizer framework.

FireOptimizer: Smarter Token Generation Through Adaptation

Large Language Models (LLMs) typically generate outputs one word (or token) at a time , which can make responses slower, especially for long outputs. To accelerate this process, a technique called speculative decoding is used. This involves a smaller, faster model—known as a draft model—that quickly generates several possible words or phrases in advance while the main LLM is still processing. The main LLM then evaluates these guesses and retains only the accurate ones. This technique is commonly referred to as speculative execution or decoding.

The limitation of this traditional approach lies in the fact that most draft models are trained on generic data (e.g., public conversations). While this approach is effective for general tasks, its accuracy or "hit rate" diminishes significantly in specialized scenarios, such as coding or financial analysis.

In contrast, Fireworks AI enhances this approach through adaptive speculative execution. Instead of using a one-size-fits-all draft model, it uses models that are customized to specific domains or user profiles. By using this adaptive approach, it optimizes predictions based on domain-specific characteristics. This method improves accuracy, increases hit rates (e.g., from 29% to 76% in code generation tasks), and reduces inference costs, leading to latency improvements of up to 3x.

Proprietary LLM serving stack: Mastering Long-Context Workloads

Processing long prompts (8K-32K tokens) is critical for applications like document analysis or code completion—but often slows systems to a crawl. Fireworks AI’s FireAttention V2 tackles this, delivering 12x faster processing for long-context tasks.

Fine-Tuning Strategies for Smaller, Efficient Models : A Practical Credit Application demo with MongoDB

Smaller, efficient models present unique opportunities for fine-tuning, allowing for specialized adaptation while maintaining resource efficiency. There has been constant research in this space as making finetuned or optimized SLM to various mediums such as cloud, on-device, dedicated hardware is becoming extremely competitive. These fine tuning techniques can be categorized as follows:

These fine tuning techniques can be categorized as follows:

The most popular technique of these is the PEFT based ones. PEFT techniques used with large pre-trained models can adapt them to new tasks by only adjusting a small fraction of their parameters. This approach helps prevent overfitting, especially with smaller datasets, and significantly reduces computational and memory requirements compared to full fine-tuning. PEFT also mitigates catastrophic forgetting in LLMs. These techniques enable efficient model customization without requiring full retraining, making them ideal for resource-constrained environments.

We can adopt the PEFT LoRA techniques combined with the availability of both trace data (data generated from the model's interactions) and labeled data (data explicitly annotated for specific tasks), enable the fine-tuning of smaller models to achieve high performance on specific tasks without the need for extensive computational resources.

To highlight the practical applications, we turn to a MongoDB credit card application demo that illustrates how MongoDB can be used for credit scoring with predictive analytics, explaining credit scoring results using generative AI and performing credit card recommendations using a RAG approach with LLM. For the purpose of this fine tuning example, we will be focusing on simplifying credit rating explanations using LLMs. The application encompasses various feature scopes i.e. user profile generation, product recommendations, and reranking with summarization tasks. Details of the application design and source code can be found at the MongoDB Credit Card Application Solution Library Page.

Step-by-Step guide for Fine-Tuning a Small Language Model with Fireworks AI

A critical aspect of the Credit Card Application demo is explaining credit scores or credit ratings to the customer in clear, everyday language. Frontier LLMs like Meta's LLaMA 3.1-405B can generate these explanations using factors such as user profile parameters, model input features, and feature importances from the model used to predict the customer's alternative credit score or credit rating. However, these tasks cannot be consistently achieved with a SLMs due to the limited number of parameters required for effective reasoning and explanation. To achieve the desired outcome, we need to leverage the fine-tuning process described above using the Fireworks AI fine-tuning platform.

The following steps will go through step-by-step, how one can make use of the platform to fine tune a SLM.

Here’s how to leverage this platform and tool:

Figure 2. LLM/SLM Fine Tuning Process

The fine-tuning process starts with the collection of relevant, task-specific data. As shown in Figure 2, MongoDB Atlas, being a general-purpose database, can be leveraged to cache LLM/SLM responses for a specific user based on their inputs to the Credit Card application (the user is allowed to simulate credit profile of user on the web UI). Below is a sample Python code snippet demonstrating how to set up a decorator to cache LLM/SLM responses in MongoDB Atlas:

class mdbcache:
   def __init__(self, function):
         self.function = function
   def __call__(self, *args, **kwargs):
         key = str(args) + str(kwargs)
         ele = ccol.find_one({"key": key})
         if ele:
            return ele["response"]
         value = self.function(*args, **kwargs)
         ccol.insert_one({"key":key, "response": value})
         return value
@mdbcache
def invoke_llm(prompt):
   """
   Invoke the language model with the given prompt with cache. The llm.invoke  method can invoke either a LLM or SLM based on the Fireworks Model ID provided at the start of applicaiton.
   Args:
         prompt (str): The prompt to pass to the LLM.
   """
   response = llm.invoke(prompt)
   return response

As shown in the diagram, the training dataset required for the fine-tuning process can be generated using a simulator. Here, we are trying to simulate a user profile with a stratified sampling approach to select an equal number of samples for all three credit ratings: Good, Normal, and Poor. For this demo, we have generated around 1300+ sample responses.

Now, the generated responses have to be transformed into the format that the Fireworks AI platform supports for starting the fine-tuning process. The file "cc_cache.jsonl" used below in the fine-tuning process can be generated by running the following code:

 from pymongo import MongoClient
 import pandas as pd
 import json
 client = MongoClient("mongodb+srv://<uid>:<pwd>@bfsi-demo.2wqno.mongodb.net/?retryWrites=true&w=majority")
 df = pd.DataFrame.from_records(client["bfsi-genai"]["cc_cache"].find({},{"_id": 0}))
 df["prompt"] = df["key"].apply(lambda x: x.strip('(').strip('"').strip(")").strip("\\"))
 del df["key"]
 df["response"] = df["response"].apply(lambda x: x.strip())
 df.to_json("cc_cache.jsonl", orient="records", lines=True)
 # transform cache to messages
 messages = []
 for item in df.iterrows():
   messages += [{"messages": [{"role": "user", "content": item["prompt"].strip("    \\")}, {"role": "assistant", "content": item["response"]}]}]
 with open("cc_cache.jsonl", "w") as f:
   for item in messages:
      f.write(json.dumps(item) + "\n")

Now that we have prepared the dataset and generated our "cc_cache.jsonl" file, we can start fine-tuning the pre-trained "llama-v3p1-8b-instruct" model by following the steps below.

1

Install firectl Install firectl, the command-line tool for managing Fireworks AI platform resources:

pip install firectl

2

Authenticate firectl with your Fireworks account:

firectl login

3

Upload your fine-tuning dataset that was created as part of data generation process to the Fireworks platform:

firectl create dataset <dataset_name> cc_cache.jsonl

4

Create a fine-tuning job with the uploaded dataset and your chosen model:

firectl create sftj --base-model accounts/fireworks/models/llama-v3p1-8b-instruct --dataset <dataset_name> --output-model ccmodel --lora-rank 8 --epochs 1

5

Monitor the fine-tuning process using the portal to ensure it runs smoothly:

Figure 3. Monitoring the fine-tuning process

6

Deploy the fine-tuned model on the Fireworks platform for inference:

firectl deploy ccmodel

The seamless integration of MongoDB and Fireworks.AI streamlines the process of enhancing AI model performance in a cost-effective and efficient manner. This is evident in the step-by-step guide to fine-tuning provided above.

After deploying the model in the fireworks platform as a serverless API, the model id (models/ft-m88hxaga-pi11m) shown in Fig 2 can be used to invoke the fine-tuned SLM model using the language model framework of your choice.

Further Reading

How Fireworks AI's platform accelerates inference speeds and throughput:

• FireAttention V2: This innovation improves the efficiency of handling long contexts during online inference, achieving up to 12 times faster processing.

• FireAttention V3: This advancement enables the use of AMD hardware as a viable alternative for GPU inference, broadening hardware options and potentially reducing costs.

• FireOptimizer: This tool allows customization of latency and quality for production inference workloads, enabling users to tailor performance to their specific needs.

MongoDB Developer Data Platform

• MongoDB Atlas

Partner Technologies

• Fireworks AI

• LangChain

MongoDB:

• Wei You Pan

• Ashwin Gangadhar

• Peyman Parsi

• Andrea Alaman Calderon

Fireworks AI:

• Benny Chen

• Ayaan Momin