Foundation Model benchmarking tool (FMBT) built using Amazon SageMaker

A key challenge with FMs is the ability to benchmark their performance in terms of inference latency, throughput and cost so as to determine which model running with what combination of the hardware and serving stack provides the best price-performance combination for a given workload.

Stated as business problem, the ask is “What is the dollar cost per transaction for a given generative AI workload that serves a given number of users while keeping the response time under a target threshold?”

But to really answer this question, we need to answer an engineering question (an optimization problem, actually) corresponding to this business problem: “What is the minimum number of instances N, of most cost optimal instance type T, that are needed to serve a workload W while keeping the average transaction latency under L seconds?”

W: = {R transactions per-minute, average prompt token length P, average generation token length G}

This foundation model benchmarking tool (a.k.a. FMBT) is a tool to answer the above engineering question and thus answer the original business question about how to get the best price performance for a given workload. Here is one of the plots generated by FMBT to help answer the above question (the numbers on the y-axis, transactions per minute and latency have been removed from the image below, you can find them in the actual plot generated on running FMBT).

Description

The FMBT is a Python package for running performance benchmarks for any model on any supported instance type (g5, p4d, Inf2). FMBT deploys models on Amazon SageMaker and use the endpoint to send inference requests to and measure metrics such as inference latency, error rate, transactions per second etc. for different combinations of instance type, inference container and settings such as tensor parallelism etc. Because FMBT works for any model therefore it can be used not only testing third party models available on SageMaker, open-source models but also proprietary models trained by enterprises on their own data.

>In a production system you may choose to deploy models outside of SageMaker such as on EC2 or EKS but even in those scenarios the benchmarking results from this tool can be used as a guide for determining the optimal instance type and serving stack (inference containers, configuration).

FMBT can be run on any AWS platform where we can run Python, such as Amazon EC2, Amazon SageMaker, or even the AWS CloudShell. It is important to run this tool on an AWS platform so that internet round trip time does not get included in the end-to-end response time latency.

The workflow for FMBT is as follows:

Create configuration file
        |
        |-----> Deploy model on SageMaker
                    |
                    |-----> Run inference against deployed endpoint(s)
                                     |
                                     |------> Create a benchmarking report

1. Create a dataset of different prompt sizes and select one or more such datasets for running the tests.

   1. Currently FMBT supports datasets from LongBench and filter out individual items from the dataset based on their size in tokens (for example, prompts less than 500 tokens, between 500 to 1000 tokens and so on and so forth). Alternatively, you can download the folder from this link to load the data.

2. Deploy any model that is deployable on SageMaker on any supported instance type (g5, p4d, Inf2).

   1. Models could be either available via SageMaker JumpStart (list available here) as well as models not available via JumpStart but still deployable on SageMaker through the low level boto3 (Python) SDK (Bring Your Own Script).
   2. Model deployment is completely configurable in terms of the inference container to use, environment variable to set, setting.properties file to provide (for inference containers such as DJL that use it) and instance type to use.

3. Benchmark FM performance in terms of inference latency, transactions per minute and dollar cost per transaction for any FM that can be deployed on SageMaker.

   1. Tests are run for each combination of the configured concurrency levels i.e. transactions (inference requests) sent to the endpoint in parallel and dataset. For example, run multiple datasets of say prompt sizes between 3000 to 4000 tokens at concurrency levels of 1, 2, 4, 6, 8 etc. so as to test how many transactions of what token length can the endpoint handle while still maintaining an acceptable level of inference latency.

4. Generate a report that compares and contrasts the performance of the model over different test configurations and stores the reports in an Amazon S3 bucket.

   1. The report is generated in the Markdown format and consists of plots, tables and text that highlight the key results and provide an overall recommendation on what is the best combination of instance type and serving stack to use for the model under stack for a dataset of interest.
   2. The report is created as an artifact of reproducible research so that anyone having access to the model, instance type and serving stack can run the code and recreate the same results and report.

5. Multiple configuration files that can be used as reference for benchmarking new models and instance types.

Getting started

FMBT is available as a Python package on PyPi and is run as a command line tool once it is installed. All data that includes metrics, reports and results are stored in an Amazon S3 bucket.

Prerequisites

Follow the prerequisites below to set up your environment before running the code:

1. Python 3.11: Setup a Python 3.11 virtual environment and install FMBT.

   python -m venv .fmbt
   pip install fmbt==1.0.0
   

