notebooks/Mistral-on-AWS-Neuron/Mathstral/Compile-Deploy-Benchmark-Mathstral-on-Neuron.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "add8a615-7d97-431c-b928-b743543a0e6b",
   "metadata": {},
   "source": [
    "# Compile, Deploy, and Benchmark Mathstral on Inferentia2 with Optimum Neuron and SageMaker\n",
    "---\n",
    "\n",
    "In this notebook, we walk through the basics of how you can get started with compiling models for AWS Neuron to deploy on Inferentia2 instances.\n",
    "AWS Neuron is an SDK with a compiler, runtime, and profiling tools that unlocks high-performance and cost-effective deep learning (DL) acceleration. It supports high-performance training on AWS Trainium-based Amazon Elastic Compute Cloud (Amazon EC2) Trn1 instances. For model deployment, it supports high-performance and low-latency inference on AWS Inferentia-based Amazon EC2 Inf1 instances and AWS Inferentia2-based Amazon EC2 Inf2 instances. With Neuron, you can use popular frameworks, such as TensorFlow and PyTorch, and optimally train and deploy machine learning (ML) models on Amazon EC2 Trn1, Inf1, and Inf2 instances, and Neuron is designed to minimize code changes and tie-in to vendor-specific solutions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d390157f-b669-410b-a62e-7cfa72d18605",
   "metadata": {},
   "source": [
    "In this sample, we will demonstrate compilation and benchmarking with the [mistralai/Mathstral-7B-v0.1](https://huggingface.co/mistralai/Mathstral-7B-v0.1) model to [Amazon SageMaker](https://aws.amazon.com/sagemaker/). We will utilize the Hugging Face LLM DLC, a purpose-built Inference Container designed to facilitate the deployment of Large Language Models (LLMs) in a secure and managed environment. This Deep Learning Container (DLC) is powered by <b>Text Generation Inference (TGI)</b>, a scalable and optimized solution for deploying and serving LLMs efficiently. Detailed instance requirements for various model sizes will also be provided to ensure optimal deployment configurations. We will be using the Ray/llmperf tool for benchmarking performance of our sagemaker endpoint with inf2."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e352b305-4e8a-4252-82a9-98b9cf818bb5",
   "metadata": {},
   "source": [
    "## Set up environment\n",
    "\n",
    "#### Local Setup (Optional)\n",
    "\n",
    "For a local server, follow these steps to execute this jupyter notebook:\n",
    "\n",
    "1. **Configure AWS CLI**: Configure [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html) with your AWS credentials. Run `aws configure` and enter your AWS Access Key ID, AWS Secret Access Key, AWS Region, and default output format.\n",
    "\n",
    "2. **Install required libraries**: Install the necessary Python libraries for working with SageMaker, such as [sagemaker](https://github.com/aws/sagemaker-python-sdk/), [boto3](https://github.com/boto/boto3), and others. You can use a Python environment manager like [conda](https://docs.conda.io/en/latest/) or [virtualenv](https://virtualenv.pypa.io/en/latest/) to manage your Python packages in your preferred IDE (e.g. [Visual Studio Code](https://code.visualstudio.com/)).\n",
    "\n",
    "3. **Create an IAM role for SageMaker**: Create an AWS Identity and Access Management (IAM) role that grants your user [SageMaker permissions](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html). \n",
    "\n",
    "By following these steps, you can set up a local Jupyter Notebook environment capable of deploying machine learning models on Amazon SageMaker using the appropriate IAM role for granting the necessary permissions.\n",
    "\n",
    "---\n",
    "\n",
    "#### Prerequisites\n",
    "\n",
    "This Jupyter Notebook can be run on a t3.medium instance (ml.t3.medium). However, to deploy mistral models to an Inf2 endpoint, you may need to request a quota increase. \n",
    "\n",
    "To request a quota increase, follow these steps:\n",
    "\n",
    "1. Navigate to the [Service Quotas console](https://console.aws.amazon.com/servicequotas/).\n",
    "2. Choose Amazon SageMaker.\n",
    "3. Review your default quota for the following resources:\n",
    "   - `ml.inf2.24xlarge` for endpoint usage `or`\n",
    "   - `ml.inf2.48xlarge` for endpoint usage\n",
    "   \n",
    "    Note that although this example showcases Mathstral ( which has been compiled to run on 12 Neuron cores) with inf2.24xlarge, you are still able to deploy and benchmark other models compiled for Neuron in our [neuron-compile-jobs](https://huggingface.co/collections/nithiyn/neuron-compile-jobs-66fc4163c5350829c9121e80) collection for Mistral models on HuggingFace.\n",
    "4. If needed, request a quota increase for these resources.\n",
    "\n",
    "---\n",
    "\n",
    "#### Requirements\n",
    "\n",
    "If using the `sagemaker` python SDK to deploy Mistral model compiled for AWS Neuron to Amazon SageMaker, we need to make sure to have an AWS account configured and the `sagemaker` python SDK installed. \n",
    "\n",
    "1. Create an Amazon SageMaker Notebook Instance \n",
    "- [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-setup-working-env.html)\n",
    "- For Notebook Instance type, choose `(ml.m5.4xlarge)`, since we will be benchmarking performance in this notebook.\n",
    "    \n",
    "2. For Select Kernel, choose [conda_pytorch_p310](https://docs.aws.amazon.com/sagemaker/latest/dg/ex1-prepare.html).\n",
    "\n",
    "3. Install the required packages.\n",
    "\n",
    "4. Set up your [HuggingFace token](https://huggingface.co/docs/transformers.js/en/guides/private): \n",
    "- User Access Tokens are the preferred way to authenticate an application to Hugging Face services.\n",
    "- To generate an access token, navigate to the Access Tokens tab in your settings and click on the New token button.\n",
    "- Choose a name for your token and click Generate a token (we recommend keeping the “Role” as read-only). You can then click the Copy button next to your newly-created token to copy it to your clipboard.\n",
    "- Copy and replace this token below in the `HF_TOKEN` and `HUGGING_FACE_HUB_TOKEN` parameter under the optimum neuron compile section and the `config` in the deployment section.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c85789-9ffd-4da8-9709-4e14d1c1fcc6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Install packages and import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce65ec1-416d-40cd-b882-43e6787e7db3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install sagemaker --upgrade --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b60518d7-0602-4fc7-baa8-d4c71b1fe4b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install transformers --upgrade --quiet\n",
    "!pip install gradio --upgrade --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eefdf2d-e151-4005-803b-b89550e4cb7b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import gradio as gr\n",
    "import json\n",
    "import os\n",
    "import sagemaker\n",
    "import sys\n",
    "import time\n",
    "from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri\n",
    "from sagemaker.huggingface import HuggingFaceModel\n",
    "\n",
    "\n",
    "print(sagemaker.__version__)\n",
    "if not sagemaker.__version__ >= \"2.232.0\": print(\"You need to upgrade or restart the kernel if you already upgraded\")\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "role = sagemaker.get_execution_role()\n",
    "bucket = sess.default_bucket()\n",
    "region = sess.boto_region_name\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0319cf46-5d9c-4233-9037-3d97038ba784",
   "metadata": {},
   "source": [
    "We use v0.025 of the tgi-neuronx image from the ecr deep learning container repository, since we need neuronx-cc> 2.15 since this is what is being used to compile our models for inf2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2efeb8-0a40-4568-95e8-b16f86f1a2e0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import get_huggingface_llm_image_uri\n",
    "# use the latest huggingface image for neuronx\n",
    "llm_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.2-optimum0.0.25-neuronx-py310-ubuntu22.04-v1.0\"\n",
    " \n",
    "# print ecr image uri\n",
    "print(f\"llm image uri: {llm_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32aaa278-3de5-46b0-88ec-b2153f8929b9",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compile models for Optimum Neuron - optional\n",
    "AWS Inferentia2 does not support dynamic shapes for inference, which means that we need to specify our sequence length and batch size ahead of time. \n",
    "For ease of use, our [Mistral-on-AWS](https://github.com/aws-samples/mistral-on-aws) team has pre-compiled these models to Neuron for your use. In order to be able to compile your models to NEFF(Neuron Executable File Format), follow the steps below:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b54f56f6-46ea-4211-9ec3-167ad4df5e98",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-block alert-warning\"> \n",
    "\n",
    "<b>NOTE:</b> This section is optional, we have already compiled this model for you in our [neuron-compile-jobs](https://huggingface.co/collections/nithiyn/neuron-compile-jobs-66fc4163c5350829c9121e80) collection for Mistral models on HuggingFace.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de60d36c-424e-4dec-a333-bbd4729ab9e9",
   "metadata": {},
   "source": [
    "### Prerequisites for Compilation\n",
    "Follow the steps below to set up your EC2 instance with the HuggingFace Neuron DLAMI from the AWS Marketplace.\n",
    "\n",
    "---\n",
    "#### Create Your EC2 instance\n",
    "##### Follow the steps here for a detailed set up of your EC2 instance: [setup](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/EC2_GetStarted.html)\n",
    "\n",
    "##### Steps:\n",
    "- Navigate to the EC2 dashboard from the AWS mgmt console and launch your instance.\n",
    "- Search for the [HuggingFace Neuron DLAMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2).\n",
    "-  Choose the instance size as `inf2.24xlarge/inf2.48xlarge` or any other AWS Neuron based instances.\n",
    "- Set the inbound rule for `ssh` to your local machine's ip address or `anywhere` (note that it is not in accordance to set this to allow trafic from any ipv4, please ensure you secure these ports once done testing.\n",
    "- Create and specify your ssh key in the instance configuration step. You will need your `.pem` file\n",
    "- Create your instance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d8f3a9f-2e08-4027-a7d8-d5cc9d2d30de",
   "metadata": {},
   "source": [
    "Once you have launched your instance, navigate to either your terminal or VSCODE and follow the steps below:\n",
    "\n",
    "<b>ssh for powershell:</b>\n",
    "```\n",
    "$PUBLIC_DNS=\"paste your public ipv4 dns here\" # public ipv4 DNS, e.g. ec2-3-80-.... from ec2 console\n",
    "$KEY_PATH=\"paste ssh key path here\" # local path to key, e.g. ssh/trn.pem\n",
    "\n",
    "ssh -i $KEY_PATH -L 8080:localhost:8080 ubuntu@$PUBLIC_DNS\n",
    "```\n",
    "<b>ssh for linux/macOS:</b>\n",
    "```\n",
    "export PUBLIC_DNS=\"paste your public ipv4 dns here\" # public ipv4 DNS, e.g. ec2-3-80-.... from ec2 console\n",
    "export KEY_PATH=\"paste ssh key path here\" # local path to key, e.g. ssh/trn.pem\n",
    "\n",
    "ssh -i $KEY_PATH -L 8080:localhost:8080 ubuntu@$PUBLIC_DNS\n",
    "``` \n",
    "You should have sshed into your EC2 instance.\n",
    "Next we can change our directory to home, navigate to \n",
    "\n",
    "```\n",
    "(aws_neuronx_venv_pytorch_2_1) ubuntu@ip-172-31-0-5:~$ cd huggingface-neuron-notebooks/\n",
    "(aws_neuronx_venv_pytorch_2_1) ubuntu@ip-172-31-0-5:~/huggingface-neuron-notebooks$ cd text-generation/\n",
    "(aws_neuronx_venv_pytorch_2_1) ubuntu@ip-172-31-0-5:~/huggingface-neuron-notebooks/text-generation$ python -m notebook --allow-root --port=8080\n",
    "```\n",
    "You should see a familiar jupyter output with a URL to the notebook.\n",
    "\n",
    "`http://localhost:8080/....`\n",
    "\n",
    "We can click on it, and a jupyter environment opens in our local browser. Upload this notebook to your jupyter environment and run the steps in the cells below by modifying it for the model you would like to compile for Neuron."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b37ba0e-c1c0-4d37-80be-3042248203e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ./mistral-model #set the name of the model as directory name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d90e0367-bf47-47f5-81ed-bcdccf000bfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf /var/tmp/neuron-compile-cache/* # clear neuron cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e64ccb94-5396-4462-ace5-e6143986fdac",
   "metadata": {},
   "outputs": [],
   "source": [
    "!optimum-cli neuron cache lookup mistralai/#look up the mistral model you would like to compile to see if it is already in the neuron persistent cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d1f4793-a2cc-4cbc-bde2-2f8f4d3cdfa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Replace the empty values below with the input shapes you would like for your model. for the input shapes used for mathstral, refer to the next section\n",
    "MODEL_ID = \"\"#HF model ID for the mistral model you would like to \n",
    "SEQUENCE_LENGTH =\"\" # Sequence length that the Neuronx-cc compiler exported model will be able to take as input.\n",
    "BATCH_SIZE =\"\" # Batch size that the Neuronx-cc compiler exported model will be able to take as input.\n",
    "NUM_CORES =\"\" # each inferentia chip has 2 cores, e.g. inf2.48xlarge has 12 chips or 24 cores\n",
    "PRECISION =\"\" # fp32/bf16/fp16 depending on the precision\n",
    "HF_MODEL_ID_TO_PUSH =\"\" # change this to your desired model id/\n",
    "HF_TOKEN =\"\" #HF_TOKEN to use that you generate in requirments\n",
    " \n",
    "# login into the huggingface hub to access gated models, like llama or mistral\n",
    "!huggingface-cli login --token $HF_TOKEN\n",
    "# compile model with optimum for batch size 4 and sequence length 2048\n",
    "!optimum-cli export neuron -m {MODEL_ID} --batch_size {BATCH_SIZE} --sequence_length {SEQUENCE_LENGTH} --num_cores {NUM_CORES} --auto_cast_type {PRECISION} ./mistral-model\n",
    "# push model to hub [repo_id] [local_path] [path_in_repo]\n",
    "!huggingface-cli upload {HF_MODEL_ID_TO_PUSH} ./mistral-model ./"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "529c838f-902e-43f6-a76d-c06356e44c2a",
   "metadata": {},
   "source": [
    "Once you run the above cells in your jupyter server, your compile job should finish and push your model to the hub under the model ID that you have specified. In the case that you would like to continue without compiling the model yourself, our Mistral-on-AWS team has created a collection of neuron compiled NEFF binaries for you to use [here](https://huggingface.co/collections/nithiyn/neuron-compile-jobs-66fc4163c5350829c9121e80). This collection is a work in progress and we will continue adding models compiled for neuron to this repository."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2774d748-97f0-4ce1-ab42-79125d40c16b",
   "metadata": {},
   "source": [
    "## Deploying Your Model to an Endpoint\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec0a6f63-b3c5-4573-bb8c-ad8674961d04",
   "metadata": {},
   "source": [
    "In this example, we are deploying the [mistralai/Mathstral-7B-v0.1](https://huggingface.co/mistralai/Mathstral-7B-v0.1) model to an Inf2.24xlarge. This model has been exported to the neuron format by the Mistral-on-AWS team using specific input_shapes and compiler parameters detailed in the paragraphs below.\n",
    "\n",
    "It has been compiled to run on an inf2.24xlarge instance on AWS. \n",
    "Note that this compilation uses 24 cores. \n",
    "\n",
    "For demonstration purposes, we have compiled it with the below input shapes, feel free to recompile as needed.\n",
    "\n",
    "These input shapes are as below:\n",
    "\n",
    "SEQUENCE_LENGTH = 4096\n",
    "\n",
    "BATCH_SIZE = 4\n",
    "\n",
    "NUM_CORES = 12\n",
    "\n",
    "PRECISION = \"bf16\"\n",
    "\n",
    "In production environments, to deploy Huggingface 🤗 Transformers models on Neuron devices, you need to compile your models and export them to a serialized format before inference. Through Ahead-Of-Time (AOT) compilation with Neuron Compiler( neuronx-cc or neuron-cc ), models are converted to serialized and optimized TorchScript modules.\n",
    "\n",
    "Although pre-compilation avoids overhead during the inference, a compiled Neuron model has some limitations:\n",
    "\n",
    "- The input shapes and data types used during the compilation cannot be changed.\n",
    "\n",
    "- Neuron models are specialized for each hardware and SDK version, which means:\n",
    "1. Models compiled with Neuron can no longer be executed in non-Neuron environment.\n",
    "2. Models compiled for inf1 (NeuronCore-v1) are not compatible with inf2 (NeuronCore-v2), and vice versa.\n",
    "3. Models compiled for an SDK version are (generally) not compatible with another SDK version\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f16becb-3b1f-4a52-aa5f-a88d5a5b2527",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import HfFolder\n",
    "from sagemaker.huggingface import HuggingFaceModel\n",
    " \n",
    "# sagemaker config\n",
    "instance_type = \"ml.inf2.24xlarge\"\n",
    "health_check_timeout=2400 # additional time to load the model\n",
    "volume_size=512 # size in GB of the EBS volume\n",
    " \n",
    "# Define Model and Endpoint configuration parameter\n",
    "config = {\n",
    "    \"HF_MODEL_ID\": \"nithiyn/mathstral-neuron\", # replace with your model id if you are using your own model\n",
    "    \"HF_NUM_CORES\": \"12\", # number of neuron cores\n",
    "    \"HF_AUTO_CAST_TYPE\": \"bf16\",  # dtype of the model\n",
    "    \"MAX_BATCH_SIZE\": \"4\", # max batch size for the model\n",
    "    'HUGGING_FACE_HUB_TOKEN': \"REPLACE WITH YOUR HF TOKEN\",\n",
    "    \"MAX_INPUT_LENGTH\": \"4000\", # max length of input text\n",
    "    \"MAX_TOTAL_TOKENS\": \"4096\", # max length of generated text\n",
    "    \"MESSAGES_API_ENABLED\": \"true\", # Enable the messages API\n",
    "}\n",
    "\n",
    "# create HuggingFaceModel with the image uri\n",
    "llm_model = HuggingFaceModel(\n",
    "  role=role,\n",
    "  image_uri=llm_image,\n",
    "  env=config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4afce912-1abd-416f-a17e-cd56784452d5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Deploy model to an endpoint\n",
    "# https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy\n",
    "llm_model._is_compiled_model = True # We have precompiled the model\n",
    " \n",
    "llm = llm_model.deploy(\n",
    "  initial_instance_count=1,\n",
    "  instance_type=instance_type,\n",
    "  container_startup_health_check_timeout=health_check_timeout,\n",
    "  volume_size=volume_size\n",
    ")\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7386b576-c42d-4c19-bf83-7d353e9b8702",
   "metadata": {},
   "source": [
    "After our endpoint is deployed we can run inference with it. We will use the predict method from the predictor to run inference on our endpoint. We can inference with different parameters to impact the generation. Parameters can be defined as in the `parameters` attribute of the payload. You can find supported parameters in the here.\n",
    "\n",
    "The `Messages API` allows us to interact with the model in a conversational way. We can define the role of the message and the content. The role can be either system, assistant or user. The system role is used to provide context to the model and the user role is used to ask questions or provide input to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6488180c-5b40-4a9b-b055-15de9dab4192",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "{\n",
    "  \"messages\": [\n",
    "    { \"role\": \"system\", \"content\": \"You are a helpful assistant.\" },\n",
    "    { \"role\": \"user\", \"content\": \"What is deep learning?\" }\n",
    "  ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e937e6da-0836-4f48-a765-d267aa96fc6f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Prompt to generate\n",
    "messages=[\n",
    "    { \"role\": \"system\", \"content\": \"You are a helpful assistant.\" },\n",
    "    { \"role\": \"user\", \"content\": \"What is deep learning?\" }\n",
    "  ]\n",
    " \n",
    "# Generation arguments\n",
    "parameters = {\n",
    "    \"model\": \"nithiyn/mathstral-neuron\", # placholder, needed\n",
    "    \"top_p\": 0.6,\n",
    "    \"temperature\": 0.9,\n",
    "    \"max_tokens\": 512,\n",
    "    \"stop\": [\"</s>\"],\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69ada3e5-4d54-4fb6-80e6-e8f2f7c2aece",
   "metadata": {},
   "source": [
    "---\n",
    "#### Streaming Responses with a Gradio Application\n",
    "\n",
    "[Amazon SageMaker supports streaming responses](https://aws.amazon.com/de/blogs/machine-learning/elevating-the-generative-ai-experience-introducing-streaming-support-in-amazon-sagemaker-hosting/) from your model. Using this capability, in the below section, let's build a gradio application to stream responses.\n",
    "\n",
    "Th code for the gradio application in the following steps can be found in [mistral_codestral.py](../gradio_neuron/mistral_codestral.py). The application will stream the responses from the model and display them in the UI. You can also use the application to test your model with your own inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c59ec28d-0ea4-4940-a843-8a12b9d6131a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# add apps directory to path ../apps/\n",
    "import sys\n",
    "sys.path.append(\"gradio_neuron\") \n",
    "from mathstral_chat import create_gradio_app\n",
    " \n",
    "# create gradio app\n",
    "create_gradio_app(\n",
    "    llm.endpoint_name,           # Sagemaker endpoint name\n",
    "    session=sess.boto_session,   # boto3 session used to send request \n",
    "    system_prompt=\"You are a helpful Assistant, called Mathstral. You are a meant to be a helpful assistant\",\n",
    "    share=True,                  # Share app publicly\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "290b8489-0d9c-40b2-96ed-935606e71eb6",
   "metadata": {},
   "source": [
    "---\n",
    "## Benchmarking with llmperf\n",
    "\n",
    "LLMPerf is a benchmarking tool designed to evaluate the performance of Large Language Models (LLMs) across various platforms, hardware configurations, and environments. It aims to standardize the process of measuring the efficiency, speed, and resource usage of LLMs by providing a set of tools, metrics, and frameworks that can be used to test different LLM implementations. Here, we have forked this repository and made modifications to the sagemaker client. Use the fork [nithiyn/llmperf](https://github.com/nithiyn/llmperf) for the purpose of this notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3125830-6499-416c-a99d-c7d82406ecb3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/nithiyn/llmperf.git\n",
    "!pip install -e llmperf/ --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf74d526-dc03-459b-91fb-11e525329e66",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# tell llmperf that we are using the messages api\n",
    "!MESSAGES_API=true python llmperf/token_benchmark_ray.py \\\n",
    "--model {llm.endpoint_name} \\\n",
    "--llm-api \"sagemaker\" \\\n",
    "--max-num-completed-requests 100 \\\n",
    "--timeout 600 \\\n",
    "--num-concurrent-requests 5 \\\n",
    "--results-dir \"results\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5c5c184-791e-442b-a658-9075bd871abf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#summarize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5df4abe9-0bd5-408f-a2df-7b291608fde1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import json\n",
    " \n",
    "# Reads the summary.json file and prints the results\n",
    "with open(glob.glob(f'results/*summary.json')[0], 'r') as file:\n",
    "    data = json.load(file)\n",
    "    \n",
    "print(\"Concurrent requests: 5\")\n",
    "print(f\"Avg. Input token length: {int(data['results_number_input_tokens_mean'])}\")\n",
    "print(f\"Avg. Output token length: {int(data['results_number_output_tokens_mean'])}\")\n",
    "print(f\"Avg. Time-to-first-Token: {data['results_ttft_s_mean']*1000:.2f}ms\")\n",
    "print(f\"Avg. Inter-Token-Latency: {data['results_inter_token_latency_s_mean']*1000:.2f}ms/token\")\n",
    "print(f\"Avg. Thorughput: {data['results_mean_output_throughput_token_per_s']:.2f} tokens/sec\")\n",
    "print(f\"Request per minute (RPM): {data['results_num_completed_requests_per_min']:.2f} req/min\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80c66237-af67-4210-9301-34bf8f336021",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Conclusion\n",
    "\n",
    "In this notebook, we've successfully gone over the process of compiling, deploying, and benchmarking Mathstral on Inferentia2."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a31da486-543c-43ad-a443-963edae16385",
   "metadata": {},
   "source": [
    "---\n",
    "## Distributors\n",
    "\n",
    "- AWS\n",
    "- Mistral"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}