examples/notebooks/mistral7b-python-onnx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Mistral-7B with MAX Engine 🏎️ on CPU\n",
    "\n",
    "**In this notebook we will walk through an example of using [Mistral-7B](https://huggingface.co/mistralai/Mistral-7B-v0.1) model with MAX Engine 🏎️ on CPU and float32. Check out the [roadmap](https://docs.modular.com/max/roadmap) for quantization and the GPU support.**\n",
    "\n",
    "The Mistral-7B-v0.1 Large Language Model (LLM) is a pretrained generative text model with 7 billion parameters. Generative text models generate the next token iteratively given a sequence of past tokens representing the input prompt plus already generated response tokens.\n",
    "\n",
    "Thus the underlying transformer model is invoked in each iteration of this loop until we reach the stopping condition (either the maximum number of generated tokens or a token designated as the end).\n",
    "\n",
    "**Caveats:**\n",
    "\n",
    "* **The model size is 28G**. Please make sure you have enough disk space to download the model and for the converted ONNX counterpart as we will use them later in this tutorial.\n",
    "* **The runtime memory requirment is around 65G** to be able finish the notebook and run all the cells.\n",
    "\n",
    "First, make sure you have installed `max` as described in the [getting started](https://docs.staging.modular.com/engine/get-started/) as well as `PyTorch` and `transformers` packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install -q torch --index-url https://download.pytorch.org/whl/cpu\n",
    "!python3 -m pip install -q transformers onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vanilla transformers\n",
    "\n",
    "Let's first see how the model generates a response using the vanilla `transformers`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "PROMPT = \"Why did the chicken cross the road?\"\n",
    "\n",
    "hf_path = \"mistralai/Mistral-7B-v0.1\"\n",
    "hfmodel = AutoModelForCausalLM.from_pretrained(hf_path)\n",
    "hftokenizer = AutoTokenizer.from_pretrained(hf_path)\n",
    "hftokenizer.pad_token = hftokenizer.eos_token\n",
    "\n",
    "# Tokenize the text prompt\n",
    "input_ids = hftokenizer(PROMPT, return_tensors=\"pt\", max_length=128, truncation=True).input_ids\n",
    "\n",
    "# Run generation\n",
    "out_ids = hfmodel.generate(input_ids=input_ids, max_new_tokens=15, do_sample=False)\n",
    "\n",
    "# De-tokenize the generated response\n",
    "response = hftokenizer.batch_decode(out_ids.numpy(), skip_special_tokens=True)[0][len(PROMPT):]\n",
    "print(\"Response:\", response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next token generation\n",
    "\n",
    "Now that we see that the model works, let's try to decompose its `hfmodel.generate` method because we will encounter it later. We should be able to get the same output as before, but by only using `forward` method of the model.\n",
    "\n",
    "The code below is a simplified version of the actual loop you can find in the transformer's source code. It starts by initializing the current sequense to the given prompt and then generates 10 subsequent tokens - these tokens constitute the response of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "from transformers.generation.logits_process import LogitsProcessorList\n",
    "\n",
    "logits_processor = LogitsProcessorList()\n",
    "\n",
    "time_start = time()\n",
    "current_seq = input_ids\n",
    "N_TOKENS = 10\n",
    "for idx in range(N_TOKENS):\n",
    "    # Run model's `forward` on the current sequence.\n",
    "    # 'logits' output would let us determine the next token for this sequence\n",
    "    outputs = hfmodel(current_seq, return_dict=True).logits\n",
    "\n",
    "    # Get the newly generated next token\n",
    "    next_token_logits = outputs[:, -1, :]\n",
    "    next_tokens_scores = logits_processor(current_seq, next_token_logits)\n",
    "    next_tokens = torch.argmax(next_tokens_scores, dim=-1)\n",
    "\n",
    "    print(hftokenizer.decode(next_tokens), end=' ', flush=True)\n",
    "\n",
    "    # Append the new token to our sequence\n",
    "    current_seq = torch.cat([current_seq, next_tokens[:, None]], dim=-1)\n",
    "\n",
    "time_finish = time()\n",
    "print(f\"Prompt: {PROMPT}\")\n",
    "print(\"Response:\", hftokenizer.batch_decode(current_seq.numpy(), skip_special_tokens=True)[0][len(PROMPT):])\n",
    "print(f\"Tokens per second: {N_TOKENS / (time_finish - time_start):.2f}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to ONNX with optimum\n",
    "\n",
    "Great! We were able to see the same response now with using only `forward` method of our model. We're now ready to use MAX Engine 🏎️ inference.\n",
    "To do that, we start by getting an ONNX version of the model. The easiest way to do it is to use HuggingFace `optimum` tool which you can install as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install -q optimum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then the conversion to ONNX. This part can take a while. Also please make sure you've enough disk space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!optimum-cli export onnx --model \"mistralai/Mistral-7B-v0.1\" \"./onnx/mistral-7b-onnx\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional: Examine the ONNX model\n",
    "\n",
    "**Caveat: 📢 if you want to run this 👇 part you will need around 95G of memory to be able to continue finishing the notebook. Otherwise, check out the included results below.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import onnx\n",
    "\n",
    "# It will take a while to load in ONNX as well as occupying 28G more memory!\n",
    "onnxmodel = onnx.load(\"./onnx/mistral-7b-onnx/model.onnx\")\n",
    "\n",
    "def print_dims(tensor):\n",
    "    dims = []\n",
    "    for dim in tensor.type.tensor_type.shape.dim:\n",
    "        if dim.HasField(\"dim_value\"):\n",
    "            dims.append(str(dim.dim_value))\n",
    "        elif dim.HasField(\"dim_param\"):\n",
    "            dims.append(str(dim.dim_param))\n",
    "    print(onnx.TensorProto.DataType.Name(tensor.type.tensor_type.elem_type), end=\" \")\n",
    "    print(\"[\", \", \".join(dims), \"]\")\n",
    "\n",
    "print(\"=== Inputs ===\")\n",
    "for input_tensor in onnxmodel.graph.input:\n",
    "    print(input_tensor.name, end=\": \")\n",
    "    print_dims(input_tensor)\n",
    "\n",
    "print(\"\\n=== Outputs ===\")\n",
    "for output_tensor in onnxmodel.graph.output:\n",
    "    print(output_tensor.name, end=\": \")\n",
    "    print_dims(output_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "=== Inputs ===\n",
    "input_ids: INT64 [ batch_size, sequence_length ]\n",
    "attention_mask: INT64 [ batch_size, past_sequence_length + 1 ]\n",
    "position_ids: INT64 [ batch_size, sequence_length ]\n",
    "past_key_values.0.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.0.value: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.1.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.1.value: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.2.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.2.value: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.3.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.3.value: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.4.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.4.value: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.5.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "past_key_values.5.value: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "\n",
    "...\n",
    "\n",
    "=== Outputs ===\n",
    "logits: FLOAT [ batch_size, sequence_length, 32000 ]\n",
    "present.0.key: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.0.value: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.1.key: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.1.value: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.2.key: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.2.value: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.3.key: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.3.value: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.4.key: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.4.value: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.5.key: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "present.5.value: FLOAT [ batch_size, 8, past_sequence_length + 1, 128 ]\n",
    "\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It might be quite surprising to see so many inputs and outputs in the model!\n",
    "\n",
    "To decode what they all mean and how they should be used we will need to look into the [documentation of the MistralModel](https://huggingface.co/docs/transformers/v4.38.2/en/model_doc/mistral#transformers.MistralModel).\n",
    "\n",
    "In short, we have the following inputs:\n",
    "* input_ids\n",
    "* position_ids\n",
    "* attention_mask\n",
    "* past_key_values\n",
    "\n",
    "And the outputs will be:\n",
    "* logits\n",
    "* present_key_value\n",
    "\n",
    "The `past_key_values.n.key` and `past_key_values.n.value` represents cached states from previous tokens in the sequence to avoid redundant recalculations, thus speeding up the generation process.\n",
    "\n",
    "**Note** that since ONNX doesn't support dictionaries as a input/output type, the `key_value` is expanded into 32 pairs of individual tensors (32 is the number of attention layers). The number of KV heads is 8 which are part of its [Grouped Query Attention (GQA)](https://arxiv.org/abs/2305.13245) mechanism. The GQA mechanism allows the model to achieve faster inference times by optimizing the attention mechanism.\n",
    "\n",
    "```\n",
    "past_key_values.0.key: FLOAT [ batch_size, 8, past_sequence_length, 128 ]\n",
    "```\n",
    "\n",
    "Moreover, the size of each key/value vector is 128.\n",
    "\n",
    "In order to use this model we will need to slightly modify our glue code to correctly weave all these values from each iteration to the next.\n",
    "Specifically, we will need to pass the `key_values` from previous iteration to the current (for the first iteration they are initializes as empty tensors).\n",
    "We will also need to correctly fill in `position_ids` and `attention_mask` tensors and update them on each iteration. We will not get into all the details of how exactly all these tensors affect the model behavior and should be used - this is an extremely interesting topic, but it is beyond the scope of this walkthrough."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MAX Engine 🏎️\n",
    "\n",
    "With that we're finally ready to use the MAX Engine 🏎️ for inference.\n",
    "\n",
    "The code modifications follows our previous approach and is quite minimal. All we need to do is to load the ONNX model (whch can take a while) into an `InferenceSession` object and instead of using the `hfmodel` we will need to use `maxmodel.execute`, and pack the input values into a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "from max import engine\n",
    "# Create an InferenceSession and load the ONNX model\n",
    "session = engine.InferenceSession()\n",
    "maxmodel = session.load(\"./onnx/mistral-7b-onnx/model.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also quickly inspect the input and output metadata that match the ONNX version above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for tensor in maxmodel.input_metadata:\n",
    "    print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')\n",
    "\n",
    "for tensor in maxmodel.output_metadata:\n",
    "    print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the how to get the response from our `maxmodel`. The token per second can be up to **2X** faster comparing to the PyTorch version. Note that this tutorial doesn't provide an accurate benchmark. For more, please check out our [performance dashboard](https://performance.modular.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {}\n",
    "N_BATCH = 1\n",
    "N_LAYERS = 32\n",
    "N_HEADS = 8\n",
    "KV_LEN = 128\n",
    "# Initialize the additional layer to 0 for the first iteration:\n",
    "for i in range(N_LAYERS):\n",
    "    inputs[f\"past_key_values.{i}.key\"] = torch.zeros([N_BATCH, N_HEADS, 0, KV_LEN], dtype=torch.float).numpy()\n",
    "    inputs[f\"past_key_values.{i}.value\"] = torch.zeros([N_BATCH, N_HEADS, 0, KV_LEN], dtype=torch.float).numpy()\n",
    "\n",
    "current_seq = input_ids\n",
    "\n",
    "time_start = time()\n",
    "for idx in range(N_TOKENS):\n",
    "    # Prepare inputs dictionary\n",
    "    inputs[\"input_ids\"] = current_seq.numpy()\n",
    "    inputs[\"position_ids\"] = torch.arange(inputs[\"input_ids\"].shape[1], dtype=torch.long).unsqueeze(0).numpy()\n",
    "    inputs[\"attention_mask\"] = torch.ones([1, inputs[f\"past_key_values.0.key\"].shape[2] + inputs[\"input_ids\"].shape[1]], dtype=torch.int64).numpy()\n",
    "\n",
    "    # Run the model with MAX engine\n",
    "    max_outputs = maxmodel.execute(**inputs)\n",
    "    outputs = torch.from_numpy(max_outputs[\"logits\"])\n",
    "\n",
    "    # Get the newly generated next token\n",
    "    next_token_logits = outputs[:, -1, :]\n",
    "    next_tokens_scores = logits_processor(current_seq, next_token_logits)\n",
    "    next_tokens = torch.argmax(next_tokens_scores, dim=-1)\n",
    "\n",
    "    print(hftokenizer.decode(next_tokens), end=' ', flush=True)\n",
    "\n",
    "    # Append the new token to our sequence\n",
    "    current_seq = torch.cat([current_seq, next_tokens[:, None]], dim=-1)\n",
    "\n",
    "    # Update the KV cache for the next iteration\n",
    "    for i in range(N_LAYERS):\n",
    "        inputs[f\"past_key_values.{i}.key\"] = max_outputs[f\"present.{i}.key\"]\n",
    "        inputs[f\"past_key_values.{i}.value\"] = max_outputs[f\"present.{i}.value\"]\n",
    "\n",
    "time_finish = time()\n",
    "\n",
    "print(f\"Prompt: {PROMPT}\")\n",
    "print(\"Response:\", hftokenizer.batch_decode(current_seq.numpy(), skip_special_tokens=True)[0][len(PROMPT):])\n",
    "print(f\"Tokens per second: {idx/(time_finish-time_start):.2f}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That is it! 🎉\n",
    "\n",
    "Serving an LLM has historically been not an easy task, but hopefully this example lifts the curtain on how this can be done. MAX Engine 🏎️ doesn't (yet) make this process easier, however, if you've already gone this path with ONNX or TorchScript, switching to MAX should be trivial and bring easy performance wins."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}