torch-neuronx/inference/hf_pretrained_gpt2_feature_extraction_on_trn1.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "9963bc02",
            "metadata": {},
            "source": [
                "# HuggingFace Pretrained GPT2 Feature Extraction on Trn1"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "f5718136",
            "metadata": {},
            "source": [
                "## Introduction\n",
                "\n",
                "This notebook demonstrates how to compile and run a HuggingFace 🤗 Transformers GPT2 model for accelerated feature extraction on Neuron. This notebook will use the [`gpt2`](https://huggingface.co/gpt2) model, which is primarily used for text generation and feature extraction. \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 \"transformers < 4.21.0\""
            ]
        },
        {
            "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. We also use define a basic wrapper that ensures the `input_ids` and `attention_mask` kwargs are passed into the GPT2 model in the correct positions without requiring a dictionary. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "8e8853d4",
            "metadata": {},
            "outputs": [],
            "source": [
                "import torch\n",
                "import torch_neuronx\n",
                "from transformers import GPT2Tokenizer, GPT2Model\n",
                "\n",
                "\n",
                "# Create a wrapper to correctly order the inputs\n",
                "class GPT2Neuron(torch.nn.Module):\n",
                "    \"\"\"\n",
                "    Ensures that `input_ids` and `attention_mask` are passed into the GPT2\n",
                "    model in the correct positions without requiring a dictionary.\n",
                "    \"\"\"\n",
                "\n",
                "    def __init__(self, model) -> None:\n",
                "        super().__init__()\n",
                "        self.model = model\n",
                "\n",
                "    def forward(self, input_ids, attention_mask):\n",
                "        return self.model(input_ids=input_ids, attention_mask=attention_mask)\n",
                "\n",
                "\n",
                "# Create the tokenizer and model\n",
                "tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
                "tokenizer.pad_token = tokenizer.eos_token # Define the padding token value\n",
                "gpt2 = GPT2Model.from_pretrained('gpt2', torchscript=True)\n",
                "model = GPT2Neuron(gpt2)\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",
                ")\n",
                "\n",
                "# Run inference on CPU\n",
                "output_cpu = model(*example)\n",
                "\n",
                "# Compile the model using the wrapper\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 feature extraction 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 outputs:    {output_cpu[0][0][0][:10]}\")\n",
                "print(f\"Neuron outputs: {output_neuron[0][0][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
}