{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74c647e6-d096-4e89-af82-a6ca26f80864",
   "metadata": {},
   "source": [
    "# AWS Inferentia2を用いたstockmark/stockmark-13bの推論\n",
    "\n",
    "このノートブックでは、`AWS Inferentia2`上で`stockmark-13b`の推論を実行する方法について説明します。\n",
    "\n",
    "## Setup\n",
    "\n",
    "`AWS Inferentia2`を搭載した`AMAZON EC2 inf2`インスタンスを起動します。AMIは`Deep Learning AMI Neuron PyTorch 1.13`の最新版を、インスタンスタイプは`inf2.8xlarge`を選びます。より高い推論パフォーマンスを得るために、`inf2.24xlarge`や`inf2.48xlarge`などの大きなインスタンスタイプを用いることもできます。\n",
    "\n",
    "参考：[Get Started with Latest Release of PyTorch Neuron (torch-neuronx)](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/neuron-setup/pytorch/neuronx/ubuntu/torch-neuronx-ubuntu20-pytorch-dlami.html#setup-torch-neuronx-ubuntu20-dlami-pytorch)\n",
    "\n",
    "インスタンスの起動後にインスタンスに接続し、まず仮想環境を作成します。\n",
    "\n",
    "`source /opt/aws_neuron_venv_pytorch/bin/activate`\n",
    "\n",
    "その後に、`transformers-neuronx`のライブラリをインストールします。\n",
    "\n",
    "`pip install transformers-neuronx --extra-index-url=https://pip.repos.neuron.amazonaws.com`\n",
    "\n",
    "参考：[Transformers NeuronX (transformers-neuronx)](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/transformers-neuronx/index.html)\n",
    "\n",
    "## Quickstart\n",
    "\n",
    "以下では、`AWS Inferentia2`上で`stockmark-13b`の推論（テキスト生成）を行うコードを紹介します。このノートブックは以下のチュートリアルをもとに作成しました。詳しくはこちらもご覧ください。\n",
    "\n",
    "参考：[Run Hugging Face meta-llama/Llama-2-13b autoregressive sampling on Inf2 & Trn1](https://github.com/aws-neuron/aws-neuron-samples/blob/master/torch-neuronx/transformers-neuronx/inference/meta-llama-2-13b-sampling.ipynb)\n",
    "\n",
    "`transformers-neuronx`を用いて推論を行うためには、事前にモデルファイルをライブラリに対応した形式に変換しておく必要があります。以下のコードは、`stockmark-13b`のモデルをHuggingface Hubからダウンロードし、`./stockmark-13b-split`というディレクトリに保存します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0475c514-6eb6-4457-aa34-62e1d42517a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers_neuronx.module import save_pretrained_split\n",
    "from transformers import LlamaForCausalLM\n",
    "\n",
    "model = LlamaForCausalLM.from_pretrained('stockmark/stockmark-13b', torch_dtype=torch.bfloat16)\n",
    "save_pretrained_split(model, './stockmark-13b-split')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a85d74d-757e-4d4c-b81a-9a5b5286207a",
   "metadata": {},
   "source": [
    "コードが終了したらJupyterLabのカーネルを再起動してください。\n",
    "\n",
    "次のコードで、`stockmark-13b`の推論を実行することができます。\n",
    "\n",
    "PyTorchでの推論と異なり、AWS Inferentia2で推論を行うために、計算グラフをコンパイルする必要があり、それに数分程度の時間がかかります。次回以降は、キャッシュされた結果が用いられます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef9a9a2d-43a7-4611-8708-e2dea60454c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "from transformers import AutoTokenizer\n",
    "from transformers_neuronx.llama.model import LlamaForSampling\n",
    "\n",
    "os.environ['NEURON_CC_FLAGS'] = '-O1'\n",
    "\n",
    "## load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained('stockmark/stockmark-13b')\n",
    "\n",
    "## load model\n",
    "model = LlamaForSampling.from_pretrained(\"./stockmark-13b-split\", batch_size=1, tp_degree=2, amp='f16')\n",
    "model.to_neuron() # compile\n",
    "\n",
    "## inference\n",
    "prompt = \"自然言語処理とは？\"\n",
    "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "\n",
    "with torch.inference_mode():\n",
    "    tokens = model.sample(inputs.input_ids, temperature = 1.0, sequence_length = 256)\n",
    "    \n",
    "generated_text = tokenizer.decode(tokens[0], skip_special_tokens=True)\n",
    "\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e19998fa-2db4-4bf2-b914-10d0885f5388",
   "metadata": {},
   "source": [
    "## Tensor parallelism\n",
    "\n",
    "`transformers-neuronx`のライブラリではテンソル並列化をサポートされています。これにより、モデルを複数のNeuronCoreに配置し、計算を並列に行うことで、推論のパフォーマンスを高めることができます。並列数（いくつのNeuronCoreを用いるか）は、モデルのロードを行う関数`LlamaForSampling.from_pretrained`の引数`tp_degree`から指定することができます。上の例では、`tp_degree=2`を指定しています。\n",
    "\n",
    "指定可能な`tp_degree`の値は、インスタンスタイプに依存します。大きなインスタンスを用いて、`tp_degree`の値を大きくすることで、latencyを小さくすることができます。\n",
    "\n",
    "| インスタンスタイプ | 指定可能なtp_degreeの最大数 |\n",
    "|:---:|:---:|\n",
    "|inf2.8xlarge|2|\n",
    "|inf2.24xlarge|12|\n",
    "|inf2.48xlarge|24|"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}