examples/notebooks/gemma2b-python-onnx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Gemma-2B with MAX Engine 🏎️ on CPU\n",
    "\n",
    "**In this notebook we will walk through an example of using [Gemma-2B](https://huggingface.co/google/gemma-2b) 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 Gemma-2B Large Language Model (LLM) is a pretrained generative text model with 2 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",
    "\n",
    "### Requirements\n",
    "\n",
    "1. Please make sure you have installed `max` as described in our [getting started](https://docs.modular.com/max/get-started/) guide\n",
    "2. Valid [huggingface access token](https://huggingface.co/docs/hub/security-tokens)\n",
    "3. Gemma-2b is a gated model. Make sure you have accepted the [terms of use](https://huggingface.co/google/gemma-2b) to be able to continue\n",
    "4. **The model size is 13G**. 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",
    "5. The runtime memory requirement is around **24G** at minimum (skipping the optional part 👇) and **50G** to run all the cells\n",
    "\n",
    "With all that, we install the required packages as follows"
   ]
  },
  {
   "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 huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Huggingface hub login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "from huggingface_hub import login\n",
    "\n",
    "access_token = getpass.getpass('Enter your Hugging Face Access Token')\n",
    "login(token=access_token)"
   ]
  },
  {
   "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",
    "hf_path = \"google/gemma-2b\"\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)\n",
    "\n",
    "# Run generation\n",
    "out_ids = hfmodel.generate(**input_ids, max_new_tokens=15, do_sample=False)\n",
    "\n",
    "# De-tokenize the generated response\n",
    "response = hftokenizer.batch_decode(out_ids, 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.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, 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.\n",
    "At the moment of writing this notebook, optimum gemma support is unreleased and but is included in the upstream [here](https://github.com/huggingface/optimum/pull/1714)\n",
    "which we can get as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install -q git+https://github.com/huggingface/optimum.git@e0cbf7dea2531603ecd0f36dfce0657baa4310ec"
   ]
  },
  {
   "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 \"google/gemma-2b\" \"./onnx/gemma-2b-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 50G 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",
    "onnxmodel = onnx.load(\"./onnx/gemma-2b-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(String(dim.dim_value))\n",
    "        elif dim.HasField(\"dim_param\"):\n",
    "            dims.append(String(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, 1, past_sequence_length, 256 ]\n",
    "past_key_values.0.value: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.1.key: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.1.value: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.2.key: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.2.value: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.3.key: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.3.value: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.4.key: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.4.value: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.5.key: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "past_key_values.5.value: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "\n",
    "...\n",
    "\n",
    "=== Outputs ===\n",
    "logits: FLOAT [ batch_size, sequence_length, 256000 ]\n",
    "present.0.key: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.0.value: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.1.key: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.1.value: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.2.key: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.2.value: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.3.key: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.3.value: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.4.key: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.4.value: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.5.key: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\n",
    "present.5.value: FLOAT [ batch_size, 1, past_sequence_length + 1, 256 ]\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",
    "\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",
    "**Note** that since ONNX doesn't support dictionaries as a input/output type, the `key_value` is expanded into 18 pairs of individual tensors (18 is the number of attention layers). The number of KV heads is 1 and the size of each key/value vector is 256.\n",
    "\n",
    "```\n",
    "past_key_values.0.key: FLOAT [ batch_size, 1, past_sequence_length, 256 ]\n",
    "```\n",
    "\n",
    "\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 lays 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/gemma-2b-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 = 18\n",
    "N_HEADS = 1\n",
    "KV_LEN = 256\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.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[\"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": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}