torch-neuronx/transformers-neuronx/inference/facebook-opt-66b-sampling.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "bfbcb86d",
            "metadata": {},
            "source": [
                "# Run Hugging Face `facebook/opt-66b` autoregressive sampling on Inf2"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "72bf3d85",
            "metadata": {},
            "source": [
                "In this example we compile and deploy the Hugging Face [facebook/opt-66b](https://huggingface.co/facebook/opt-66b) model for tensor parallel inference on Neuron using the `transformers-neuronx` package.\n",
                "\n",
                "The example has the following main sections:\n",
                "1. Set up the Jupyter Notebook\n",
                "1. Install dependencies\n",
                "1. Download and construct the model\n",
                "1. Split the model `state_dict` into multiple files\n",
                "1. Perform autoregressive sampling using tensor parallelism\n",
                "\n",
                "This Jupyter Notebook should be run on an Inf2 instance (`inf2.48xlarge`)."
            ]
        },
        {
            "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": "ad5ea26f",
            "metadata": {},
            "source": [
                "## Set up the Jupyter Notebook"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "458bf0a2",
            "metadata": {},
            "source": [
                "The following steps set up Jupyter Notebook and launch this tutorial:\n",
                "1. Clone the [AWS Neuron Samples](https://github.com/aws-neuron/aws-neuron-samples) repo to your instance using\n",
                "```\n",
                "git clone https://github.com/aws-neuron/aws-neuron-samples.git\n",
                "```\n",
                "2. Navigate to the `transformers-neuronx` inference samples folder\n",
                "```\n",
                "cd aws-neuron-samples/torch-neuronx/transformers-neuronx/inference\n",
                "```\n",
                "3. Follow the instructions in [Jupyter Notebook QuickStart](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/notebook/setup-jupyter-notebook-steps-troubleshooting.html) to run Jupyter Notebook on your instance.\n",
                "4. Locate this tutorial in your Jupyter Notebook session (`facebook-opt-66b-sampling.ipynb`) and launch it. Follow the rest of the instructions in this tutorial. "
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e805bd56",
            "metadata": {},
            "source": [
                "## Install Dependencies"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "0de11e14",
            "metadata": {},
            "source": [
                "This tutorial requires the following pip packages:\n",
                "\n",
                " - `torch-neuronx`\n",
                " - `neuronx-cc`\n",
                " - `transformers`\n",
                " - `transformers-neuronx`\n",
                "\n",
                "Most of these packages will be installed when configuring your environment using the [torch-neuronx inference setup guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/torch-neuronx.html#setup-torch-neuronx). The additional dependencies must be installed here:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "c69d635e",
            "metadata": {},
            "outputs": [],
            "source": [
                "!pip install transformers-neuronx transformers"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "7fea6460",
            "metadata": {},
            "source": [
                "## Download and construct the model"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "7ff897c1",
            "metadata": {},
            "source": [
                "We download and construct the `facebook/opt-66b` model using the Hugging Face\n",
                "`from_pretrained` method.\n",
                "\n",
                "Note that downloading the model checkpoint from the Hugging Face repo will take a few minutes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "e7b59499",
            "metadata": {},
            "outputs": [],
            "source": [
                "from transformers.models.opt import OPTForCausalLM\n",
                "\n",
                "hf_model = OPTForCausalLM.from_pretrained('facebook/opt-66b', low_cpu_mem_usage=True)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a54a8fb0",
            "metadata": {},
            "source": [
                "## Split the model state_dict into multiple files"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "bd5755e7",
            "metadata": {},
            "source": [
                "For the sake of reducing host memory usage, it is recommended to save the model `state_dict` as\n",
                "multiple files, as opposed to one monolithic file given by `torch.save`. This \"split-format\"\n",
                "`state_dict` can be created using the `save_pretrained_split` function. With this checkpoint format,\n",
                "the Neuron model loader can load parameters to the Neuron device high-bandwidth memory (HBM) directly\n",
                "by keeping at most one layer of model parameters in the CPU main memory.\n",
                "\n",
                "To reduce memory usage during compilation and deployment, we cast the attention and mlp to `float16` precision before saving them. We keep the layernorms in `float32`. To do this, we implement a callback function that casts each layer in the model.\n",
                "\n",
                "Note that splitting the model weights will take a few minutes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "1d116f54",
            "metadata": {},
            "outputs": [],
            "source": [
                "import torch\n",
                "from transformers_neuronx.module import save_pretrained_split\n",
                "\n",
                "def amp_callback(model, dtype):\n",
                "    # cast attention and mlp to low precision only; layernorms stay as f32\n",
                "    for block in model.model.decoder.layers:\n",
                "        block.self_attn.to(dtype)\n",
                "        block.fc1.to(dtype)\n",
                "        block.fc2.to(dtype)\n",
                "    model.lm_head.to(dtype)\n",
                "\n",
                "amp_callback(hf_model, torch.float16)\n",
                "save_pretrained_split(hf_model, './opt-66b-split')"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "cbf53ff2",
            "metadata": {},
            "source": [
                "## Perform autoregressive sampling using tensor parallelism"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "8c609ff6",
            "metadata": {},
            "source": [
                "Now we have all the necessary files for running `facebook/opt-66b` autoregressive sampling. \n",
                "\n",
                "To get a large language model working on Inf2 & Trn1, tensor parallelism is used to split weights and data across multiple NeuronCores. Each NeuronCore has 16GB of memory. As a rule of thumb, the total space required per NeuronCore will be at least `2 * number of model parameters` for a `float16` casted model. In reality, the total space required is often greater due to the key value cache, which grows with sequence lenght. This memory usage determines the minimum viable instance size since the amount of memory that will be allocated on one NeuronCore is directly proportional to the parallelism degree (`tp_degree`), or rather the number of physical NeuronCores per instance. The parallelism degree must be chosen to ensure that the memory usage per NeuronCore will be less than the physical 16GB limit. While this determines the minimum instance sizing, further decreasing the memory usage per NeuronCore by using a larger instance and a higher `tp_degree` should result in a faster model\n",
                "\n",
                "We will use the Neuron `OPTForSampling` class to implement tensor parallelism. The default model config supports sampling up to sequence length 2048, and we set batch size to 2. Tensor-parallelism is enabled through the \n",
                "`tp_degree=24` argument. Internally, the Neuron tensor manipulator will shard and duplicate tensors to multiple\n",
                "NeuronCores (24 in this case) to enable tensor-parallel computations on multiple NeuronCores. The model computational graph is compiled by `neuronx-cc` for optimized inference on Neuron.\n",
                "\n",
                "Note that compiling the model will take a long time. However, compilation is one-time process so it will not impact deployment performance. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "d98a9cb7",
            "metadata": {},
            "outputs": [],
            "source": [
                "import time\n",
                "import torch\n",
                "from transformers import AutoTokenizer\n",
                "from transformers_neuronx.opt.model import OPTForSampling\n",
                "\n",
                "# load facebook/opt-66b to NeuronCores with 24-way tensor parallel\n",
                "# enable float16 casting\n",
                "neuron_model = OPTForSampling.from_pretrained('./opt-66b-split', batch_size=2, tp_degree=24, amp='f16')\n",
                "neuron_model.to_neuron()\n",
                "\n",
                "# construct a tokenizer and encode prompt text\n",
                "tokenizer = AutoTokenizer.from_pretrained('facebook/opt-66b')\n",
                "batch_prompts = [\n",
                "    \"Hello, I'm a language model,\",\n",
                "    \"Welcome to Amazon Elastic Compute Cloud,\",\n",
                "]\n",
                "input_ids = torch.as_tensor([tokenizer.encode(text) for text in batch_prompts])\n",
                "\n",
                "with torch.inference_mode():\n",
                "    start = time.time()\n",
                "    generated_sequences = neuron_model.sample(input_ids, sequence_length=2048)\n",
                "    elapsed = time.time() - start\n",
                "\n",
                "generated_sequences = [tokenizer.decode(seq) for seq in generated_sequences]\n",
                "print(f'generated sequences {generated_sequences} in {elapsed} seconds')"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4666d038",
            "metadata": {},
            "source": [
                "The `facebook/opt-66b` model is not supported on `trn1.32xlarge`. The `facebook/opt-66b` model number of attention heads is 72, so the largest tensor parallel degree that is supported on Trn1 and is divisor of 72 is 8. However, the total size of the `facebook/opt-66b` model and its key-value caches is greater than 8 * 16 GB in `float16`, so it will not fit into memory using this tensor parallel degree. Thus, `inf2.48xlarge` with tensor parallel degree = 24 is the supported configuration for the `facebook/opt-66b` model."
            ]
        }
    ],
    "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
}