torch-neuronx/inference/hf_pretrained_bert_inference_on_trn1.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "9963bc02",
            "metadata": {},
            "source": [
                "# HuggingFace Pretrained BERT Inference on Trn1"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "f5718136",
            "metadata": {},
            "source": [
                "## Introduction\n",
                "\n",
                "This notebook demonstrates how to compile and run a HuggingFace 🤗 Transformers BERT model for accelerated inference on Neuron. This notebook will use the [`bert-base-uncased`](https://huggingface.co/bert-base-uncased) model, which is primarily used for masked language modeling and next sentence prediction. \n",
                "\n",
                "This Jupyter notebook should be run on a Trn1 instance (`trn1.2xlarge` or larger)."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Verify that this Jupyter notebook is running the Python kernel environment that was set up according to the [PyTorch Installation Guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/torch-neuronx.html#setup-torch-neuronx). You can select the kernel from the 'Kernel -> Change Kernel' option on the top of this Jupyter notebook page."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "07bb46a3",
            "metadata": {},
            "source": [
                "## Install Dependencies\n",
                "This tutorial requires the following pip packages:\n",
                "\n",
                "- `torch-neuronx`\n",
                "- `neuronx-cc`\n",
                "- `transformers`\n",
                "\n",
                "Most of these packages will be installed when configuring your environment using the Trn1 setup guide. The additional dependencies must be installed here:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "220cf78c",
            "metadata": {},
            "outputs": [],
            "source": [
                "%env TOKENIZERS_PARALLELISM=True #Supresses tokenizer warnings making errors easier to detect\n",
                "!pip install -U transformers"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5a7dbb7b",
            "metadata": {},
            "source": [
                "## Compile the model into an AWS Neuron optimized TorchScript"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "caf9362d",
            "metadata": {},
            "source": [
                "In the following section, we load the model and tokenizer, get s sample input, run inference on CPU, compile the model for Neuron using `torch_neuronx.trace()` and save the optimized model as `TorchScript`.\n",
                "\n",
                "`torch_neuronx.trace()` expects a tensor or tuple of tensor inputs to use for tracing, so we unpack the tokenzier output. Additionally, the input shape that's used duing compilation must match the input shape that's used during inference. To handle this, we pad the inputs to the maximum size that we will see during inference."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "8e8853d4",
            "metadata": {},
            "outputs": [],
            "source": [
                "import torch\n",
                "import torch_neuronx\n",
                "from transformers import BertTokenizer, BertModel\n",
                "\n",
                "# Create the tokenizer and model\n",
                "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
                "model = BertModel.from_pretrained(\"bert-base-uncased\")\n",
                "model.eval()\n",
                "\n",
                "# Get an example input\n",
                "text = \"Replace me by any text you'd like.\"\n",
                "\n",
                "encoded_input = tokenizer(\n",
                "    text,\n",
                "    max_length=128,\n",
                "    padding='max_length',\n",
                "    truncation=True,\n",
                "    return_tensors='pt'\n",
                ")\n",
                "\n",
                "example = (\n",
                "    encoded_input['input_ids'],\n",
                "    encoded_input['attention_mask'],\n",
                "    encoded_input['token_type_ids']\n",
                ")\n",
                "\n",
                "# Run inference on CPU\n",
                "output_cpu = model(*example)\n",
                "\n",
                "# Compile the model\n",
                "model_neuron = torch_neuronx.trace(model, example)\n",
                "\n",
                "# Save the TorchScript for inference deployment\n",
                "filename = 'model.pt'\n",
                "torch.jit.save(model_neuron, filename)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a7056884",
            "metadata": {},
            "source": [
                "## Run inference and compare results\n",
                "\n",
                "In this section we load the compiled model, run inference on Neuron, and compare the CPU and Neuron outputs."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "27b56563",
            "metadata": {},
            "outputs": [],
            "source": [
                "# Load the TorchScript compiled model\n",
                "model_neuron = torch.jit.load(filename)\n",
                "\n",
                "# Run inference using the Neuron model\n",
                "output_neuron = model_neuron(*example)\n",
                "\n",
                "# Compare the results\n",
                "print(f\"CPU last_hidden_state:    {output_cpu['last_hidden_state'][0][0][:10]}\")\n",
                "print(f\"Neuron last_hidden_state: {output_neuron['last_hidden_state'][0][0][:10]}\")\n",
                "print(f\"CPU pooler_output:        {output_cpu['pooler_output'][0][:10]}\")\n",
                "print(f\"Neuron pooler_output:     {output_neuron['pooler_output'][0][:10]}\")"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python (Neuron PyTorch)",
            "language": "python",
            "name": "pytorch_venv"
        },
        "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.8.10"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}