torch-neuronx/transformers-neuronx/inference/meta-llama-2-13b-sampling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "59913016-f89e-4a0e-9afe-b3a06e9112d5",
   "metadata": {},
   "source": [
    "# Run Hugging Face `meta-llama/Llama-2-13b` autoregressive sampling on Inf2 & Trn1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8454655-ec27-45e3-8da7-f82b744321ee",
   "metadata": {},
   "source": [
    "In this example we compile and deploy the Hugging Face [meta-llama/Llama-2-13b](https://huggingface.co/meta-llama/Llama-2-13b) 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 the model\n",
    "1. Perform autoregressive sampling using tensor parallelism\n",
    "\n",
    "This Jupyter Notebook can be run on an Inf2 instance (`inf2.48xlarge`) or Trn1 instance (`trn1.32xlarge`)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af2b7693-2950-41fc-a038-17cba44bf003",
   "metadata": {},
   "source": [
    "## Set up the Jupyter Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c47ef383-0dea-4423-8c38-29c73927fd78",
   "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 (`meta-llama-2-13b-sampling.ipynb`) and launch it. Follow the rest of the instructions in this tutorial. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a727963e-8178-4d2a-a5cd-a4f2bf00197e",
   "metadata": {},
   "source": [
    "## Install Dependencies\n",
    "This tutorial requires the following pip packages:\n",
    "\n",
    " - `torch-neuronx`\n",
    " - `neuronx-cc`\n",
    " - `sentencepiece`\n",
    " - `transformers`\n",
    " - `transformers-neuronx`\n",
    "\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": "3c4899b2-39b2-4309-b7df-48fe74b56eb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers-neuronx sentencepiece"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14400e26-2058-44b0-b680-b1cee57203aa",
   "metadata": {},
   "source": [
    "## Download the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e233a69-5658-4180-8f6c-91f377a01001",
   "metadata": {},
   "source": [
    "Use of the Llama 2 model is governed by the Meta license and must be downloaded and converted to the standard Hugging Face format prior to running this sample. \n",
    "\n",
    "Follow the steps described in [meta-llama/Llama-2-13b](https://huggingface.co/meta-llama/Llama-2-13b) to get access to the Llama 2 model from Meta and download the weights and tokenizer.\n",
    "\n",
    "After gaining access to the model checkpoints, you should be able to use the already converted checkpoints. Otherwise, if you are converting your own model, feel free to use the [conversion script](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py). The script can be called with the following (example) command:\n",
    "```\n",
    "python src/transformers/models/llama/convert_llama_weights_to_hf.py \\\n",
    "    --input_dir /path/to/downloaded/llama/weights --model_size 13Bf --output_dir ./Llama-2-13b --safe_serialization\n",
    "```\n",
    "\n",
    "Note: For the purposes of this sample we assume you have saved the Llama-2-13b model in a directory called `Llama-2-13b` with the following format:\n",
    "```\n",
    "Llama-2-13b/\n",
    "├── config.json\n",
    "├── generation_config.json\n",
    "├── model-00001-of-00003.safetensors\n",
    "├── model-00002-of-00003.safetensors\n",
    "├── model-00003-of-00003.safetensors\n",
    "├── model.safetensors.index.json\n",
    "├── special_tokens_map.json\n",
    "├── tokenizer.json\n",
    "├── tokenizer.model\n",
    "└── tokenizer_config.json\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e1ededb-e0d6-4c1d-aac8-bc3d29bd6ebe",
   "metadata": {},
   "source": [
    "## Perform autoregressive sampling using tensor parallelism"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a87b9f-2948-4db9-946f-b618533f03a7",
   "metadata": {},
   "source": [
    "Now we have all of the necessary files for running `meta-llama/Llama-2-13b` autoregressive sampling. \n",
    "\n",
    "The memory required to host any model can be computed with:\n",
    "```\n",
    "total memory = bytes per parameter * number of parameters\n",
    "```\n",
    "When using `float16` casted weights for a 13 billion parameter model, this works out to `2 * 13B` or ~26GB of weights. Each NeuronCore has 16GB of memory which means that a 26GB model cannot fit on a single NeuronCore. In reality, the total space required is often greater than just the number of parameters due to caching attention layer projections (KV caching). This caching mechanism grows memory allocations linearly with sequence length and batch size.\n",
    "\n",
    "To get very large language models to fit on Inf2 & Trn1, tensor parallelism is used to split weights, data, and compute across multiple NeuronCores. The number of NeuronCores that the weights are split across can be controlled by setting the `tp_degree` parameter. This parallelism degree must be chosen to ensure that the memory usage per NeuronCore will be less than the physical 16GB limit. When configuring tensor parallelism, the memory per NeuronCore can be computed with:\n",
    "\n",
    "```\n",
    "memory per core = (bytes per parameter * number of parameters) / tp_degree\n",
    "```\n",
    "\n",
    "This can be used to compute the minimum instance sizing by ensuring that the value selected for `tp_degree` results in less than 16GB allocated per NeuronCore.\n",
    "\n",
    "Note that increasing the `tp_degree` beyond the minimum requirement almost always results in a faster model. Increasing the tensor parallelism degree improves memory bandwidth which improves model performance. To optimize performance it's recommended to use the highest tensor parallelism degree that is supported by the instance. In this sample we use tensor parallelism degree 24 to optimize performance on `inf2.48xlarge`, but this should be changed to 32 if you are using a `trn1.32xlarge`. \n",
    "\n",
    "We will use the Neuron `LlamaForSampling` class to implement tensor parallelism for the Llama 2 model. The default model config supports sampling up to sequence length 2048. Tensor parallelism is enabled through the argument `tp_degree=24`. We enable `float16` casting with the `amp='f16'` flag. The model computational graph is compiled by `neuronx-cc` for optimized inference on Neuron."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc93079d-7d17-48fd-bf9d-7176bc061a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "from transformers import AutoTokenizer\n",
    "from transformers_neuronx import LlamaForSampling\n",
    "\n",
    "# load meta-llama/Llama-2-13b to the NeuronCores with 24-way tensor parallelism and run compilation\n",
    "neuron_model = LlamaForSampling.from_pretrained('Llama-2-13b', batch_size=1, tp_degree=24, amp='f16')\n",
    "neuron_model.to_neuron()\n",
    "\n",
    "# construct a tokenizer and encode prompt text\n",
    "tokenizer = AutoTokenizer.from_pretrained('Llama-2-13b')\n",
    "prompt = \"Hello, I'm a language model,\"\n",
    "input_ids = tokenizer.encode(prompt, return_tensors=\"pt\")\n",
    "\n",
    "# run inference with top-k sampling\n",
    "with torch.inference_mode():\n",
    "    start = time.time()\n",
    "    generated_sequences = neuron_model.sample(input_ids, sequence_length=2048, top_k=50)\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": "94ac4991-7606-4c2f-90af-230998b0de20",
   "metadata": {},
   "source": [
    "## Save and load the compiled model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be6a4ba9-40fd-4544-81ab-9fd249f22e4d",
   "metadata": {},
   "source": [
    "The ```save``` and ```load``` functions can be used to save and load compiled model artifacts respectively. Loading compiled model artifacts from a provided directory will avoid model recompilation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07176c29-b30b-4d16-8291-3bd0142dc42d",
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_model.save('./neuron_artifacts') # can be copied and used on a different neuron instance\n",
    "del neuron_model\n",
    "neuron_model = LlamaForSampling.from_pretrained('Llama-2-13b', batch_size=1, tp_degree=24, amp='f16')\n",
    "neuron_model.load('neuron_artifacts') # Load the compiled Neuron artifacts\n",
    "neuron_model.to_neuron() # will skip compile\n",
    "\n",
    "with torch.inference_mode():\n",
    "    start = time.time()\n",
    "    generated_sequences = neuron_model.sample(input_ids, sequence_length=2048, top_k=50)\n",
    "    elapsed = time.time() - start\n",
    "\n",
    "print(f'generated sequences {generated_sequences} in {elapsed} seconds')"
   ]
  }
 ],
 "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}