notebooks/deploy-mixtral.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deploy Mixtral 8x7B on Amazon SageMaker\n",
    "\n",
    "[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) is the open LLM from [Mistral AI](https://huggingface.co/mistralai). The Mixtral-8x7B is a Sparse Mixture of Experts it has a similar architecture to Mistral 7B, but comes with a twist: it’s actually 8 “expert” models in one. If you want to learn more about MoEs check out [Mixture of Experts Explained](https://huggingface.co/blog/moe). \n",
    "\n",
    "In this blog you will learn how to deploy [mistralai/Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) model to Amazon SageMaker. We are going to use the Hugging Face LLM DLC is a new purpose-built Inference Container to easily deploy LLMs in a secure and managed environment. The DLC is powered by [Text Generation Inference (TGI)](https://github.com/huggingface/text-generation-inference) a scalelable, optimized solution for deploying and serving Large Language Models (LLMs). The Blog post also includes Hardware requirements for the different model sizes. \n",
    "\n",
    "In the blog will cover how to:\n",
    "1. [Setup development environment](#1-setup-development-environment)\n",
    "2. [Retrieve the new Hugging Face LLM DLC](#2-retrieve-the-new-hugging-face-llm-dlc)\n",
    "3. [Hardware requirements](#3-hardware-requirements)\n",
    "4. [Deploy Mixtral 8x7B to Amazon SageMaker](#4-deploy-mixtral-8x7b-to-amazon-sagemaker)\n",
    "5. [Run inference and chat with the model](#5-run-inference-and-chat-with-the-model)\n",
    "6. [Clean up](#5-clean-up)\n",
    "\n",
    "Lets get started!\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup development environment\n",
    "\n",
    "We are going to use the `sagemaker` python SDK to deploy Mixtral to Amazon SageMaker. We need to make sure to have an AWS account configured and the `sagemaker` python SDK installed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: update when container is added to sagemaker sdk\n",
    "!pip install \"sagemaker>=2.199.0\" --upgrade --quiet"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are going to use Sagemaker in a local environment. You need access to an IAM Role with the required permissions for Sagemaker. You can find [here](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) more about it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "sess = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it not exists\n",
    "sagemaker_session_bucket=None\n",
    "if sagemaker_session_bucket is None and sess is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sess.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker session region: {sess.boto_region_name}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Retrieve the new Hugging Face LLM DLC\n",
    "\n",
    "Compared to deploying regular Hugging Face models we first need to retrieve the container uri and provide it to our `HuggingFaceModel` model class with a `image_uri` pointing to the image. To retrieve the new Hugging Face LLM DLC in Amazon SageMaker, we can use the `get_huggingface_llm_image_uri` method provided by the `sagemaker` SDK. This method allows us to retrieve the URI for the desired Hugging Face LLM DLC based on the specified `backend`, `session`, `region`, and `version`. You can find the available versions [here](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-text-generation-inference-containers)\n",
    "\n",
    "_Note: At the time of writing this blog post the latest version of the Hugging Face LLM DLC is not yet available via the `get_huggingface_llm_image_uri` method. We are going to use the raw container uri instead._\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMMENT IN WHEN PR (https://github.com/aws/sagemaker-python-sdk/pull/4314) IS MERGED\n",
    "# from sagemaker.huggingface import get_huggingface_llm_image_uri\n",
    "\n",
    "# # retrieve the llm image uri\n",
    "# llm_image = get_huggingface_llm_image_uri(\n",
    "#   \"huggingface\",\n",
    "#   version=\"1.3.1\"\n",
    "# )\n",
    "\n",
    "region_mapping = {\n",
    "    \"af-south-1\": \"626614931356\",\n",
    "    \"il-central-1\": \"780543022126\",\n",
    "    \"ap-east-1\": \"871362719292\",\n",
    "    \"ap-northeast-1\": \"763104351884\",\n",
    "    \"ap-northeast-2\": \"763104351884\",\n",
    "    \"ap-northeast-3\": \"364406365360\",\n",
    "    \"ap-south-1\": \"763104351884\",\n",
    "    \"ap-south-2\": \"772153158452\",\n",
    "    \"ap-southeast-1\": \"763104351884\",\n",
    "    \"ap-southeast-2\": \"763104351884\",\n",
    "    \"ap-southeast-3\": \"907027046896\",\n",
    "    \"ap-southeast-4\": \"457447274322\",\n",
    "    \"ca-central-1\": \"763104351884\",\n",
    "    \"cn-north-1\": \"727897471807\",\n",
    "    \"cn-northwest-1\": \"727897471807\",\n",
    "    \"eu-central-1\": \"763104351884\",\n",
    "    \"eu-central-2\": \"380420809688\",\n",
    "    \"eu-north-1\": \"763104351884\",\n",
    "    \"eu-west-1\": \"763104351884\",\n",
    "    \"eu-west-2\": \"763104351884\",\n",
    "    \"eu-west-3\": \"763104351884\",\n",
    "    \"eu-south-1\": \"692866216735\",\n",
    "    \"eu-south-2\": \"503227376785\",\n",
    "    \"me-south-1\": \"217643126080\",\n",
    "    \"me-central-1\": \"914824155844\",\n",
    "    \"sa-east-1\": \"763104351884\",\n",
    "    \"us-east-1\": \"763104351884\",\n",
    "    \"us-east-2\": \"763104351884\",\n",
    "    \"us-gov-east-1\": \"446045086412\",\n",
    "    \"us-gov-west-1\": \"442386744353\",\n",
    "    \"us-iso-east-1\": \"886529160074\",\n",
    "    \"us-isob-east-1\": \"094389454867\",\n",
    "    \"us-west-1\": \"763104351884\",\n",
    "    \"us-west-2\": \"763104351884\",\n",
    "}\n",
    "\n",
    "llm_image = f\"{region_mapping[sess.boto_region_name]}.dkr.ecr.{sess.boto_region_name}.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.3.1-gpu-py310-cu121-ubuntu20.04-v1.0\"\n",
    "\n",
    "# print ecr image uri\n",
    "print(f\"llm image uri: {llm_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Hardware requirements\n",
    "\n",
    "[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) is 45B parameter Mixture-of-Experts (MoE) model. This means we’ll need to have enough VRAM to hold a dense 47B parameter model. Why 47B parameters and not 8 x 7B = 56B? That’s because in MoE models, only the FFN layers are treated as individual experts. This allows Mixtral 8x7B to achieve the latency of ~12B models by being 45B parameters. \n",
    "\n",
    "> Note: The Hugging Face LLM DLC version 1.3.1 does not yet support quantization techniques with AWQ and GPTQ. We expect to support these techniques in the next releases. We will update this blog post once the new version is available. \n",
    "\n",
    "| Model                                                                       | Instance Type       | Quantization | NUM_GPUS | \n",
    "|-----------------------------------------------------------------------------|---------------------|--------------|----------|\n",
    "| [Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) | `(ml.)p4d.24xlarge` | `-`          | 8        | \n",
    "| [Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) | `(ml.)g5.48xlarge`  | `-`          | 8        | \n",
    "\n",
    "\n",
    "The parameter tested for `MAX_INPUT_LENGTH`, `MAX_BATCH_PREFILL_TOKENS`, `MAX_TOTAL_TOKENS` and `MAX_BATCH_TOTAL_TOKENS` are the same for both instances. You can tweak them for your needs. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Deploy Mixtral 8x7B to Amazon SageMaker\n",
    "\n",
    "To deploy [Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) to Amazon SageMaker we create a `HuggingFaceModel` model class and define our endpoint configuration including the `hf_model_id`, `instance_type` etc. We will use a `g5.45xlarge` instance type, which has 8 NVIDIA A10G GPUs and 192GB of GPU memory. You need atleast > 100GB of GPU memory to run Mixtral 8x7B in float16 with decent input length. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from sagemaker.huggingface import HuggingFaceModel\n",
    "\n",
    "# sagemaker config\n",
    "instance_type = \"ml.g5.48xlarge\"\n",
    "number_of_gpu = 8\n",
    "health_check_timeout = 300\n",
    "\n",
    "# Define Model and Endpoint configuration parameter\n",
    "config = {\n",
    "  'HF_MODEL_ID': \"mistralai/Mixtral-8x7B-Instruct-v0.1\", # model_id from hf.co/models\n",
    "  'SM_NUM_GPUS': json.dumps(number_of_gpu), # Number of GPU used per replica\n",
    "  'MAX_INPUT_LENGTH': json.dumps(24000),  # Max length of input text\n",
    "  'MAX_BATCH_PREFILL_TOKENS': json.dumps(32000),  # Number of tokens for the prefill operation.\n",
    "  'MAX_TOTAL_TOKENS': json.dumps(32000),  # Max length of the generation (including input text)\n",
    "  'MAX_BATCH_TOTAL_TOKENS': json.dumps(512000),  # Limits the number of tokens that can be processed in parallel during the generation\n",
    "  # ,'HF_MODEL_QUANTIZE': \"awq\", # comment in to quantize not supported yet\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",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we have created the `HuggingFaceModel` we can deploy it to Amazon SageMaker using the `deploy` method. We will deploy the model with the `ml.g5.48xlarge` instance type. TGI will automatically distribute and shard the model across all GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deploy model to an endpoint\n",
    "# https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy\n",
    "llm = llm_model.deploy(\n",
    "  initial_instance_count=1,\n",
    "  instance_type=instance_type,\n",
    "  container_startup_health_check_timeout=health_check_timeout, # 10 minutes to be able to load the model\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SageMaker will now create our endpoint and deploy the model to it. This can takes a 10-15 minutes. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Run inference and chat with the model\n",
    "\n",
    "After our endpoint is deployed we can run inference on 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](https://www.philschmid.de/sagemaker-llama-llm#5-run-inference-and-chat-with-the-model) or in the open api specification of the TGI in the [swagger documentation](https://huggingface.github.io/text-generation-inference/)\n",
    "\n",
    "The `mistralai/Mixtral-8x7B-Instruct-v0.1` is a conversational chat model meaning we can chat with it using the following prompt:\n",
    "  \n",
    "```\n",
    "<s> [INST] User Instruction 1 [/INST] Model answer 1</s> [INST] User instruction 2 [/INST]\n",
    "```\n",
    "Lets see, if Mixtral can come up with some cool ideas for the summer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt to generate\n",
    "prompt=f\"<s> [INST] What are cool ideas for the summer to do? List 5. [/INST] \"\n",
    "\n",
    "# Generation arguments\n",
    "payload = {\n",
    "    \"do_sample\": True,\n",
    "    \"top_p\": 0.6,\n",
    "    \"temperature\": 0.9,\n",
    "    \"top_k\": 50,\n",
    "    \"max_new_tokens\": 1024,\n",
    "    \"repetition_penalty\": 1.03,\n",
    "    \"return_full_text\": False,\n",
    "    \"stop\": [\"</s>\"]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay lets test it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = llm.predict({\"inputs\":prompt, \"parameters\":payload})\n",
    "\n",
    "print(chat[0][\"generated_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome, we tested infernece now lets build a cool demo which support streaming responses. [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. We can use this to stream responses, we can leverage this to create a streaming gradio application with a better user experience.\n",
    "\n",
    "We created a sample application that you can use to test your model. You can find the code in [gradio-app.py](../demo/sagemaker_chat.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,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add apps directory to path ../apps/\n",
    "import sys\n",
    "sys.path.append(\"../demo\") \n",
    "from sagemaker_chat import create_gradio_app\n",
    "\n",
    "# Generation arguments\n",
    "parameters = {\n",
    "    \"do_sample\": True,\n",
    "    \"top_p\": 0.6,\n",
    "    \"temperature\": 0.9,\n",
    "    \"top_k\": 50,\n",
    "    \"max_new_tokens\": 1024,\n",
    "    \"repetition_penalty\": 1.03,\n",
    "    \"return_full_text\": False,\n",
    "    \"stop\": [\"</s>\"]\n",
    "}\n",
    "\n",
    "# define format function for our input\n",
    "def format_prompt(message, history, system_prompt):\n",
    "    prompt = \"\"\n",
    "    for user_prompt, bot_response in history:\n",
    "        prompt = f\"<s> [INST] {user_prompt} [/INST] {bot_response}</s>\"\n",
    "        prompt += f\"### Instruction\\n{user_prompt}\\n\\n\"\n",
    "        prompt += f\"### Answer\\n{bot_response}\\n\\n\"  # Response already contains \"Falcon: \"\n",
    "    # add new user prompt if history is not empty\n",
    "    if len(history) > 0:\n",
    "        prompt += f\" [INST] {message} [/INST] \"    \n",
    "    else:\n",
    "        prompt += f\"<s> [INST] {message} [/INST] \"\n",
    "    return prompt\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",
    "    parameters=parameters,       # Request parameters\n",
    "    system_prompt=None,          # System prompt to use -> Mistral does not support system prompts\n",
    "    format_prompt=format_prompt, # Function to format prompt\n",
    "    concurrency_count=4,         # Number of concurrent requests\n",
    "    share=True,                  # Share app publicly\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![gradio](../assets/gradio.png)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Clean up\n",
    "\n",
    "To clean up, we can delete the model and endpoint.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm.delete_model()\n",
    "llm.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hf",
   "language": "python",
   "name": "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5fcf248a74081676ead7e77f54b2c239ba2921b952f7cbcdbbe5427323165924"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}