2. S3 buckets for test data, scripts, and results: Create two buckets within your AWS account:

   • Read bucket: This bucket contains tokenizer files, prompt template, source data and deployment scripts stored in a directory structure as shown below. FMBT needs to have read access to this bucket.

     s3://
         ├── source_data/
         ├── source_data/.json
         ├── prompt_template/
         ├── prompt_template/prompt_template.txt
         ├── scripts/
         ├── scripts/.py
         ├── tokenizer/
         ├── tokenizer/tokenizer.json
         ├── tokenizer/config.json
     

     • The details of the bucket structure is as follows:

       1. Source Data Directory: Create a source_data directory that stores the dataset you want to benchmark with. FMBT uses Q&A datasets from the LongBench dataset or alternatively from this link. Support for bring your own dataset will be added soon.

          • Download the different files specified in the LongBench dataset into the source_data directory. Following is a good list to get started with:

            • 2wikimqa
            • hotpotqa
            • narrativeqa
            • triviaqa

            Store these files in the source_data directory.

       2. Prompt Template Directory: Create a prompt_template directory that contains a prompt_template.txt file. This .txt file contains the prompt template that your specific model supports. FMBT already supports the prompt template compatible with Llama models.

       3. Scripts Directory: FMBT also supports a bring your own script (BYOS) mode for deploying models that are not natively available via SageMaker JumpStart i.e. anything not included in this list. Here are the steps to use BYOS.

          1. Create a Python script to deploy your model on a SageMaker endpoint. This script needs to have a deploy function that 1_deploy_model.ipynb can invoke. See p4d_hf_tgi.py for reference.

          2. Place your deployment script in the scripts directory in your read bucket. If your script deploys a model directly from HuggingFace and needs to have access to a HuggingFace auth token, then create a file called hf_token.txt and put the auth token in that file. The .gitignore file in this repo has rules to not commit the hf_token.txt to the repo. Today, FMBT provides inference scripts for:

             • All SageMaker Jumpstart Models
             • Text-Generation-Inference (TGI) container supported models
             • Deep Java Library DeepSpeed container supported models

             Deployment scripts for the options above are available in the scripts directory, you can use these as reference for creating your own deployment scripts as well.

       4. Tokenizer Directory: Place the tokenizer.json, config.json and any other files required for your model's tokenizer in the tokenizer directory. The tokenizer for your model should be compatible with the tokenizers package. FMBT uses AutoTokenizer.from_pretrained to load the tokenizer. >As an example, to use the Llama 2 Tokenizer for counting prompt and generation tokens for the Llama 2 family of models: Accept the License here: meta approval form and download the tokenizer.json and config.json files from Hugging Face website and place them in the tokenizer directory.

   • Write bucket: All prompt payloads, model endpoint and metrics generated by FMBT are stored in this bucket. FMBT requires write permissions to store the results in this bucket. No directory structure needs to be pre-created in this bucket, everything is created by FMBT at runtime.

     s3://
         ├── 
         ├── /data
         ├── /data/metrics
         ├── /data/models
         ├── /data/prompts
     
     

Steps to run

1. pip install the FMBT package from PyPi.

2. Create a config file using one of the config files available here.

   1. The configuration file is a YAML file containing configuration for all steps of the benchmarking process. It is recommended to create a copy of an existing config file and tweak it as necessary to create a new one for your experiment.

3. Create the read and write buckets as mentioned in the prerequisites section. Mention the respective directories for your read and write buckets within the config files.

4. Run the FMBT tool from the command line.

   # the config file path could be an S3 path and https path 
   # or even a path to a file on the local filesystem
   fmbt --config-file \path\to\config\file
   

5. Depending upon the experiments in the config file, the FMBT run may take a few minutes to several hours. Once the run completes, you can find the report and metrics in the write S3 bucket set in the config file. The report is generated as a markdown file called report.md and is available in the metrics directory in the write S3 bucket.

Results

Here is a screenshot of the report.md file generated by FMBT.

Building the FMBT Python package

The following steps describe how to build the FMBT Python package.

1. Clone the FMBT repo from GitHub.

2. Make any code changes as needed.

3. Install poetry.

   pip install poetry
   

4. Change directory to the FMBT repo directory and run poetry build.

   poetry build
   

5. The .whl file is generated in the dist folder. Install the .whl in your current Python environment.

   pip install .\dist\fmbt-X.Y.Z.tar.gz
   

6. Run FMBT as usual through the fmbt CLI command.

Pending enhancements

The following enhancements are identified as next steps for FMBT.

1. [Highest priority] Convert FMBT to a Python package and publish on PyPi.

2. Containerize FMBT and provide instructions for running the container on EC2.

3. Add code to determine the cost of running an entire experiment and include it in the final report. This would only include the cost of running the SageMaker endpoints based on hourly public pricing (the cost of running this code on a notebook or a EC2 is trivial in comparison and can be ignored).

4. Support for a custom token counter. Currently only the LLama tokenizer is supported but we want to allow users to bring their own token counting logic for different models.

5. Support for different payload formats that might be needed for different inference containers. Currently the HF TGI container, and DJL Deep Speed container on SageMaker both use the same format but in future other containers might need a different payload format.

6. Emit live metrics so that they can be monitored through Grafana via live dashboard.

7. Allow users to publish their experiment configs and results by doing a POST to an AWS Lambda that writes results to a common S3 bucket that can serve as storage for a simple website.

8. Create a leaderboard of model benchmarks.

View the ISSUES on github and add any you might think be an beneficial iteration to this benchmarking harness.

Security

See CONTRIBUTING for more information.

License

This library is licensed under the MIT-0 License. See the LICENSE file.