torch-neuronx/training/hf_bert_jp/bert-jp-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39d73d16-fce1-4340-8857-9a7815f0773f",
   "metadata": {},
   "source": [
    "# 日本語 BERT Base Model Fine-tuning & Deployment on Inferentia2\n",
    "本 Notebook の元ネタのブログはこちらから\n",
    "+ https://aws.amazon.com/jp/blogs/news/aws-trainium-amazon-ec2-trn1-ml-training-part1/\n",
    "+ https://aws.amazon.com/jp/blogs/news/aws-trainium-amazon-ec2-trn1-ml-training-part2/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0641e6b-67e8-4f10-86ef-1ce00c20162b",
   "metadata": {},
   "source": [
    "## 事前準備"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eaf706fb-f236-4b87-870f-4e7f81806dcc",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com\n",
      "Requirement already satisfied: pip in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (23.2.1)\n",
      "Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com\n",
      "Requirement already satisfied: transformers[ja]==4.27.4 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (4.27.4)\n",
      "Requirement already satisfied: datasets in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (2.14.4)\n",
      "Requirement already satisfied: filelock in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (3.12.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (0.14.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (1.21.6)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (21.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (5.4.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (2023.5.5)\n",
      "Requirement already satisfied: requests in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (2.30.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (0.12.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (4.65.0)\n",
      "Requirement already satisfied: fugashi>=1.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (1.2.1)\n",
      "Requirement already satisfied: ipadic<2.0,>=1.0.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (1.0.0)\n",
      "Requirement already satisfied: unidic-lite>=1.0.7 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (1.0.8)\n",
      "Requirement already satisfied: unidic>=1.0.2 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (1.1.0)\n",
      "Requirement already satisfied: sudachipy>=0.6.6 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (0.6.7)\n",
      "Requirement already satisfied: sudachidict-core>=20220729 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (20230110)\n",
      "Requirement already satisfied: rhoknp>=1.1.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from transformers[ja]==4.27.4) (1.3.0)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (12.0.0)\n",
      "Requirement already satisfied: dill<0.3.8,>=0.3.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (0.3.5.1)\n",
      "Requirement already satisfied: pandas in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (2.0.1)\n",
      "Requirement already satisfied: xxhash in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (3.2.0)\n",
      "Requirement already satisfied: multiprocess in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (0.70.13)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (2023.4.0)\n",
      "Requirement already satisfied: aiohttp in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from datasets) (3.8.4)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (23.1.0)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (3.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (6.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (1.9.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (1.3.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from huggingface-hub<1.0,>=0.11.0->transformers[ja]==4.27.4) (4.5.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from packaging>=20.0->transformers[ja]==4.27.4) (3.0.9)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from requests->transformers[ja]==4.27.4) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from requests->transformers[ja]==4.27.4) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from requests->transformers[ja]==4.27.4) (2022.12.7)\n",
      "Requirement already satisfied: wasabi<1.0.0,>=0.6.0 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from unidic>=1.0.2->transformers[ja]==4.27.4) (0.10.1)\n",
      "Requirement already satisfied: plac<2.0.0,>=1.1.3 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from unidic>=1.0.2->transformers[ja]==4.27.4) (1.3.5)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from pandas->datasets) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from pandas->datasets) (2023.3)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from pandas->datasets) (2023.3)\n",
      "Requirement already satisfied: six>=1.5 in /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "%env TOKENIZERS_PARALLELISM=True #Supresses tokenizer warnings making errors easier to detect\n",
    "!pip install -U pip\n",
    "!pip install -U transformers[ja]==4.27.4 datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b8aa82ee-39d2-4e6c-9294-b920a1284890",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aws-neuronx-runtime-discovery 2.9\n",
      "libneuronxla                  0.5.391\n",
      "neuron-cc                     1.15.0.0+eec0c3604\n",
      "neuronperf                    1.8.0.0+5dccae385\n",
      "neuronx-cc                    2.8.0.25+a3ad0f342\n",
      "neuronx-hwm                   2.8.0.3+2b7c6da39\n",
      "optimum-neuron                0.0.9.dev0\n",
      "torch                         1.13.1\n",
      "torch-neuron                  1.13.1.2.8.9.0\n",
      "torch-neuronx                 1.13.1.1.9.0\n",
      "torch-xla                     1.13.1+torchneuron8\n",
      "torchmetrics                  1.0.1\n",
      "torchvision                   0.14.1\n",
      "transformers-neuronx          0.5.58\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "!pip list | grep \"neuron\\|torch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b2168c77-6ac7-4f71-95c7-06e6df4386fd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!sudo rmmod neuron; sudo modprobe neuron"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87668961-d047-4603-a3b1-a05544713952",
   "metadata": {},
   "source": [
    "## データセットの準備\n",
    "本テストでは、Huggingface Hub で利用可能な以下のセンチメント（感情）データセットのうち、日本語のサブセットを使用します。\n",
    "https://huggingface.co/datasets/tyqiangz/multilingual-sentiments\n",
    "\n",
    "本テストではテキストデータをPositiveかNegativeに分類する 2 クラスの分類問題として扱うことにします。元々のデータセットは positive(LABEL_0)、neutral(LABEL_1)、negative(LABEL_2)としてラベル付けされていますが、neutralのデータは使用しないこととし、ラベルをpositive(LABEL_0)、negative(LABEL_1)として再定義します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53189385-ef49-45a3-b87a-9955c91395ec",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['text', 'source', 'label'],\n",
      "        num_rows: 120000\n",
      "    })\n",
      "    validation: Dataset({\n",
      "        features: ['text', 'source', 'label'],\n",
      "        num_rows: 3000\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['text', 'source', 'label'],\n",
      "        num_rows: 3000\n",
      "    })\n",
      "})\n",
      "{'text': Value(dtype='string', id=None), 'source': Value(dtype='string', id=None), 'label': ClassLabel(names=['positive', 'neutral', 'negative'], id=None)}\n",
      "{'text': '箱を開けただけで予想以上に匂いが甘くて 私には合わないかもです。 でも もったいないので今使ってるのが無くなったら使ってみます。', 'labels': 1}\n",
      "{'text': '色々な変換アダプが売ってますけど長さも結構ありますし重宝しています。主にタブレットから→テレビに使ってるんですけど子供達に動画を見せる時タブレットを直接見せるよりテレビで大きな画面で見せた方が良いですね。なんの設定も要らないので直接繋げればすぐ使える所も気に入っています。', 'labels': 0}\n"
     ]
    }
   ],
   "source": [
    "from transformers import BertJapaneseTokenizer\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Prepare dataset\n",
    "dataset = load_dataset(\"tyqiangz/multilingual-sentiments\", \"japanese\")\n",
    "print(dataset)\n",
    "\n",
    "print(dataset[\"train\"].features)\n",
    "\n",
    "dataset = dataset.remove_columns([\"source\"])\n",
    "dataset = dataset.filter(lambda dataset: dataset[\"label\"] != 1)\n",
    "dataset = dataset.map(lambda dataset: {\"labels\": int(dataset[\"label\"] == 2)}, remove_columns=[\"label\"])\n",
    "\n",
    "print(dataset[\"train\"][20000])\n",
    "print(dataset[\"train\"][50000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0481f80-52cf-4255-a76e-c49a948247ec",
   "metadata": {},
   "source": [
    "次に、文章テキストのままだとモデルのトレーニングはできないため、テキストを意味のある単位で分割（トークナイズ）した上で数値に変換します。トークナイザーには MeCab ベースの BertJapaneseTokenizer を利用しました。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a3f7c32e-c0c0-4a1e-92c2-4b6490f68917",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c207ad64a5f4fb682edf381402f3aa2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/4000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b360c980d1474e3e9576ee8c244c5e92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/256 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "MODEL_NAME = \"cl-tohoku/bert-base-japanese-whole-word-masking\"\n",
    "tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], padding=\"max_length\", max_length=128, truncation=True)\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "tokenized_datasets = tokenized_datasets.remove_columns([\"text\"])\n",
    "\n",
    "train_dataset = tokenized_datasets[\"train\"].shuffle().select(range(4000))\n",
    "eval_dataset = tokenized_datasets[\"test\"].shuffle().select(range(256))\n",
    "\n",
    "# Save dataset\n",
    "train_dataset.save_to_disk(\"./train/\")\n",
    "eval_dataset.save_to_disk(\"./test/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4005724-1c52-4426-b7cb-16996462d29d",
   "metadata": {},
   "source": [
    "実際にどのように変換されているのか、以下のスクリプトを実行し確認してみましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "229a9403-e9b7-407e-a6ce-71ecd69697d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'text': '色々な変換アダプが売ってますけど長さも結構ありますし重宝しています。主にタブレットから→テレビに使ってるんですけど子供達に動画を見せる時タブレットを直接見せるよりテレビで大きな画面で見せた方が良いですね。なんの設定も要らないので直接繋げればすぐ使える所も気に入っています。', 'labels': 0}\n",
      "Tokenize: ['色々', 'な', '変換', 'アダ', '##プ', 'が', '売っ', 'て', 'ます', 'けど', '長', 'さ', 'も', '結', '##構', 'あり', 'ます', 'し', '重宝', 'し', 'て', 'い', 'ます', '。', '主', 'に', 'タブレット', 'から', '→', 'テレビ', 'に', '使っ', 'てる', 'ん', 'です', 'けど', '子供', '達', 'に', '動画', 'を', '見せる', '時', 'タブレット', 'を', '直接', '見せる', 'より', 'テレビ', 'で', '大きな', '画面', 'で', '見せ', 'た', '方', 'が', '良い', 'です', 'ね', '。', 'なん', 'の', '設定', 'も', '要', '##ら', 'ない', 'ので', '直接', '繋', '##げ', '##れ', 'ば', 'すぐ', '使える', '所', 'も', '気に入っ', 'て', 'い', 'ます', '。']\n",
      "Encode: [2, 16166, 18, 4618, 6568, 28526, 14, 18544, 16, 2610, 11218, 177, 26, 28, 224, 28902, 130, 2610, 15, 26736, 15, 16, 21, 2610, 8, 137, 7, 21024, 40, 2592, 571, 7, 2110, 7134, 1058, 2992, 11218, 1803, 784, 7, 4884, 11, 7237, 72, 21024, 11, 2031, 7237, 221, 571, 12, 1200, 4180, 12, 2685, 10, 283, 14, 3614, 2992, 1852, 8, 4508, 5, 1374, 28, 598, 28469, 80, 947, 2031, 3510, 28764, 28461, 312, 2459, 10265, 233, 28, 18759, 16, 21, 2610, 8, 3]\n"
     ]
    }
   ],
   "source": [
    "index = 50000\n",
    "print(dataset[\"train\"][index])\n",
    "print('Tokenize:', tokenizer.tokenize(dataset[\"train\"]['text'][index]))\n",
    "print('Encode:', tokenizer.encode(dataset[\"train\"]['text'][index]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59b00108-8105-4190-b6ad-27effb7092d5",
   "metadata": {},
   "source": [
    "## Trainer API を使用した トレーニング（ファインチューニング）実行\n",
    "Transformers には Trainer という便利なクラスがあり、Torch Neuron からも利用可能です。 ここでは Trainer API を利用してトレーニングを実行していきます。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36508a19-4f9d-483c-b118-bbf66b2e85ef",
   "metadata": {},
   "source": [
    "### neuron_parallel_compile による事前コンパイル\n",
    "トレーニングの各ステップでは、グラフがトレースされ、トレースされたグラフが以前のものと異なる場合は、再度計算グラフのコンパイルが発生します。大規模なモデルの場合、各グラフのコンパイル時間が長くなることがあり、トレーニング時間の中で占めるコンパイル時間がボトルネックとなってしまう場合もあり得ます。このコンパイル時間を短縮するため、PyTorch Neuron では neuron_parallel_compile ユーティリティが提供されています。neuron_parallel_compile は、スクリプトの試行からグラフを抽出し並列事前コンパイルを実施、コンパイル結果（NEFF : Neuron Executable File Format）をキャッシュとしてディスク上に保持します。\n",
    "\n",
    "では実際に事前コンパイルを実行してみましょう。以下の内容でbert-jp-precompile.pyというファイル名の Python スクリプトを作成し実行します。スクリプトは基本的にこの後実行するトレーニングスクリプトと同じ内容ですが、neuron_parallel_compileはグラフの高速コンパイルのみを目的とし実際の演算は実行されず、出力結果は無効となります。トレーニング実行中も必要に応じてグラフはコンパイルされるため、この事前コンパイルのプロセスはスキップし、次のトレーニング処理に直接進んでも問題はありません。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e6acfe7-0dc2-4341-af96-f837bc627583",
   "metadata": {},
   "source": [
    "コンパイル時間を短縮するためデータセット、epoch 数を制限している点にご注意ください。コンパイル結果は `/var/tmp/neuron-compile-cache/` 以下に保存されています。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "70c8f769-91a2-4966-8b55-2e1df787d716",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting bert-jp-precompile.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile bert-jp-precompile.py\n",
    "\n",
    "from transformers import BertForSequenceClassification, Trainer, TrainingArguments\n",
    "from datasets import load_from_disk\n",
    "import torch, torch_xla.core.xla_model as xm\n",
    "import os\n",
    "\n",
    "os.environ[\"NEURON_CC_FLAGS\"] = \"--model-type=transformer\"\n",
    "\n",
    "device = xm.xla_device()\n",
    "\n",
    "MODEL_NAME = \"cl-tohoku/bert-base-japanese-whole-word-masking\"\n",
    "model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2).to(device)\n",
    "\n",
    "train_dataset = load_from_disk(\"./train/\").with_format(\"torch\")\n",
    "train_dataset = train_dataset.select(range(64))\n",
    "\n",
    "eval_dataset = load_from_disk(\"./test/\").with_format(\"torch\")\n",
    "eval_dataset = eval_dataset.select(range(64))\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    num_train_epochs = 2,\n",
    "    learning_rate = 5e-5,\n",
    "    per_device_train_batch_size = 8,\n",
    "    per_device_eval_batch_size = 8,\n",
    "    output_dir = \"./results\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_args,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = eval_dataset,\n",
    ")\n",
    "\n",
    "train_result = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6db73bea-dcee-4d21-badc-ea7844418195",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-08-14 08:09:07.000974:  715601  INFO ||PARALLEL_COMPILE||: Removing existing workdir /tmp/parallel_compile_workdir\n",
      "2023-08-14 08:09:07.000977:  715601  INFO ||PARALLEL_COMPILE||: Running trial run (add option to terminate trial run early; also ignore trial run's generated outputs, i.e. loss, checkpoints)\n",
      "Some weights of the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking were not used when initializing BertForSequenceClassification: ['cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "  0%|          | 0/16 [00:00<?, ?it/s]2023-08-14 08:09:20.000543: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:09:20.000544: INFO ||NCC_WRAPPER||: Extracting graphs (/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_15770068861141756492+abb26765/model.hlo.pb) for ahead-of-time parallel compilation. No compilation was done.\n",
      "  6%|▋         | 1/16 [00:00<00:04,  3.09it/s]2023-08-14 08:09:21.000668: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:09:21.000669: INFO ||NCC_WRAPPER||: Extracting graphs (/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_16268669036401171959+abb26765/model.hlo.pb) for ahead-of-time parallel compilation. No compilation was done.\n",
      " 12%|█▎        | 2/16 [00:02<00:17,  1.23s/it]2023-08-14 08:09:23.000591: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:09:23.000592: INFO ||NCC_WRAPPER||: Extracting graphs (/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_14409091931058726813+abb26765/model.hlo.pb) for ahead-of-time parallel compilation. No compilation was done.\n",
      " 94%|█████████▍| 15/16 [00:05<00:00,  6.83it/s]2023-08-14 08:09:26.000074: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:09:26.000075: INFO ||NCC_WRAPPER||: Extracting graphs (/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_12573384035295665077+abb26765/model.hlo.pb) for ahead-of-time parallel compilation. No compilation was done.\n",
      "100%|██████████| 16/16 [00:05<00:00,  2.73it/s]\n",
      "{'train_runtime': 6.5838, 'train_samples_per_second': 19.442, 'train_steps_per_second': 2.43, 'train_loss': 0.0, 'epoch': 2.0}\n",
      "2023-08-14 08:09:26.000814:  715601  INFO ||PARALLEL_COMPILE||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:09:30.000797:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:31.000249:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:31.000750:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:31.000807:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:31.000878:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      ".2023-08-14 08:09:32.000022:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:32.000188:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      "..2023-08-14 08:09:32.000394:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 4, locked are 0, failed are 0, done are 51, total is 55\n",
      ".\n",
      "Compiler status PASS\n",
      "2023-08-14 08:09:34.000852:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 1, locked are 3, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:35.000168:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 1, locked are 3, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:35.000169:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 1, locked are 3, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:35.000169:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 1, locked are 3, failed are 0, done are 51, total is 55\n",
      "2023-08-14 08:09:37.000837:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 3, failed are 0, done are 52, total is 55\n",
      "2023-08-14 08:09:37.000838:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 3, failed are 0, done are 52, total is 55\n",
      "2023-08-14 08:09:37.000838:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:09:37.000908:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:09:38.000287:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 3, failed are 0, done are 52, total is 55\n",
      "2023-08-14 08:09:38.000287:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:09:38.000287:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 3, failed are 0, done are 52, total is 55\n",
      "2023-08-14 08:09:38.000287:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:09:38.000550:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 3, failed are 0, done are 52, total is 55\n",
      "2023-08-14 08:09:38.000550:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "...\n",
      "Compiler status PASS\n",
      "2023-08-14 08:10:04.000407:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 2, failed are 0, done are 53, total is 55\n",
      "2023-08-14 08:10:04.000407:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "........................\n",
      "Compiler status PASS\n",
      "2023-08-14 08:13:59.000853:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 1, failed are 0, done are 54, total is 55\n",
      "2023-08-14 08:13:59.000853:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "...\n",
      "Compiler status PASS\n",
      "2023-08-14 08:15:08.000105:  715601  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 55, total is 55\n",
      "2023-08-14 08:15:08.000106:  715601  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:15:08.000106:  715601  INFO ||PARALLEL_COMPILE||: {\n",
      "    \"compilation_summary\": {\n",
      "        \"SUCCESS\": 4\n",
      "    },\n",
      "    \"compilation_report\": {\n",
      "        \"/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_15770068861141756492+abb26765/model.hlo.pb\": {\n",
      "            \"status\": \"SUCCESS\",\n",
      "            \"retry\": 0\n",
      "        },\n",
      "        \"/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_12573384035295665077+abb26765/model.hlo.pb\": {\n",
      "            \"status\": \"SUCCESS\",\n",
      "            \"retry\": 0\n",
      "        },\n",
      "        \"/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_16268669036401171959+abb26765/model.hlo.pb\": {\n",
      "            \"status\": \"SUCCESS\",\n",
      "            \"retry\": 0\n",
      "        },\n",
      "        \"/var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_14409091931058726813+abb26765/model.hlo.pb\": {\n",
      "            \"status\": \"SUCCESS\",\n",
      "            \"retry\": 0\n",
      "        }\n",
      "    }\n",
      "}\n",
      "2023-08-14 08:15:08.000106:  715601  INFO ||PARALLEL_COMPILE||: Total graphs: 4\n",
      "2023-08-14 08:15:08.000106:  715601  INFO ||PARALLEL_COMPILE||: Total successful compilations: 4\n",
      "2023-08-14 08:15:08.000106:  715601  INFO ||PARALLEL_COMPILE||: Total failed compilations: 0\n",
      "\n",
      "real\t6m0.309s\n",
      "user\t15m6.311s\n",
      "sys\t1m3.244s\n"
     ]
    }
   ],
   "source": [
    "!time XLA_USE_BF16=1 neuron_parallel_compile python3 bert-jp-precompile.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "049a9a17-1e83-4628-a6f4-29fb5dff2d04",
   "metadata": {},
   "source": [
    "### シングルワーカーでのトレーニング実行\n",
    "次に実際にトレーニングを実行してみます。事前コンパイルを実行した場合でも、追加のコンパイルが発生することがあります。一通りのコンパイルが終了した後、2度目以降の実行では、Neuron コアの恩恵を受けた高速トレーニングを体験できます。以下の内容で bert-jp-single.py というファイル名の Python スクリプトを作成し実行してみましょう。\n",
    "\n",
    "先程の事前コンパイルとは異なり、今回は実際にトレーニングを実行するため、用意したデータセット全てに対して epoch = 10 で実行します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "78747ec6-3e55-4b1a-b57f-a8f656b069a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting bert-jp-single.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile bert-jp-single.py\n",
    "\n",
    "from transformers import BertForSequenceClassification, BertJapaneseTokenizer, Trainer, TrainingArguments\n",
    "from datasets import load_from_disk\n",
    "import torch, torch_xla.core.xla_model as xm\n",
    "import os\n",
    "\n",
    "os.environ[\"NEURON_CC_FLAGS\"] = \"--model-type=transformer\"\n",
    "\n",
    "device = xm.xla_device()\n",
    "\n",
    "MODEL_NAME = \"cl-tohoku/bert-base-japanese-whole-word-masking\"\n",
    "tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)\n",
    "model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2).to(device)\n",
    "\n",
    "train_dataset = load_from_disk(\"./train/\").with_format(\"torch\")\n",
    "eval_dataset = load_from_disk(\"./test/\").with_format(\"torch\")\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    num_train_epochs = 10,\n",
    "    learning_rate = 5e-5,\n",
    "    per_device_train_batch_size = 8,\n",
    "    per_device_eval_batch_size = 8,\n",
    "    output_dir = \"./results\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_args,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = eval_dataset,\n",
    "    tokenizer = tokenizer,\n",
    ")\n",
    "\n",
    "train_result = trainer.train()\n",
    "print(train_result)\n",
    "\n",
    "eval_result = trainer.evaluate()\n",
    "print(eval_result)\n",
    "\n",
    "trainer.save_model(\"./results\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "13e1319d-30c2-4b9d-be36-07ae7a64fa57",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking were not used when initializing BertForSequenceClassification: ['cls.predictions.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "  0%|                                                  | 0/5000 [00:00<?, ?it/s]2023-08-14 08:15:20.000679: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:15:20.000681: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_15770068861141756492+abb26765/model.neff. Exiting with a successfully compiled graph.\n",
      "  0%|                                          | 1/5000 [00:00<25:42,  3.24it/s]2023-08-14 08:15:21.000792: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:15:21.000871: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_16268669036401171959+abb26765/model.neff. Exiting with a successfully compiled graph.\n",
      "  0%|                                        | 2/5000 [00:02<1:44:31,  1.25s/it]2023-08-14 08:15:24.000441: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:15:24.000534: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_14409091931058726813+abb26765/model.neff. Exiting with a successfully compiled graph.\n",
      " 10%|████                                    | 500/5000 [00:44<05:46, 12.97it/s]2023-08-14 08:16:05.000406: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      ".\n",
      "Compiler status PASS\n",
      "{'loss': 0.29, 'learning_rate': 4.5e-05, 'epoch': 1.0}                          \n",
      " 10%|████                                    | 500/5000 [01:01<05:46, 12.97it/s]2023-08-14 08:16:22.000141: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      ".\n",
      "Compiler status PASS\n",
      "{'loss': 0.1135, 'learning_rate': 4e-05, 'epoch': 2.0}                          \n",
      "{'loss': 0.0393, 'learning_rate': 3.5e-05, 'epoch': 3.0}                        \n",
      "{'loss': 0.0201, 'learning_rate': 3e-05, 'epoch': 4.0}                          \n",
      "{'loss': 0.0107, 'learning_rate': 2.5e-05, 'epoch': 5.0}                        \n",
      "{'loss': 0.0028, 'learning_rate': 2e-05, 'epoch': 6.0}                          \n",
      "{'loss': 0.0015, 'learning_rate': 1.5e-05, 'epoch': 7.0}                        \n",
      "{'loss': 0.0014, 'learning_rate': 1e-05, 'epoch': 8.0}                          \n",
      "{'loss': 0.0, 'learning_rate': 5e-06, 'epoch': 9.0}                             \n",
      "{'loss': 0.0, 'learning_rate': 0.0, 'epoch': 10.0}                              \n",
      "{'train_runtime': 490.9411, 'train_samples_per_second': 81.476, 'train_steps_per_second': 10.185, 'train_loss': 0.047935736083984375, 'epoch': 10.0}\n",
      "100%|███████████████████████████████████████| 5000/5000 [08:10<00:00, 10.20it/s]\n",
      "TrainOutput(global_step=5000, training_loss=0.047935736083984375, metrics={'train_runtime': 490.9411, 'train_samples_per_second': 81.476, 'train_steps_per_second': 10.185, 'train_loss': 0.047935736083984375, 'epoch': 10.0})\n",
      "2023-08-14 08:23:31.000055: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "..\n",
      "Compiler status PASS\n",
      "2023-08-14 08:24:03.000422: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      ".\n",
      "Compiler status PASS\n",
      "100%|███████████████████████████████████████████| 32/32 [00:00<00:00, 36.30it/s]\n",
      "{'eval_loss': 0.3578033447265625, 'eval_runtime': 39.847, 'eval_samples_per_second': 6.425, 'eval_steps_per_second': 0.803, 'epoch': 10.0}\n",
      "\n",
      "real\t9m6.026s\n",
      "user\t9m23.274s\n",
      "sys\t1m22.760s\n"
     ]
    }
   ],
   "source": [
    "!time XLA_USE_BF16=1 python3 bert-jp-single.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aae6d30-f843-45fd-b6ef-65a65f61763e",
   "metadata": {},
   "source": [
    "ステップ数 5000 のトレーニングが 9~10分程で完了しました。\n",
    "\n",
    "トレーニング実行中に、AWS Neuron で提供される `neuron-top` ツールを利用すると、Neuron コア及び vCPU の利用率、アクセラレータメモリ、ホストメモリの利用状況等を確認することができます。inf2.xlarge には、一つの Inferentia2 チップ、チップ内に二つの Neuron コアが搭載されています。結果から、二つある Neuron コア（NC0 及び NC1）のうち一つの Neuron コアのみが利用されていることが分かります。まだ最適化の余地はありそうです。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f4bb724-139b-4a2b-9e2c-c66023c65768",
   "metadata": {},
   "source": [
    "生成されたモデルから期待通りの出力が得られるか確認しておきましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7e5ea0d5-d506-42b5-9f24-0438bc5e742c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'label': 'LABEL_0', 'score': 0.9999897480010986}]\n",
      "[{'label': 'LABEL_1', 'score': 0.9999812841415405}]\n"
     ]
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "classifier = pipeline(\"text-classification\", model = \"./results/\")\n",
    "\n",
    "print(classifier(\"大変すばらしい商品でした。感激です。\"))\n",
    "print(classifier(\"期待していた商品とは異なりました。残念です。\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dfdbb2d-d477-42e3-9967-f64fba883b0d",
   "metadata": {},
   "source": [
    "期待通りの出力を得られることが確認できたようです。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28270199-5a61-4ec1-a045-833abc6c16e7",
   "metadata": {},
   "source": [
    "## torchrun を用いたマルチワーカーでのトレーニング実行\n",
    "それでは、先程のトレーニングスクリプトに変更を加え、二つある Neuron コアを有効活用してみましょう。複数の Neuron コアを利用したマルチワーカーで実行するためには `torchrun` コマンドを利用します。`torchrun` コマンドに対して、オプション `--nproc_per_node` で利用する Neuron コアの数（並列実行するワーカー数）を指定します。trn1.2xlarge (inf2.xlarge) では 2 を、trn1.32xlargeの場合は 2, 8, 32 が指定可能です。\n",
    "\n",
    "`torchrun` を利用したデータパラレルトレーニングを実行するにあたって、先程のスクリプトに一部変更を加えた `bert-jp-dual.py` というファイル名のスクリプトを作成し実行します。\n",
    "\n",
    "それでは変更後のスクリプトを利用して　inf2.xlarge 上の二つ Neuron コアを利用したトレーニングを実行してみましょう。シングルワーカーでのトレーニング結果と比較し `Total train batch size` の値が倍の 16 に、`Total optimization steps` が半分の 2500 となっている点を確認できると思います。\n",
    "\n",
    "シングルワーカー時の手順同様、まずは事前コンパイルを実行し、その後に実際のトレーニングを実行します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f2837ccc-aac9-47d7-803e-3c75afccc426",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting bert-jp-dual-precompile.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile bert-jp-dual-precompile.py\n",
    "\n",
    "from transformers import BertForSequenceClassification, BertJapaneseTokenizer, Trainer, TrainingArguments\n",
    "from datasets import load_from_disk\n",
    "import torch, torch_xla.distributed.xla_backend\n",
    "import os\n",
    "\n",
    "os.environ[\"NEURON_CC_FLAGS\"] = \"--model-type=transformer\"\n",
    "\n",
    "MODEL_NAME = \"cl-tohoku/bert-base-japanese-whole-word-masking\"\n",
    "model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)\n",
    "\n",
    "train_dataset = load_from_disk(\"./train/\").with_format(\"torch\")\n",
    "train_dataset = train_dataset.select(range(64))\n",
    "\n",
    "eval_dataset = load_from_disk(\"./test/\").with_format(\"torch\")\n",
    "eval_dataset = eval_dataset.select(range(64))\n",
    "\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    num_train_epochs = 2,\n",
    "    learning_rate = 5e-5,\n",
    "    per_device_train_batch_size = 8,\n",
    "    per_device_eval_batch_size = 8,\n",
    "    output_dir = \"./results\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_args,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = eval_dataset,\n",
    ")\n",
    "\n",
    "train_result = trainer.train()\n",
    "print(train_result)\n",
    "\n",
    "eval_result = trainer.evaluate()\n",
    "print(eval_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "50d164bb-deb5-4c34-bafa-8c66b2291ff5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-08-14 08:24:16.000996:  904568  INFO ||PARALLEL_COMPILE||: Removing existing workdir /tmp/parallel_compile_workdir\n",
      "2023-08-14 08:24:16.000998:  904568  INFO ||PARALLEL_COMPILE||: Running trial run (add option to terminate trial run early; also ignore trial run's generated outputs, i.e. loss, checkpoints)\n",
      "WARNING:torch.distributed.run:\n",
      "*****************************************\n",
      "Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. \n",
      "*****************************************\n",
      "Some weights of the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking were not used when initializing BertForSequenceClassification: ['cls.predictions.decoder.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Some weights of the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking were not used when initializing BertForSequenceClassification: ['cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "/opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "2023-08-14 08:24:32.000310: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:32.000312: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_15770068861141756492+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      " 12%|█▎        | 1/8 [00:00<00:00,  8.26it/s]2023-08-14 08:24:33.000463: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:33.000526: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_4976758461586571876+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      " 25%|██▌       | 2/8 [00:01<00:05,  1.14it/s]2023-08-14 08:24:35.000400: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:35.000477: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_1129850769929414091+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      " 88%|████████▊ | 7/8 [00:03<00:00,  2.50it/s]2023-08-14 08:24:37.000229: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:37.000244: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_12573384035295665077+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "2023-08-14 08:24:38.000045: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:38.000099: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_16102519548021684275+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "{'train_runtime': 5.2387, 'train_samples_per_second': 24.434, 'train_steps_per_second': 1.527, 'train_loss': -4.237517714500427e-08, 'epoch': 2.0}\n",
      "100%|██████████| 8/8 [00:05<00:00,  2.50it/s]TrainOutput(global_step=8, training_loss=1.0, metrics={'train_runtime': 5.2386, 'train_samples_per_second': 24.434, 'train_steps_per_second': 1.527, 'total_flos': 9.183549615799121e-41, 'train_loss': 1.0, 'epoch': 2.0})\n",
      "100%|██████████| 8/8 [00:05<00:00,  1.38it/s]\n",
      "TrainOutput(global_step=8, training_loss=-4.237517714500427e-08, metrics={'train_runtime': 5.2387, 'train_samples_per_second': 24.434, 'train_steps_per_second': 1.527, 'train_loss': -4.237517714500427e-08, 'epoch': 2.0})\n",
      "2023-08-14 08:24:38.000985: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:38.000999: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_2296725661428058300+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "2023-08-14 08:24:39.000242: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:39.000243: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_13837218930383410181+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "2023-08-14 08:24:39.000757: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:39.000835: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_9481639979634615951+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "  0%|          | 0/4 [00:00<?, ?it/s]{'eval_loss': -1.4657210065138902e+30, 'eval_runtime': 2.4316, 'eval_samples_per_second': 26.321, 'eval_steps_per_second': 1.645, 'epoch': 2.0}\n",
      "100%|██████████| 4/4 [00:00<00:00, 55.01it/s]\n",
      "{'eval_loss': -1.4657210065138902e+30, 'eval_runtime': 2.4306, 'eval_samples_per_second': 26.331, 'eval_steps_per_second': 1.646, 'epoch': 2.0}\n",
      "2023-08-14 08:24:43.000337:  904568  INFO ||PARALLEL_COMPILE||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:24:49.000714:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:49.000782:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:49.000795:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:49.000853:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:49.000853:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:49.000854:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:49.000939:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:49.000939:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:49.000998:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:49.000999:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:50.000070:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:50.000128:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:50.000157:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:50.000157:  904568  INFO ||PARALLEL_COMPILE||: Current remaining items are 0, locked are 0, failed are 0, done are 59, total is 59\n",
      "2023-08-14 08:24:50.000157:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:50.000157:  904568  INFO ||PARALLEL_COMPILE||: Progress Status: No more compile jobs!\n",
      "2023-08-14 08:24:50.000158:  904568  INFO ||PARALLEL_COMPILE||: {\n",
      "    \"compilation_summary\": {},\n",
      "    \"compilation_report\": {}\n",
      "}\n",
      "2023-08-14 08:24:50.000158:  904568  INFO ||PARALLEL_COMPILE||: Total graphs: 0\n",
      "2023-08-14 08:24:50.000158:  904568  INFO ||PARALLEL_COMPILE||: Total successful compilations: 0\n",
      "2023-08-14 08:24:50.000158:  904568  INFO ||PARALLEL_COMPILE||: Total failed compilations: 0\n",
      "\n",
      "real\t0m33.343s\n",
      "user\t0m16.984s\n",
      "sys\t0m4.249s\n"
     ]
    }
   ],
   "source": [
    "!time XLA_USE_BF16=1 neuron_parallel_compile torchrun --nproc_per_node=2 bert-jp-dual-precompile.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "60bddfc6-2cb6-4383-b095-e6cb58a9e68c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting bert-jp-dual.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile bert-jp-dual.py\n",
    "\n",
    "from transformers import BertForSequenceClassification, BertJapaneseTokenizer, Trainer, TrainingArguments\n",
    "from datasets import load_from_disk\n",
    "import torch, torch_xla.distributed.xla_backend\n",
    "import os\n",
    "\n",
    "os.environ[\"NEURON_CC_FLAGS\"] = \"--model-type=transformer\"\n",
    "\n",
    "MODEL_NAME = \"cl-tohoku/bert-base-japanese-whole-word-masking\"\n",
    "tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)\n",
    "model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)\n",
    "\n",
    "train_dataset = load_from_disk(\"./train/\").with_format(\"torch\")\n",
    "eval_dataset = load_from_disk(\"./test/\").with_format(\"torch\")\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    num_train_epochs = 10,\n",
    "    learning_rate = 5e-5,\n",
    "    per_device_train_batch_size = 8,\n",
    "    per_device_eval_batch_size = 8,\n",
    "    output_dir = \"./results\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_args,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = eval_dataset,\n",
    "    tokenizer = tokenizer,\n",
    ")\n",
    "\n",
    "train_result = trainer.train()\n",
    "print(train_result)\n",
    "\n",
    "eval_result = trainer.evaluate()\n",
    "print(eval_result)\n",
    "\n",
    "trainer.save_model(\"./results\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f73d8ec0-162b-4d6a-bebf-a0b991b08213",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:torch.distributed.run:\n",
      "*****************************************\n",
      "Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. \n",
      "*****************************************\n",
      "Some weights of the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking were not used when initializing BertForSequenceClassification: ['cls.predictions.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Some weights of the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking were not used when initializing BertForSequenceClassification: ['cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at cl-tohoku/bert-base-japanese-whole-word-masking and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "/opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "2023-08-14 08:25:05.000350: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:25:05.000352: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_15770068861141756492+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "  0%|                                          | 1/2500 [00:00<04:56,  8.43it/s]2023-08-14 08:25:06.000478: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:25:06.000542: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_4976758461586571876+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "2023-Aug-14 08:25:07.0366 906211:906257 [1] nccl_net_ofi_init:1415 CCOM WARN NET/OFI aws-ofi-nccl initialization failed\n",
      "2023-Aug-14 08:25:07.0366 906211:906257 [1] init.cc:138 CCOM WARN OFI plugin initNet() failed is EFA enabled?\n",
      "  0%|                                          | 2/2500 [00:01<36:49,  1.13it/s]2023-08-14 08:25:09.000040: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:25:09.000117: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_1129850769929414091+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      " 20%|████████                                | 500/2500 [00:39<02:25, 13.71it/s]2023-08-14 08:25:45.000802: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:25:45.000802: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_569571592325889951+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "2023-08-14 08:25:45.000976: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      ".\n",
      "Compiler status PASS\n",
      "{'loss': 0.221, 'learning_rate': 4e-05, 'epoch': 2.0}                           \n",
      " 20%|████████                                | 500/2500 [00:46<02:25, 13.71it/s]2023-08-14 08:25:52.000407: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      " 20%|████████                                | 500/2500 [00:50<02:25, 13.71it/s].Non-output memory location with no reader {_convert-t87}@SB<0,0>(1x4)#Internal DebugInfo: <_convert||UNDEF||[1, 1, 1]>\n",
      "\n",
      "Compiler status PASS\n",
      "2023-08-14 08:26:09.000822: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:26:09.000823: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_931450859643720011+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "{'loss': 0.0389, 'learning_rate': 3e-05, 'epoch': 4.0}                          \n",
      "{'loss': 0.0122, 'learning_rate': 2e-05, 'epoch': 6.0}                          \n",
      "{'loss': 0.0032, 'learning_rate': 1e-05, 'epoch': 8.0}                          \n",
      "{'loss': 0.002, 'learning_rate': 0.0, 'epoch': 10.0}                            \n",
      "100%|███████████████████████████████████████| 2500/2500 [04:01<00:00, 16.01it/s]2023-08-14 08:29:07.000336: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      ".Non-output memory location with no reader {_convert-t106}@SB<0,0>(1x4)#Internal DebugInfo: <_convert||UNDEF||[1, 1, 1]>\n",
      "\n",
      "Compiler status PASS\n",
      "{'train_runtime': 242.0433, 'train_samples_per_second': 165.26, 'train_steps_per_second': 10.329, 'train_loss': 0.05544990234375, 'epoch': 10.0}\n",
      "TrainOutput(global_step=2500, training_loss=0.05544990234375, metrics={'train_runtime': 242.0146, 'train_samples_per_second': 165.279, 'train_steps_per_second': 10.33, 'total_flos': 2627832790384640.0, 'train_loss': 0.05544990234375, 'epoch': 10.0})\n",
      "100%|███████████████████████████████████████| 2500/2500 [04:07<00:00, 10.09it/s]\n",
      "TrainOutput(global_step=2500, training_loss=0.05544990234375, metrics={'train_runtime': 242.0433, 'train_samples_per_second': 165.26, 'train_steps_per_second': 10.329, 'train_loss': 0.05544990234375, 'epoch': 10.0})\n",
      "2023-08-14 08:29:13.000996: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:29:14.000009: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_2296725661428058300+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      "2023-08-14 08:29:14.000304: INFO ||NCC_WRAPPER||: Compile cache path: /var/tmp/neuron-compile-cache\n",
      "2023-08-14 08:29:14.000305: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.8.0.25+a3ad0f342/MODULE_13837218930383410181+d41d8cd9/model.neff. Exiting with a successfully compiled graph.\n",
      " 88%|█████████████████████████████████████▋     | 14/16 [00:00<00:00, 31.64it/s]{'eval_loss': 0.39816153049468994, 'eval_runtime': 1.1355, 'eval_samples_per_second': 225.458, 'eval_steps_per_second': 14.091, 'epoch': 10.0}\n",
      "100%|███████████████████████████████████████████| 16/16 [00:00<00:00, 32.58it/s]\n",
      "{'eval_loss': 0.39816153049468994, 'eval_runtime': 1.1343, 'eval_samples_per_second': 225.685, 'eval_steps_per_second': 14.105, 'epoch': 10.0}\n",
      "\n",
      "real\t4m31.237s\n",
      "user\t6m6.730s\n",
      "sys\t0m18.259s\n"
     ]
    }
   ],
   "source": [
    "!time XLA_USE_BF16=1 torchrun --nproc_per_node=2 bert-jp-dual.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "582fc974-a3fa-45ad-ba0b-1d3ac99d8f33",
   "metadata": {},
   "source": [
    "トレーニング実行中の neuron-top の出力も確認してみましょう。今度は二つの Neuron コアが利用されている事が確認できると思います。トレーニングに要する実行時間も 5~6 分に削減されました。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "293c4c41-9152-4b02-91fa-d5897079427f",
   "metadata": {},
   "source": [
    "## 推論実行\n",
    "先ほどは生成されたモデルから期待通りの出力が得られるかどうかCPU上で推論実行し、結果を確認しました。ここでは生成されたモデルをinf2.xlarge上で推論実行します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0b881d0f-28b5-49e5-849f-3f0266daa97c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU paraphrase logits: [[ 4.5357184 -4.789145 ]]\n",
      "CPU paraphrase logits: [[-4.6640625  4.8995476]]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch_neuronx\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "import transformers\n",
    "\n",
    "from transformers import BertForSequenceClassification, BertJapaneseTokenizer\n",
    "\n",
    "def encode(tokenizer, *inputs, max_length=128, batch_size=1):\n",
    "    tokens = tokenizer.encode_plus(\n",
    "        *inputs,\n",
    "        max_length=max_length,\n",
    "        padding='max_length',\n",
    "        truncation=True,\n",
    "        return_tensors=\"pt\"\n",
    "    )\n",
    "    return (\n",
    "        torch.repeat_interleave(tokens['input_ids'], batch_size, 0),\n",
    "        torch.repeat_interleave(tokens['attention_mask'], batch_size, 0),\n",
    "        torch.repeat_interleave(tokens['token_type_ids'], batch_size, 0),\n",
    "    )\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"./results/\")\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\"./results/\", torchscript=True)\n",
    "\n",
    "\n",
    "sequence = \"大変すばらしい商品でした。感激です。\"\n",
    "paraphrase = encode(tokenizer, sequence)\n",
    "cpu_paraphrase_logits = model(*paraphrase)[0]\n",
    "print('CPU paraphrase logits:', cpu_paraphrase_logits.detach().numpy())\n",
    "\n",
    "sequence = \"期待していた商品とは異なりました。残念です。\"\n",
    "paraphrase = encode(tokenizer, sequence)\n",
    "cpu_paraphrase_logits = model(*paraphrase)[0]\n",
    "print('CPU paraphrase logits:', cpu_paraphrase_logits.detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e608f7d-16cb-4dc0-b7ac-253129e64df2",
   "metadata": {},
   "source": [
    "### Compile the model for Neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "09bc7d66-7fdb-42d8-91f3-98708982fee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = 'model.pt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ee754ea0-97ae-4cd4-8184-4cd3b7fa6fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_neuron = torch_neuronx.trace(model, paraphrase)\n",
    "\n",
    "# Save the TorchScript for inference deployment\n",
    "torch.jit.save(model_neuron, filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f805991-1929-4899-aa6d-46682b7e3acb",
   "metadata": {},
   "source": [
    "### Run inference and compare results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b81e935f-0eab-4128-85db-047e5022e2c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neuron paraphrase logits: [[ 4.5336947 -4.7886157]]\n",
      "Neuron paraphrase logits: [[-4.663004   4.9002094]]\n"
     ]
    }
   ],
   "source": [
    "model_neuron = torch.jit.load(filename)\n",
    "\n",
    "sequence = \"大変すばらしい商品でした。感激です。\"\n",
    "paraphrase = encode(tokenizer, sequence)\n",
    "neuron_paraphrase_logits = model_neuron(*paraphrase)[0]\n",
    "print('Neuron paraphrase logits:', neuron_paraphrase_logits.detach().numpy())\n",
    "\n",
    "sequence = \"期待していた商品とは異なりました。残念です。\"\n",
    "paraphrase = encode(tokenizer, sequence)\n",
    "neuron_paraphrase_logits = model_neuron(*paraphrase)[0]\n",
    "print('Neuron paraphrase logits:', neuron_paraphrase_logits.detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22e3299b-f8ec-4083-8813-54b693270cb0",
   "metadata": {},
   "source": [
    "CPUで推論実行した結果と同様の結果が得られている事が確認できました。推論性能を評価する方法は以下のサンプルをご参照下さい。\n",
    "+ https://github.com/aws-neuron/aws-neuron-sdk/blob/master/src/examples/pytorch/torch-neuronx/bert-base-cased-finetuned-mrpc-inference-on-trn1-tutorial.ipynb"
   ]
  }
 ],
 "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
}