SageMaker JumpStart Foundation Models - Benchmark Latency and Throughput

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Welcome to Amazon SageMaker JumpStart! You can use SageMaker JumpStart to solve many Machine Learning tasks through one-click in SageMaker Studio, or through SageMaker Python SDK.

When testing a large language model for production use cases, common questions arise, such as:

• What is the inference latency for my expected payloads?
• How much throughput does this model configuration provide for my expected payloads?
• What is the inference throughput and latency for my expected concurrency load, i.e., the number of concurrent requests that have invoked the endpoint?
• How much does it cost to generate 1 million tokens?
• How does instance type selection (e.g., ml.g5.2xlarge) affect latency and throughput?
• How does modification of the deployment configuration (e.g., tensor parallel degree) affect latency and throughput?

Given these questions, you may notice that inference latency and throughput depend on numerous factors, to include payload, number of concurrent requests, instance type, deployment configuration, and more. In this notebook, we demonstrate how you can run your own latency and throughput benchmark for SageMaker JumpStart endpoints. This benchmarking process involves running load tests with various concurrent request values for each payload and deployed endpoint.

The primary inputs to this benchmarking tool include the models to benchmark and the payloads used to invoke endpoints.

• MODELS: A dictionary mapping a unique name to benchmarking configuration. The model can be defined in 3 different ways. Each model value should be a dictionary with the following keys:
  • jumpstart_model_specs key: requires model_args and optionally deploy_args definitions to use with a SageMaker SDK JumpStartModel constructor and deploy methods, respectfully. This should be used to deploy and benchmark a JumpStart model.
  • model_specs key: requires image_uri_args, model_args, and deploy_args definitions to use with a SageMaker SDK Model constructor and deploy methods. This should be used to deploy and benchmark a non-JumpStart model.
  • endpoint_name key: provide the endpoint name of a pre-deployed model to benchmark.
  • huggingface_model_id key: to compute metrics with respect to model tokens, provide the HuggingFace Model ID with an appropriate tokenizer to use.
• PAYLOADS: A dictionary mapping a unique name to a payload of interest. The benchmarking tool will serially run a concurrency probe against each payload.

For this notebook, we deploy and benchmark Falcon 7B and Mistral 7B models. We will use two payloads, with input sequence lengths of 512 and 6400 words. Since Falcon 7B can only support up to a 2k context, we expect only Mistral 7B benchmarking to be successful for the large context payload.

Now create a Benchmarker object and run benchmarking for all models. This will first concurrently create a Predictor for all models. If endpoint_name is specified in the MODELS definition or provided in the JSON metrics file of a previous run, the endpoint will be attached to a Predictor. Otherwise, a new endpoint will be deployed. Once an endpoint is in service, it will begin the load test concurrency probe. A concurrency probe will be executed concurrently for all models. For each model, the probe will sweep concurrent request values, performing a load test at each unique value, until an error occurs. These errors may be validation checks (e.g., endpoint is overloaded, input sequence length unsupported, etc.), SageMaker invocation timeout, or any other potential model error. The concurrency probe for each specified payload will run serially for each model. When the probe has completed, all computed metrics will be saved in a JSON file for downstream analysis.

Now that benchmarking is complete, let's load the results into a Pandas DataFrame and create a pivot table that shows throughput, p90 latency, and cost to generate one million tokens. Many variations of these metrics are recorded in the DataFrame, so please extract any information relevant to your benchmarking effort.

Finally, please remember to clean up all model and endpoint resources to avoid incurring additional costs after your benchmarking is complete.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.