gpt-j-t4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6f680dc0",
   "metadata": {},
   "source": [
    "# Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb0e23de",
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo apt-get install -y nvidia-driver-470"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3d5f6690",
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting torch\n",
      "  Downloading torch-1.9.0-cp38-cp38-manylinux1_x86_64.whl (831.4 MB)\n",
      "\u001b[K     |████████████▍                   | 321.7 MB 147.5 MB/s eta 0:00:04"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[K     |██████████████████████████▉     | 696.7 MB 138.6 MB/s eta 0:00:01"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[K     |████████████████████████████████| 831.4 MB 10 kB/s \n",
      "\u001b[?25hCollecting typing-extensions\n",
      "  Downloading typing_extensions-3.10.0.0-py3-none-any.whl (26 kB)\n",
      "Installing collected packages: typing-extensions, torch\n",
      "Successfully installed torch-1.9.0 typing-extensions-3.10.0.0\n",
      "fatal: destination path 'mesh-transformer-jax' already exists and is not an empty directory.\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting git+https://github.com/deepmind/dm-haiku (from -r mesh-transformer-jax/requirements.txt (line 10))\n",
      "  Cloning https://github.com/deepmind/dm-haiku to /tmp/pip-req-build-gui19b49\n",
      "  Running command git clone -q https://github.com/deepmind/dm-haiku /tmp/pip-req-build-gui19b49\n",
      "Collecting git+https://github.com/EleutherAI/lm-evaluation-harness/ (from -r mesh-transformer-jax/requirements.txt (line 11))\n",
      "  Cloning https://github.com/EleutherAI/lm-evaluation-harness/ to /tmp/pip-req-build-w0imism7\n",
      "  Running command git clone -q https://github.com/EleutherAI/lm-evaluation-harness/ /tmp/pip-req-build-w0imism7\n",
      "Collecting numpy~=1.19.5\n",
      "  Downloading numpy-1.19.5-cp38-cp38-manylinux2010_x86_64.whl (14.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 14.9 MB 22.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting transformers~=4.4.2\n",
      "  Downloading transformers-4.4.2-py3-none-any.whl (2.0 MB)\n",
      "\u001b[K     |████████████████████████████████| 2.0 MB 124.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tqdm~=4.45.0\n",
      "  Downloading tqdm-4.45.0-py2.py3-none-any.whl (60 kB)\n",
      "\u001b[K     |████████████████████████████████| 60 kB 13.4 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting setuptools~=51.3.3\n",
      "  Downloading setuptools-51.3.3-py3-none-any.whl (786 kB)\n",
      "\u001b[K     |████████████████████████████████| 786 kB 98.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting wandb~=0.10.22\n",
      "  Downloading wandb-0.10.32-py2.py3-none-any.whl (1.8 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.8 MB 110.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting einops~=0.3.0\n",
      "  Downloading einops-0.3.0-py2.py3-none-any.whl (25 kB)\n",
      "Collecting requests~=2.25.1\n",
      "  Downloading requests-2.25.1-py2.py3-none-any.whl (61 kB)\n",
      "\u001b[K     |████████████████████████████████| 61 kB 16.3 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting fabric~=2.6.0\n",
      "  Downloading fabric-2.6.0-py2.py3-none-any.whl (53 kB)\n",
      "\u001b[K     |████████████████████████████████| 53 kB 4.0 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting optax==0.0.6\n",
      "  Downloading optax-0.0.6-py3-none-any.whl (96 kB)\n",
      "\u001b[K     |████████████████████████████████| 96 kB 10.6 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting ray~=1.2.0\n",
      "  Downloading ray-1.2.0-cp38-cp38-manylinux2014_x86_64.whl (47.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 47.3 MB 84.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting jax~=0.2.12\n",
      "  Downloading jax-0.2.14.tar.gz (669 kB)\n",
      "\u001b[K     |████████████████████████████████| 669 kB 89.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting Flask~=1.1.2\n",
      "  Downloading Flask-1.1.4-py2.py3-none-any.whl (94 kB)\n",
      "\u001b[K     |████████████████████████████████| 94 kB 7.4 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting cloudpickle~=1.3.0\n",
      "  Downloading cloudpickle-1.3.0-py2.py3-none-any.whl (26 kB)\n",
      "Collecting tensorflow-cpu~=2.4.1\n",
      "  Downloading tensorflow_cpu-2.4.2-cp38-cp38-manylinux2010_x86_64.whl (144.4 MB)\n",
      "\u001b[K     |████████████████████████████████| 144.4 MB 117.4 MB/s eta 0:00:01   |███████▍                        | 33.1 MB 34.4 MB/s eta 0:00:04\n",
      "\u001b[?25hCollecting google-cloud-storage~=1.36.2\n",
      "  Downloading google_cloud_storage-1.36.2-py2.py3-none-any.whl (97 kB)\n",
      "\u001b[K     |████████████████████████████████| 97 kB 13.2 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting smart_open[gcs]\n",
      "  Downloading smart_open-5.1.0-py3-none-any.whl (57 kB)\n",
      "\u001b[K     |████████████████████████████████| 57 kB 11.8 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting func_timeout\n",
      "  Downloading func_timeout-4.3.5.tar.gz (44 kB)\n",
      "\u001b[K     |████████████████████████████████| 44 kB 5.1 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting ftfy\n",
      "  Downloading ftfy-6.0.3.tar.gz (64 kB)\n",
      "\u001b[K     |████████████████████████████████| 64 kB 6.6 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting absl-py>=0.7.1\n",
      "  Downloading absl_py-0.13.0-py3-none-any.whl (132 kB)\n",
      "\u001b[K     |████████████████████████████████| 132 kB 114.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting jmp>=0.0.2\n",
      "  Downloading jmp-0.0.2-py3-none-any.whl (16 kB)\n",
      "Collecting tabulate>=0.8.9\n",
      "  Downloading tabulate-0.8.9-py3-none-any.whl (25 kB)\n",
      "Collecting best_download>=0.0.6\n",
      "  Downloading best_download-0.0.6-py3-none-any.whl (4.5 kB)\n",
      "Collecting black==20.8b1\n",
      "  Downloading black-20.8b1.tar.gz (1.1 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.1 MB 78.7 MB/s eta 0:00:01\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h    Preparing wheel metadata ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting click>=7.1\n",
      "  Downloading click-8.0.1-py3-none-any.whl (97 kB)\n",
      "\u001b[K     |████████████████████████████████| 97 kB 14.2 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting datasets>=1.2.1\n",
      "  Downloading datasets-1.8.0-py3-none-any.whl (237 kB)\n",
      "\u001b[K     |████████████████████████████████| 237 kB 134.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting jsonlines==2.0.0\n",
      "  Downloading jsonlines-2.0.0-py3-none-any.whl (6.3 kB)\n",
      "Collecting lm_dataformat==0.0.19\n",
      "  Downloading lm_dataformat-0.0.19-py3-none-any.whl (5.4 kB)\n",
      "Collecting mock==4.0.3\n",
      "  Downloading mock-4.0.3-py3-none-any.whl (28 kB)\n",
      "Collecting numexpr==2.7.2\n",
      "  Downloading numexpr-2.7.2-cp38-cp38-manylinux2010_x86_64.whl (472 kB)\n",
      "\u001b[K     |████████████████████████████████| 472 kB 90.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting openai==0.6.4\n",
      "  Downloading openai-0.6.4.tar.gz (159 kB)\n",
      "\u001b[K     |████████████████████████████████| 159 kB 131.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pybind11==2.6.2\n",
      "  Downloading pybind11-2.6.2-py2.py3-none-any.whl (191 kB)\n",
      "\u001b[K     |████████████████████████████████| 191 kB 119.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pycountry==20.7.3\n",
      "  Downloading pycountry-20.7.3.tar.gz (10.1 MB)\n",
      "\u001b[K     |████████████████████████████████| 10.1 MB 54.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pytablewriter==0.58.0\n",
      "  Downloading pytablewriter-0.58.0-py3-none-any.whl (96 kB)\n",
      "\u001b[K     |████████████████████████████████| 96 kB 12.8 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting pytest==6.2.3\n",
      "  Downloading pytest-6.2.3-py3-none-any.whl (280 kB)\n",
      "\u001b[K     |████████████████████████████████| 280 kB 110.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting sacrebleu==1.5.0\n",
      "  Downloading sacrebleu-1.5.0-py3-none-any.whl (65 kB)\n",
      "\u001b[K     |████████████████████████████████| 65 kB 8.8 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting scikit-learn>=0.24.1\n",
      "  Downloading scikit_learn-0.24.2-cp38-cp38-manylinux2010_x86_64.whl (24.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 24.9 MB 53.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting sqlitedict==1.6.0\n",
      "  Downloading sqlitedict-1.6.0.tar.gz (29 kB)\n",
      "Requirement already satisfied: torch>=1.7 in /home/ubuntu/.local/lib/python3.8/site-packages (from lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.9.0)\n",
      "Collecting tqdm-multiprocess==0.0.11\n",
      "  Downloading tqdm_multiprocess-0.0.11-py3-none-any.whl (9.8 kB)\n",
      "Collecting zstandard==0.15.2\n",
      "  Downloading zstandard-0.15.2-cp38-cp38-manylinux2014_x86_64.whl (2.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 2.2 MB 52.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting jaxlib>=0.1.37\n",
      "  Downloading jaxlib-0.1.67-cp38-none-manylinux2010_x86_64.whl (45.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 45.3 MB 67.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting chex>=0.0.4\n",
      "  Downloading chex-0.0.7-py3-none-any.whl (52 kB)\n",
      "\u001b[K     |████████████████████████████████| 52 kB 2.3 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting regex>=2020.1.8\n",
      "  Downloading regex-2021.4.4-cp38-cp38-manylinux2014_x86_64.whl (733 kB)\n",
      "\u001b[K     |████████████████████████████████| 733 kB 109.8 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pathspec<1,>=0.6\n",
      "  Downloading pathspec-0.8.1-py2.py3-none-any.whl (28 kB)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4 in /home/ubuntu/.local/lib/python3.8/site-packages (from black==20.8b1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (3.10.0.0)\n",
      "Collecting typed-ast>=1.4.0\n",
      "  Downloading typed_ast-1.4.3-cp38-cp38-manylinux1_x86_64.whl (774 kB)\n",
      "\u001b[K     |████████████████████████████████| 774 kB 107.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting mypy-extensions>=0.4.3\n",
      "  Downloading mypy_extensions-0.4.3-py2.py3-none-any.whl (4.5 kB)\n",
      "Collecting appdirs\n",
      "  Downloading appdirs-1.4.4-py2.py3-none-any.whl (9.6 kB)\n",
      "Collecting toml>=0.10.1\n",
      "  Downloading toml-0.10.2-py2.py3-none-any.whl (16 kB)\n",
      "Collecting ujson\n",
      "  Downloading ujson-4.0.2-cp38-cp38-manylinux1_x86_64.whl (181 kB)\n",
      "\u001b[K     |████████████████████████████████| 181 kB 129.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting msgfy<1,>=0.1.0\n",
      "  Downloading msgfy-0.1.0-py3-none-any.whl (4.3 kB)\n",
      "Collecting tcolorpy<1,>=0.0.5\n",
      "  Downloading tcolorpy-0.1.0-py3-none-any.whl (8.0 kB)\n",
      "Collecting pathvalidate<3,>=2.3.0\n",
      "  Downloading pathvalidate-2.4.1-py3-none-any.whl (18 kB)\n",
      "Collecting tabledata<2,>=1.1.3\n",
      "  Downloading tabledata-1.1.4-py3-none-any.whl (11 kB)\n",
      "Collecting typepy[datetime]<2,>=1.1.1\n",
      "  Downloading typepy-1.1.5-py3-none-any.whl (30 kB)\n",
      "Collecting mbstrdecoder<2,>=1.0.0\n",
      "  Downloading mbstrdecoder-1.0.1-py3-none-any.whl (7.8 kB)\n",
      "Collecting DataProperty<2,>=0.50.0\n",
      "  Downloading DataProperty-0.50.1-py3-none-any.whl (26 kB)\n",
      "Requirement already satisfied: packaging in /home/ubuntu/.local/lib/python3.8/site-packages (from pytest==6.2.3->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (20.9)\n",
      "Collecting py>=1.8.2\n",
      "  Downloading py-1.10.0-py2.py3-none-any.whl (97 kB)\n",
      "\u001b[K     |████████████████████████████████| 97 kB 16.6 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting iniconfig\n",
      "  Downloading iniconfig-1.1.1-py2.py3-none-any.whl (5.0 kB)\n",
      "Collecting pluggy<1.0.0a1,>=0.12\n",
      "  Downloading pluggy-0.13.1-py2.py3-none-any.whl (18 kB)\n",
      "Requirement already satisfied: attrs>=19.2.0 in /usr/lib/python3/dist-packages (from pytest==6.2.3->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (19.3.0)\n",
      "Collecting portalocker\n",
      "  Downloading portalocker-2.3.0-py2.py3-none-any.whl (15 kB)\n",
      "Requirement already satisfied: colorama in /usr/lib/python3/dist-packages (from tqdm-multiprocess==0.0.11->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (0.4.3)\n",
      "Collecting tokenizers<0.11,>=0.10.1\n",
      "  Downloading tokenizers-0.10.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (3.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 3.3 MB 40.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting sacremoses\n",
      "  Downloading sacremoses-0.0.45-py3-none-any.whl (895 kB)\n",
      "\u001b[K     |████████████████████████████████| 895 kB 83.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting filelock\n",
      "  Downloading filelock-3.0.12-py3-none-any.whl (7.6 kB)\n",
      "Collecting configparser>=3.8.1\n",
      "  Downloading configparser-5.0.2-py3-none-any.whl (19 kB)\n",
      "Collecting GitPython>=1.0.0\n",
      "  Downloading GitPython-3.1.14-py3-none-any.whl (159 kB)\n",
      "\u001b[K     |████████████████████████████████| 159 kB 127.8 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting promise<3,>=2.0\n",
      "  Downloading promise-2.3.tar.gz (19 kB)\n",
      "Collecting shortuuid>=0.5.0\n",
      "  Downloading shortuuid-1.0.1-py3-none-any.whl (7.5 kB)\n",
      "Requirement already satisfied: python-dateutil>=2.6.1 in /home/ubuntu/.local/lib/python3.8/site-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (2.8.1)\n",
      "Collecting sentry-sdk>=0.4.0\n",
      "  Downloading sentry_sdk-1.1.0-py2.py3-none-any.whl (131 kB)\n",
      "\u001b[K     |████████████████████████████████| 131 kB 113.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting subprocess32>=3.5.3\n",
      "  Downloading subprocess32-3.5.4.tar.gz (97 kB)\n",
      "\u001b[K     |████████████████████████████████| 97 kB 14.7 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting protobuf>=3.12.0\n",
      "  Downloading protobuf-3.17.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.0 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.0 MB 71.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting docker-pycreds>=0.4.0\n",
      "  Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)\n",
      "Collecting pathtools\n",
      "  Downloading pathtools-0.1.2.tar.gz (11 kB)\n",
      "Collecting psutil>=5.0.0\n",
      "  Downloading psutil-5.8.0-cp38-cp38-manylinux2010_x86_64.whl (296 kB)\n",
      "\u001b[K     |████████████████████████████████| 296 kB 117.3 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: six>=1.13.0 in /usr/lib/python3/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (1.14.0)\n",
      "Requirement already satisfied: PyYAML in /usr/lib/python3/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (5.3.1)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (2.8)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (1.25.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (2019.11.28)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (3.0.4)\n",
      "Collecting pathlib2\n",
      "  Downloading pathlib2-2.3.5-py2.py3-none-any.whl (18 kB)\n",
      "Collecting paramiko>=2.4\n",
      "  Downloading paramiko-2.7.2-py2.py3-none-any.whl (206 kB)\n",
      "\u001b[K     |████████████████████████████████| 206 kB 127.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting invoke<2.0,>=1.3\n",
      "  Downloading invoke-1.5.0-py3-none-any.whl (211 kB)\n",
      "\u001b[K     |████████████████████████████████| 211 kB 126.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting colorful\n",
      "  Downloading colorful-0.5.4-py2.py3-none-any.whl (201 kB)\n",
      "\u001b[K     |████████████████████████████████| 201 kB 132.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting msgpack<2.0.0,>=1.0.0\n",
      "  Downloading msgpack-1.0.2-cp38-cp38-manylinux1_x86_64.whl (302 kB)\n",
      "\u001b[K     |████████████████████████████████| 302 kB 127.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting redis>=3.5.0\n",
      "  Downloading redis-3.5.3-py2.py3-none-any.whl (72 kB)\n",
      "\u001b[K     |████████████████████████████████| 72 kB 1.3 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting opencensus\n",
      "  Downloading opencensus-0.7.13-py2.py3-none-any.whl (127 kB)\n",
      "\u001b[K     |████████████████████████████████| 127 kB 134.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting py-spy>=0.2.0\n",
      "  Downloading py_spy-0.3.7-py2.py3-none-manylinux1_x86_64.whl (3.1 MB)\n",
      "\u001b[K     |████████████████████████████████| 3.1 MB 53.8 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting grpcio>=1.28.1\n",
      "  Downloading grpcio-1.38.0-cp38-cp38-manylinux2014_x86_64.whl (4.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.2 MB 69.2 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: prometheus-client>=0.7.1 in /home/ubuntu/.local/lib/python3.8/site-packages (from ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (0.11.0)\n",
      "Collecting gpustat\n",
      "  Downloading gpustat-0.6.0.tar.gz (78 kB)\n",
      "\u001b[K     |████████████████████████████████| 78 kB 15.8 MB/s  eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: jsonschema in /usr/lib/python3/dist-packages (from ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (3.2.0)\n",
      "Collecting aiohttp-cors\n",
      "  Downloading aiohttp_cors-0.7.0-py3-none-any.whl (27 kB)\n",
      "Collecting aioredis\n",
      "  Downloading aioredis-1.3.1-py3-none-any.whl (65 kB)\n",
      "\u001b[K     |████████████████████████████████| 65 kB 460 kB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting aiohttp\n",
      "  Downloading aiohttp-3.7.4.post0-cp38-cp38-manylinux2014_x86_64.whl (1.5 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.5 MB 52.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting opt_einsum\n",
      "  Downloading opt_einsum-3.3.0-py3-none-any.whl (65 kB)\n",
      "\u001b[K     |████████████████████████████████| 65 kB 310 kB/s s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: Jinja2<3.0,>=2.10.1 in /usr/lib/python3/dist-packages (from Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 14)) (2.10.1)\n",
      "Collecting itsdangerous<2.0,>=0.24\n",
      "  Downloading itsdangerous-1.1.0-py2.py3-none-any.whl (16 kB)\n",
      "Collecting Werkzeug<2.0,>=0.15\n",
      "  Downloading Werkzeug-1.0.1-py2.py3-none-any.whl (298 kB)\n",
      "\u001b[K     |████████████████████████████████| 298 kB 122.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting click>=7.1\n",
      "  Downloading click-7.1.2-py2.py3-none-any.whl (82 kB)\n",
      "\u001b[K     |████████████████████████████████| 82 kB 74 kB/s /s eta 0:00:01\n",
      "\u001b[?25hCollecting termcolor~=1.1.0\n",
      "  Downloading termcolor-1.1.0.tar.gz (3.9 kB)\n",
      "Collecting wrapt~=1.12.1\n",
      "  Downloading wrapt-1.12.1.tar.gz (27 kB)\n",
      "Collecting grpcio>=1.28.1\n",
      "  Downloading grpcio-1.32.0-cp38-cp38-manylinux2014_x86_64.whl (3.8 MB)\n",
      "\u001b[K     |████████████████████████████████| 3.8 MB 48.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting h5py~=2.10.0\n",
      "  Downloading h5py-2.10.0-cp38-cp38-manylinux1_x86_64.whl (2.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 2.9 MB 124.1 MB/s eta 0:00:01\u001b[K     |████████████████████████████████| 2.9 MB 124.1 MB/s \n",
      "\u001b[?25hCollecting gast==0.3.3\n",
      "  Downloading gast-0.3.3-py2.py3-none-any.whl (9.7 kB)\n",
      "Collecting flatbuffers~=1.12.0\n",
      "  Downloading flatbuffers-1.12-py2.py3-none-any.whl (15 kB)\n",
      "Collecting tensorflow-estimator<2.5.0,>=2.4.0\n",
      "  Downloading tensorflow_estimator-2.4.0-py2.py3-none-any.whl (462 kB)\n",
      "\u001b[K     |████████████████████████████████| 462 kB 97.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting google-pasta~=0.2\n",
      "  Downloading google_pasta-0.2.0-py3-none-any.whl (57 kB)\n",
      "\u001b[K     |████████████████████████████████| 57 kB 606 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting astunparse~=1.6.3\n",
      "  Downloading astunparse-1.6.3-py2.py3-none-any.whl (12 kB)\n",
      "Collecting typing-extensions>=3.7.4\n",
      "  Downloading typing_extensions-3.7.4.3-py3-none-any.whl (22 kB)\n",
      "Collecting wheel~=0.35\n",
      "  Using cached wheel-0.36.2-py2.py3-none-any.whl (35 kB)\n",
      "Collecting tensorboard~=2.4\n",
      "  Downloading tensorboard-2.5.0-py3-none-any.whl (6.0 MB)\n",
      "\u001b[K     |████████████████████████████████| 6.0 MB 50.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting keras-preprocessing~=1.1.2\n",
      "  Downloading Keras_Preprocessing-1.1.2-py2.py3-none-any.whl (42 kB)\n",
      "\u001b[K     |████████████████████████████████| 42 kB 78 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting six>=1.13.0\n",
      "  Downloading six-1.15.0-py2.py3-none-any.whl (10 kB)\n",
      "Collecting google-auth<2.0dev,>=1.11.0\n",
      "  Downloading google_auth-1.31.0-py2.py3-none-any.whl (147 kB)\n",
      "\u001b[K     |████████████████████████████████| 147 kB 120.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting google-cloud-core<2.0dev,>=1.4.1\n",
      "  Downloading google_cloud_core-1.7.0-py2.py3-none-any.whl (28 kB)\n",
      "Collecting google-resumable-media<2.0dev,>=1.2.0\n",
      "  Downloading google_resumable_media-1.3.0-py2.py3-none-any.whl (75 kB)\n",
      "\u001b[K     |████████████████████████████████| 75 kB 326 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting rehash\n",
      "  Downloading rehash-1.0.0-py2.py3-none-any.whl (8.7 kB)\n",
      "Collecting dm-tree>=0.1.5\n",
      "  Downloading dm_tree-0.1.6-cp38-cp38-manylinux_2_24_x86_64.whl (94 kB)\n",
      "\u001b[K     |████████████████████████████████| 94 kB 259 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting toolz>=0.9.0\n",
      "  Downloading toolz-0.11.1-py3-none-any.whl (55 kB)\n",
      "\u001b[K     |████████████████████████████████| 55 kB 242 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting huggingface-hub<0.1.0\n",
      "  Downloading huggingface_hub-0.0.10-py3-none-any.whl (37 kB)\n",
      "Collecting multiprocess\n",
      "  Downloading multiprocess-0.70.12.2-py38-none-any.whl (128 kB)\n",
      "\u001b[K     |████████████████████████████████| 128 kB 131.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pandas\n",
      "  Downloading pandas-1.2.4-cp38-cp38-manylinux1_x86_64.whl (9.7 MB)\n",
      "\u001b[K     |████████████████████████████████| 9.7 MB 51.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting dill\n",
      "  Downloading dill-0.3.4-py2.py3-none-any.whl (86 kB)\n",
      "\u001b[K     |████████████████████████████████| 86 kB 522 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting xxhash\n",
      "  Downloading xxhash-2.0.2-cp38-cp38-manylinux2010_x86_64.whl (243 kB)\n",
      "\u001b[K     |████████████████████████████████| 243 kB 119.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting fsspec\n",
      "  Downloading fsspec-2021.6.0-py3-none-any.whl (114 kB)\n",
      "\u001b[K     |████████████████████████████████| 114 kB 114.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pyarrow<4.0.0,>=1.0.0\n",
      "  Downloading pyarrow-3.0.0-cp38-cp38-manylinux2014_x86_64.whl (20.7 MB)\n",
      "\u001b[K     |████████████████████████████████| 20.7 MB 132.8 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting gitdb<5,>=4.0.1\n",
      "  Downloading gitdb-4.0.7-py3-none-any.whl (63 kB)\n",
      "\u001b[K     |████████████████████████████████| 63 kB 147 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting smmap<5,>=3.0.1\n",
      "  Downloading smmap-4.0.0-py2.py3-none-any.whl (24 kB)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/lib/python3/dist-packages (from google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (0.2.1)\n",
      "Collecting rsa<5,>=3.1.4\n",
      "  Downloading rsa-4.7.2-py3-none-any.whl (34 kB)\n",
      "Collecting cachetools<5.0,>=2.0.0\n",
      "  Downloading cachetools-4.2.2-py3-none-any.whl (11 kB)\n",
      "Collecting google-api-core<2.0.0dev,>=1.21.0\n",
      "  Downloading google_api_core-1.30.0-py2.py3-none-any.whl (93 kB)\n",
      "\u001b[K     |████████████████████████████████| 93 kB 791 kB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting googleapis-common-protos<2.0dev,>=1.6.0\n",
      "  Downloading googleapis_common_protos-1.53.0-py2.py3-none-any.whl (198 kB)\n",
      "\u001b[K     |████████████████████████████████| 198 kB 119.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: pytz in /home/ubuntu/.local/lib/python3.8/site-packages (from google-api-core<2.0.0dev,>=1.21.0->google-cloud-core<2.0dev,>=1.4.1->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (2021.1)\n",
      "Collecting google-crc32c<2.0dev,>=1.0\n",
      "  Downloading google_crc32c-1.1.2-cp38-cp38-manylinux2014_x86_64.whl (38 kB)\n",
      "Requirement already satisfied: cffi>=1.0.0 in /home/ubuntu/.local/lib/python3.8/site-packages (from google-crc32c<2.0dev,>=1.0->google-resumable-media<2.0dev,>=1.2.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (1.14.5)\n",
      "Requirement already satisfied: pycparser in /home/ubuntu/.local/lib/python3.8/site-packages (from cffi>=1.0.0->google-crc32c<2.0dev,>=1.0->google-resumable-media<2.0dev,>=1.2.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (2.20)\n",
      "Collecting scipy\n",
      "  Downloading scipy-1.6.3-cp38-cp38-manylinux1_x86_64.whl (27.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 27.2 MB 107.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: pyparsing>=2.0.2 in /home/ubuntu/.local/lib/python3.8/site-packages (from packaging->pytest==6.2.3->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (2.4.7)\n",
      "Collecting bcrypt>=3.1.3\n",
      "  Downloading bcrypt-3.2.0-cp36-abi3-manylinux2010_x86_64.whl (63 kB)\n",
      "\u001b[K     |████████████████████████████████| 63 kB 4.9 MB/s s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: cryptography>=2.5 in /usr/lib/python3/dist-packages (from paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 8)) (2.8)\n",
      "Requirement already satisfied: pynacl>=1.0.1 in /usr/lib/python3/dist-packages (from paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 8)) (1.3.0)\n",
      "Requirement already satisfied: pyasn1>=0.1.3 in /usr/lib/python3/dist-packages (from rsa<5,>=3.1.4->google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (0.4.2)\n",
      "Collecting threadpoolctl>=2.0.0\n",
      "  Downloading threadpoolctl-2.1.0-py3-none-any.whl (12 kB)\n",
      "Collecting joblib>=0.11\n",
      "  Downloading joblib-1.0.1-py3-none-any.whl (303 kB)\n",
      "\u001b[K     |████████████████████████████████| 303 kB 123.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting markdown>=2.6.8\n",
      "  Downloading Markdown-3.3.4-py3-none-any.whl (97 kB)\n",
      "\u001b[K     |████████████████████████████████| 97 kB 15.3 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting google-auth-oauthlib<0.5,>=0.4.1\n",
      "  Downloading google_auth_oauthlib-0.4.4-py2.py3-none-any.whl (18 kB)\n",
      "Collecting tensorboard-plugin-wit>=1.6.0\n",
      "  Downloading tensorboard_plugin_wit-1.8.0-py3-none-any.whl (781 kB)\n",
      "\u001b[K     |████████████████████████████████| 781 kB 115.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tensorboard-data-server<0.7.0,>=0.6.0\n",
      "  Downloading tensorboard_data_server-0.6.1-py3-none-manylinux2010_x86_64.whl (4.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.9 MB 115.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting requests-oauthlib>=0.7.0\n",
      "  Downloading requests_oauthlib-1.3.0-py2.py3-none-any.whl (23 kB)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /usr/lib/python3/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (3.1.0)\n",
      "Requirement already satisfied: wcwidth in /home/ubuntu/.local/lib/python3.8/site-packages (from ftfy->-r mesh-transformer-jax/requirements.txt (line 20)) (0.2.5)\n",
      "Collecting async-timeout<4.0,>=3.0\n",
      "  Downloading async_timeout-3.0.1-py3-none-any.whl (8.2 kB)\n",
      "Collecting multidict<7.0,>=4.5\n",
      "  Downloading multidict-5.1.0-cp38-cp38-manylinux2014_x86_64.whl (159 kB)\n",
      "\u001b[K     |████████████████████████████████| 159 kB 118.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting yarl<2.0,>=1.0\n",
      "  Downloading yarl-1.6.3-cp38-cp38-manylinux2014_x86_64.whl (324 kB)\n",
      "\u001b[K     |████████████████████████████████| 324 kB 124.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting hiredis\n",
      "  Downloading hiredis-2.0.0-cp38-cp38-manylinux2010_x86_64.whl (85 kB)\n",
      "\u001b[K     |████████████████████████████████| 85 kB 9.8 MB/s s eta 0:00:01\n",
      "\u001b[?25hCollecting blessings>=1.6\n",
      "  Downloading blessings-1.7-py3-none-any.whl (18 kB)\n",
      "Collecting nvidia-ml-py3>=7.352.0\n",
      "  Downloading nvidia-ml-py3-7.352.0.tar.gz (19 kB)\n",
      "Collecting opencensus-context==0.1.2\n",
      "  Downloading opencensus_context-0.1.2-py2.py3-none-any.whl (4.4 kB)\n",
      "Building wheels for collected packages: dm-haiku, lm-eval-harness, black, openai, pycountry, sqlitedict, jax, promise, subprocess32, termcolor, wrapt, func-timeout, ftfy, gpustat, nvidia-ml-py3, pathtools\n",
      "  Building wheel for dm-haiku (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for dm-haiku: filename=dm_haiku-0.0.5.dev0-py3-none-any.whl size=296288 sha256=fc659c1f4209f8db5b50b7db9fe6b4c3c2428b32d3f8a5557f38d71801f24c5e\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-78jde1o0/wheels/c7/4d/89/b159f184ad7c9e95672c342eafcc176ad92ee0c77f27f3bd23\n",
      "  Building wheel for lm-eval-harness (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for lm-eval-harness: filename=lm_eval_harness-0.0.1-py3-none-any.whl size=96969 sha256=2e86f0d0843fad37ab63f4f17af4dc33748792015f578c31086a16a90d31e401\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-78jde1o0/wheels/f9/bc/f7/9a79ab0e8c58f8037034286f323bf021e90804b5dbe9d0bc46\n",
      "  Building wheel for black (PEP 517) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for black: filename=black-20.8b1-py3-none-any.whl size=124195 sha256=02335a475bf259353ba058f057fbe26470fccd6490cd3022f0490030cddd6350\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/95/a4/59/10cd5378d52f92cdb45025f040e4686e10ae5217961c25fd66\n",
      "  Building wheel for openai (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for openai: filename=openai-0.6.4-py3-none-any.whl size=172202 sha256=22dd3cd4da97ae2507e138b8efa0af30f3eb388d53167ca00d3f0195fec00a99\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/8f/40/fd/04d52b10d0484bba0acfb7b43ca07516ede10b657a3231d325\n",
      "  Building wheel for pycountry (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for pycountry: filename=pycountry-20.7.3-py2.py3-none-any.whl size=10746863 sha256=133d0f9bbc55b11cfda027c74f40e2092ec5239ce6a4b85297cc2fd6399554e0\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/09/eb/0d/4ee773c6a4aadc2a43cb5c1d07f268f13c4cdc0eec88e7c1ef\n",
      "  Building wheel for sqlitedict (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for sqlitedict: filename=sqlitedict-1.6.0-py3-none-any.whl size=14688 sha256=3f8260c6ef7e35324388144782ad350ff3616fb2888eff30ac95436bec2b9344\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/7a/a6/e9/1522a2a8f3f81d20faf3b4f25ea05572dc1dab9b8202b1367c\n",
      "  Building wheel for jax (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for jax: filename=jax-0.2.14-py3-none-any.whl size=771353 sha256=566c6369615ca66e11411ec78da916bcd8bb0cabf3295a2203d96a7f81d716ed\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/93/a0/85/5a45d2ecf77637337831631269deea8cd7b45d6960f046b979\n",
      "  Building wheel for promise (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for promise: filename=promise-2.3-py3-none-any.whl size=21493 sha256=1ec094b1742b4a99451280f203ca8a8897e7d118e207c9091b668f293da095aa\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/54/aa/01/724885182f93150035a2a91bce34a12877e8067a97baaf5dc8\n",
      "  Building wheel for subprocess32 (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for subprocess32: filename=subprocess32-3.5.4-py3-none-any.whl size=6489 sha256=a464885ab98158c2f492af2c137459932d94944abd5933c21b1e3b24effde64e\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/9f/69/d1/50b39b308a87998eaf5c1d9095e5a5bd2ad98501e2b7936d36\n",
      "  Building wheel for termcolor (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for termcolor: filename=termcolor-1.1.0-py3-none-any.whl size=4830 sha256=0a935896f1250baea7f6d7d7d0f248b65120b3a8a559a2375bcfe1d7314d4781\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/a0/16/9c/5473df82468f958445479c59e784896fa24f4a5fc024b0f501\n",
      "  Building wheel for wrapt (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for wrapt: filename=wrapt-1.12.1-cp38-cp38-linux_x86_64.whl size=78556 sha256=a10d61246b2429fdcc59583d0b59ac099fba77f2e880b3c354b69cc4f20d0cb3\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/5f/fd/9e/b6cf5890494cb8ef0b5eaff72e5d55a70fb56316007d6dfe73\n",
      "  Building wheel for func-timeout (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for func-timeout: filename=func_timeout-4.3.5-py3-none-any.whl size=15077 sha256=2b1d18fe5b536b32be3e5c97816e8a660ec91ff2269500fdc65a9e09d0f71057\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/a8/92/ca/5bbab358275e310af23b73fc32ebf37d6a7a08c87c8d2cdbc1\n",
      "  Building wheel for ftfy (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for ftfy: filename=ftfy-6.0.3-py3-none-any.whl size=41913 sha256=e9d9b1cbddb278df1038a39d1c53908169799b1f20508335b6c59c733e1aeb87\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/7f/40/63/4bf603cec3ecc4a26985405834cb47eb8368bfa59e15dde046\n",
      "  Building wheel for gpustat (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for gpustat: filename=gpustat-0.6.0-py3-none-any.whl size=12617 sha256=2d3eb4c7d0a2cdcc11eb2d877b9fec6cf6612b1ccb58a0b799af151ca57cf9f9\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/0d/d9/80/b6cbcdc9946c7b50ce35441cc9e7d8c5a9d066469ba99bae44\n",
      "  Building wheel for nvidia-ml-py3 (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for nvidia-ml-py3: filename=nvidia_ml_py3-7.352.0-py3-none-any.whl size=19189 sha256=a37540d1824983713d5b08f51ea6feb7888f2c577b06c5b3a35087ebf9ef2803\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/b9/b1/68/cb4feab29709d4155310d29a421389665dcab9eb3b679b527b\n",
      "  Building wheel for pathtools (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for pathtools: filename=pathtools-0.1.2-py3-none-any.whl size=8784 sha256=c1245451e141fcbeefa592652231484181c2d7f24c966864d92748deb6616d1b\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/4c/8e/7e/72fbc243e1aeecae64a96875432e70d4e92f3d2d18123be004\n",
      "Successfully built dm-haiku lm-eval-harness black openai pycountry sqlitedict jax promise subprocess32 termcolor wrapt func-timeout ftfy gpustat nvidia-ml-py3 pathtools\n",
      "Installing collected packages: six, mbstrdecoder, typepy, setuptools, rsa, protobuf, cachetools, requests, numpy, multidict, googleapis-common-protos, google-auth, yarl, typing-extensions, tqdm, smmap, scipy, requests-oauthlib, regex, opt-einsum, joblib, google-crc32c, google-api-core, flatbuffers, filelock, dill, DataProperty, click, async-timeout, absl-py, zstandard, xxhash, wheel, Werkzeug, ujson, typed-ast, toolz, toml, tokenizers, threadpoolctl, tensorboard-plugin-wit, tensorboard-data-server, tcolorpy, tabledata, sacremoses, rehash, pyarrow, py, psutil, portalocker, pluggy, pathvalidate, pathspec, pandas, opencensus-context, nvidia-ml-py3, mypy-extensions, multiprocess, msgfy, markdown, jsonlines, jaxlib, jax, iniconfig, huggingface-hub, hiredis, grpcio, google-resumable-media, google-cloud-core, google-auth-oauthlib, gitdb, fsspec, dm-tree, blessings, bcrypt, appdirs, aiohttp, wrapt, transformers, tqdm-multiprocess, termcolor, tensorflow-estimator, tensorboard, tabulate, subprocess32, sqlitedict, smart-open, shortuuid, sentry-sdk, scikit-learn, sacrebleu, redis, pytest, pytablewriter, pycountry, pybind11, py-spy, promise, pathtools, pathlib2, paramiko, opencensus, openai, numexpr, msgpack, mock, lm-dataformat, keras-preprocessing, jmp, itsdangerous, invoke, h5py, gpustat, google-pasta, google-cloud-storage, GitPython, gast, docker-pycreds, datasets, configparser, colorful, chex, black, best-download, astunparse, aioredis, aiohttp-cors, wandb, tensorflow-cpu, ray, optax, lm-eval-harness, func-timeout, ftfy, Flask, fabric, einops, dm-haiku, cloudpickle\n",
      "  Attempting uninstall: typing-extensions\n",
      "    Found existing installation: typing-extensions 3.10.0.0\n",
      "    Uninstalling typing-extensions-3.10.0.0:\n",
      "      Successfully uninstalled typing-extensions-3.10.0.0\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "launchpadlib 1.10.13 requires testresources, which is not installed.\u001b[0m\n",
      "Successfully installed DataProperty-0.50.1 Flask-1.1.4 GitPython-3.1.14 Werkzeug-1.0.1 absl-py-0.13.0 aiohttp-3.7.4.post0 aiohttp-cors-0.7.0 aioredis-1.3.1 appdirs-1.4.4 astunparse-1.6.3 async-timeout-3.0.1 bcrypt-3.2.0 best-download-0.0.6 black-20.8b1 blessings-1.7 cachetools-4.2.2 chex-0.0.7 click-7.1.2 cloudpickle-1.3.0 colorful-0.5.4 configparser-5.0.2 datasets-1.8.0 dill-0.3.4 dm-haiku-0.0.5.dev0 dm-tree-0.1.6 docker-pycreds-0.4.0 einops-0.3.0 fabric-2.6.0 filelock-3.0.12 flatbuffers-1.12 fsspec-2021.6.0 ftfy-6.0.3 func-timeout-4.3.5 gast-0.3.3 gitdb-4.0.7 google-api-core-1.30.0 google-auth-1.31.0 google-auth-oauthlib-0.4.4 google-cloud-core-1.7.0 google-cloud-storage-1.36.2 google-crc32c-1.1.2 google-pasta-0.2.0 google-resumable-media-1.3.0 googleapis-common-protos-1.53.0 gpustat-0.6.0 grpcio-1.32.0 h5py-2.10.0 hiredis-2.0.0 huggingface-hub-0.0.10 iniconfig-1.1.1 invoke-1.5.0 itsdangerous-1.1.0 jax-0.2.14 jaxlib-0.1.67 jmp-0.0.2 joblib-1.0.1 jsonlines-2.0.0 keras-preprocessing-1.1.2 lm-dataformat-0.0.19 lm-eval-harness-0.0.1 markdown-3.3.4 mbstrdecoder-1.0.1 mock-4.0.3 msgfy-0.1.0 msgpack-1.0.2 multidict-5.1.0 multiprocess-0.70.12.2 mypy-extensions-0.4.3 numexpr-2.7.2 numpy-1.19.5 nvidia-ml-py3-7.352.0 openai-0.6.4 opencensus-0.7.13 opencensus-context-0.1.2 opt-einsum-3.3.0 optax-0.0.6 pandas-1.2.4 paramiko-2.7.2 pathlib2-2.3.5 pathspec-0.8.1 pathtools-0.1.2 pathvalidate-2.4.1 pluggy-0.13.1 portalocker-2.3.0 promise-2.3 protobuf-3.17.3 psutil-5.8.0 py-1.10.0 py-spy-0.3.7 pyarrow-3.0.0 pybind11-2.6.2 pycountry-20.7.3 pytablewriter-0.58.0 pytest-6.2.3 ray-1.2.0 redis-3.5.3 regex-2021.4.4 rehash-1.0.0 requests-2.25.1 requests-oauthlib-1.3.0 rsa-4.7.2 sacrebleu-1.5.0 sacremoses-0.0.45 scikit-learn-0.24.2 scipy-1.6.3 sentry-sdk-1.1.0 setuptools-51.3.3 shortuuid-1.0.1 six-1.15.0 smart-open-5.1.0 smmap-4.0.0 sqlitedict-1.6.0 subprocess32-3.5.4 tabledata-1.1.4 tabulate-0.8.9 tcolorpy-0.1.0 tensorboard-2.5.0 tensorboard-data-server-0.6.1 tensorboard-plugin-wit-1.8.0 tensorflow-cpu-2.4.2 tensorflow-estimator-2.4.0 termcolor-1.1.0 threadpoolctl-2.1.0 tokenizers-0.10.3 toml-0.10.2 toolz-0.11.1 tqdm-4.45.0 tqdm-multiprocess-0.0.11 transformers-4.4.2 typed-ast-1.4.3 typepy-1.1.5 typing-extensions-3.7.4.3 ujson-4.0.2 wandb-0.10.32 wheel-0.36.2 wrapt-1.12.1 xxhash-2.0.2 yarl-1.6.3 zstandard-0.15.2\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html\n",
      "Processing ./mesh-transformer-jax\n",
      "\u001b[33m  DEPRECATION: A future pip version will change local packages to be built in-place without first copying to a temporary directory. We recommend you use --use-feature=in-tree-build to test your packages with this new behavior before it becomes the default.\n",
      "   pip 21.3 will remove support for this functionality. You can find discussion regarding this at https://github.com/pypa/pip/issues/7555.\u001b[0m\n",
      "Collecting jax==0.2.12\n",
      "  Downloading jax-0.2.12.tar.gz (590 kB)\n",
      "\u001b[K     |████████████████████████████████| 590 kB 23.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting jaxlib==0.1.67+cuda111\n",
      "  Downloading https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.67%2Bcuda111-cp38-none-manylinux2010_x86_64.whl (194.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 194.6 MB 135.7 MB/s eta 0:00:01     |███████████████████████████▊    | 168.8 MB 148.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.12 in /home/ubuntu/.local/lib/python3.8/site-packages (from jax==0.2.12) (1.19.5)\n",
      "Requirement already satisfied: absl-py in /home/ubuntu/.local/lib/python3.8/site-packages (from jax==0.2.12) (0.13.0)\n",
      "Requirement already satisfied: opt_einsum in /home/ubuntu/.local/lib/python3.8/site-packages (from jax==0.2.12) (3.3.0)\n",
      "Requirement already satisfied: scipy in /home/ubuntu/.local/lib/python3.8/site-packages (from jaxlib==0.1.67+cuda111) (1.6.3)\n",
      "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /home/ubuntu/.local/lib/python3.8/site-packages (from jaxlib==0.1.67+cuda111) (1.12)\n",
      "Requirement already satisfied: six in /home/ubuntu/.local/lib/python3.8/site-packages (from absl-py->jax==0.2.12) (1.15.0)\n",
      "Building wheels for collected packages: mesh-transformer, jax\n",
      "  Building wheel for mesh-transformer (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for mesh-transformer: filename=mesh_transformer-0.0.0-py3-none-any.whl size=20007 sha256=a5df3168e2535c5ab1d1aa969541b58cf26a4debe3e4a2a1159667f034e87442\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/d0/b9/ce/d0870b31fde1e4231775e930342b0a28404a48f4b571d16c9d\n",
      "  Building wheel for jax (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for jax: filename=jax-0.2.12-py3-none-any.whl size=682484 sha256=7537069264670941fb3ddce1ad9a3b4ad06142a30af3c79dac0675aaa248a7ef\n",
      "  Stored in directory: /home/ubuntu/.cache/pip/wheels/13/c3/89/d192f2a08b9d115f694d0a1e9a334f9516fdaa6cf6e74923cd\n",
      "Successfully built mesh-transformer jax\n",
      "Installing collected packages: mesh-transformer, jaxlib, jax\n",
      "  Attempting uninstall: jaxlib\n",
      "    Found existing installation: jaxlib 0.1.67\n",
      "    Uninstalling jaxlib-0.1.67:\n",
      "      Successfully uninstalled jaxlib-0.1.67\n",
      "  Attempting uninstall: jax\n",
      "    Found existing installation: jax 0.2.14\n",
      "    Uninstalling jax-0.2.14:\n",
      "      Successfully uninstalled jax-0.2.14\n",
      "Successfully installed jax-0.2.12 jaxlib-0.1.67+cuda111 mesh-transformer-0.0.0\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting git+https://github.com/finetuneanon/transformers@gpt-j\n",
      "  Cloning https://github.com/finetuneanon/transformers (to revision gpt-j) to /tmp/pip-req-build-zii4tbn0\n",
      "  Running command git clone -q https://github.com/finetuneanon/transformers /tmp/pip-req-build-zii4tbn0\n",
      "  Running command git checkout -b gpt-j --track origin/gpt-j\n",
      "  Switched to a new branch 'gpt-j'\n",
      "  Branch 'gpt-j' set up to track remote branch 'gpt-j' from 'origin'.\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h    Preparing wheel metadata ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: regex!=2019.12.17 in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (2021.4.4)\n",
      "Requirement already satisfied: filelock in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (3.0.12)\n",
      "Collecting huggingface-hub==0.0.8\n",
      "  Downloading huggingface_hub-0.0.8-py3-none-any.whl (34 kB)\n",
      "Requirement already satisfied: sacremoses in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (0.0.45)\n",
      "Requirement already satisfied: tqdm>=4.27 in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (4.45.0)\n",
      "Requirement already satisfied: requests in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (2.25.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (1.19.5)\n",
      "Requirement already satisfied: einops==0.3.0 in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (0.3.0)\n",
      "Requirement already satisfied: tokenizers<0.11,>=0.10.1 in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (0.10.3)\n",
      "Requirement already satisfied: packaging in /home/ubuntu/.local/lib/python3.8/site-packages (from transformers==4.7.0.dev0) (20.9)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /home/ubuntu/.local/lib/python3.8/site-packages (from packaging->transformers==4.7.0.dev0) (2.4.7)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/lib/python3/dist-packages (from requests->transformers==4.7.0.dev0) (2.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->transformers==4.7.0.dev0) (2019.11.28)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/lib/python3/dist-packages (from requests->transformers==4.7.0.dev0) (3.0.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/lib/python3/dist-packages (from requests->transformers==4.7.0.dev0) (1.25.8)\n",
      "Requirement already satisfied: click in /home/ubuntu/.local/lib/python3.8/site-packages (from sacremoses->transformers==4.7.0.dev0) (7.1.2)\n",
      "Requirement already satisfied: six in /home/ubuntu/.local/lib/python3.8/site-packages (from sacremoses->transformers==4.7.0.dev0) (1.15.0)\n",
      "Requirement already satisfied: joblib in /home/ubuntu/.local/lib/python3.8/site-packages (from sacremoses->transformers==4.7.0.dev0) (1.0.1)\n",
      "Building wheels for collected packages: transformers\n",
      "  Building wheel for transformers (PEP 517) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for transformers: filename=transformers-4.7.0.dev0-py3-none-any.whl size=2405405 sha256=1a5665ce7e1bd571545584e8322cc0ca2a7caefb1e4f2ee9312e4edb8eadf0e4\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-3uo9l5xt/wheels/de/c8/69/3053dc726555b5d24bc6498573032f7193ac034e22925b94db\n",
      "Successfully built transformers\n",
      "Installing collected packages: huggingface-hub, transformers\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.0.10\n",
      "    Uninstalling huggingface-hub-0.0.10:\n",
      "      Successfully uninstalled huggingface-hub-0.0.10\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.4.2\n",
      "    Uninstalling transformers-4.4.2:\n",
      "      Successfully uninstalled transformers-4.4.2\n",
      "Successfully installed huggingface-hub-0.0.8 transformers-4.7.0.dev0\n"
     ]
    }
   ],
   "source": [
    "!pip3 install torch \n",
    "\n",
    "!git clone https://github.com/kingoflolz/mesh-transformer-jax.git\n",
    "!pip3 install -r mesh-transformer-jax/requirements.txt\n",
    "\n",
    "# jax 0.2.12 is required due to a regression with xmap in 0.2.13\n",
    "!pip3 install mesh-transformer-jax/ jax==0.2.12 jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
    "!pip3 install git+https://github.com/finetuneanon/transformers@gpt-j"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c347a8e",
   "metadata": {},
   "source": [
    "## Get the actual GPT-J model [Source](https://github.com/kingoflolz/mesh-transformer-jax/blob/master/colab_demo.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "971f231c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo apt install zstd\n",
    "\n",
    "# the \"slim\" version contain only bf16 weights and no optimizer parameters, which minimizes bandwidth and memory\n",
    "!time wget -c https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd\n",
    "\n",
    "!time tar -I zstd -xf step_383500_slim.tar.zstd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d046b59b",
   "metadata": {},
   "source": [
    "## Convert to torch [Source](https://gist.github.com/finetuneanon/ee196c6cd16af1de4ca444862414683a) Credit:[finetuneanon](https://github.com/finetuneanon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8db4de",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir gpt-j-hf\n",
    "!curl https://gist.githubusercontent.com/finetuneanon/a55bdb3f5881e361faef0e96e1d41f09/raw/e5a38dad34ff42bbad188afd5e4fdb2ab2eacb6d/gpt-j-6b.json > gpt-j-hf/config.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a4e21c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "import io\n",
    "import os\n",
    "\n",
    "torch.set_printoptions(linewidth=130, sci_mode=False)\n",
    "np.set_printoptions(linewidth=130, suppress=True)\n",
    "\n",
    "layers = 28\n",
    "total_shards = 8\n",
    "ckpt_dir = \"step_383500/\"\n",
    "output_dir = \"j6b_ckpt\"\n",
    "\n",
    "def reshard(x, old_shape):\n",
    "    if len(x.shape) == 1:\n",
    "        # print(\"epoch\")\n",
    "        # print(x)\n",
    "        out = x[0:1]\n",
    "\n",
    "    elif len(x.shape) == 2:\n",
    "        #print(f\"LN/bias {x.shape}\")\n",
    "        #print(x[:, :16])\n",
    "\n",
    "        if (x[1:] == x[-1]).all():\n",
    "            #print(\"LN\")\n",
    "            if (x[1:] == 0).all() or (x[1:] == 1).all():\n",
    "                out = x[0:1]\n",
    "            else:\n",
    "                #print(\"shard bias\")\n",
    "                out = x[0:1] * 8#* x.shape[0] / old_shape[0]\n",
    "        else:\n",
    "            #print(\"bias\")\n",
    "            out = x.reshape(old_shape)\n",
    "\n",
    "        #print(out[:, :16])\n",
    "\n",
    "    elif len(x.shape) == 3:\n",
    "        #print(f\"weight {x.shape}\")\n",
    "        if x.shape[0] * x.shape[2] == old_shape[2]:\n",
    "            #print(\"case 1\")\n",
    "            out = jnp.transpose(x, (1, 0, 2)).reshape(old_shape)\n",
    "        elif x.shape[0] * x.shape[1] == old_shape[1]:\n",
    "            #print(\"case 2\")\n",
    "            out = x.reshape(old_shape)\n",
    "        else:\n",
    "            raise Exception(f\"unimplemented, {x.shape}, {old_shape}\")\n",
    "    else:\n",
    "        raise Exception(f\"unimplemented, {x}\")\n",
    "    #flattened, structure = jax.tree_flatten(out)\n",
    "    #return flattened\n",
    "    return out\n",
    "\n",
    "def get_old_shape(t, dim=2):\n",
    "    if len(t.shape) == 3:\n",
    "        shard_shape = t.shape\n",
    "        if dim == 1:\n",
    "            return (shard_shape[0] * shard_shape[1], shard_shape[2])\n",
    "        elif dim == 2:\n",
    "            return (shard_shape[1], shard_shape[0] * shard_shape[2])\n",
    "        else:\n",
    "            raise ValueError(f\"unsupported dim {dim}\")\n",
    "    if len(t.shape) == 2:\n",
    "        return (t.shape[1] * t.shape[0],)\n",
    "    else:\n",
    "        raise ValueError(f\"unsupported shape {t.shape}\")\n",
    "\n",
    "def read_shard(ckpt_dir):\n",
    "    global part\n",
    "    out = []\n",
    "    idx = part\n",
    "    file_path = ckpt_dir + f\"{idx}.npz\"\n",
    "    #print(f\"-- {file_path}\")\n",
    "    with open(file_path, \"rb\") as f:\n",
    "        buf = f.read()\n",
    "        f_io = io.BytesIO(buf)\n",
    "        deserialized = np.load(f_io)\n",
    "        for i in deserialized:\n",
    "            out.append(deserialized[i])\n",
    "            #print(deserialized[i].shape)\n",
    "    return out\n",
    "\n",
    "def save(ckpt):\n",
    "    try: os.mkdir(output_dir)\n",
    "    except: pass\n",
    "    checkpoint = {}\n",
    "    for i, x in enumerate(ckpt.items()):\n",
    "        checkpoint[x[0]] = f\"{output_dir}/b{i}.pt\"\n",
    "        torch.save(x[1], f\"{output_dir}/b{i}.pt\")\n",
    "    torch.save(checkpoint, f\"{output_dir}/m.pt\")\n",
    "\n",
    "unshard = None\n",
    "transforms = [(\"transformer.wte.bias\", None, None), (\"transformer.wte.weight\", unshard, 1)]\n",
    "\n",
    "checkpoint = {}\n",
    "\n",
    "layer_names = sorted(map(str, range(layers)))\n",
    "for layer in layer_names:\n",
    "    checkpoint[f\"transformer.h.{layer}.attn.attention.bias\"] = torch.tril(torch.ones(1, 1, 2048, 2048))\n",
    "    checkpoint[f\"transformer.h.{layer}.attn.attention.masked_bias\"] = torch.tensor(-1e9)\n",
    "    transforms.extend([\n",
    "        (f\"transformer.h.{layer}.attn.attention.q_proj.weight\", unshard, 2),\n",
    "        (f\"transformer.h.{layer}.attn.attention.v_proj.weight\", unshard, 2),\n",
    "        (f\"transformer.h.{layer}.attn.attention.k_proj.weight\", unshard, 2),\n",
    "        (f\"transformer.h.{layer}.attn.attention.out_proj.weight\", unshard, 1),\n",
    "        (f\"transformer.h.{layer}.mlp.c_fc.bias\", unshard, 1),\n",
    "        (f\"transformer.h.{layer}.mlp.c_fc.weight\", unshard, 2),\n",
    "        (f\"transformer.h.{layer}.mlp.c_proj.bias\", None, None),\n",
    "        (f\"transformer.h.{layer}.mlp.c_proj.weight\", unshard, 1),\n",
    "        (f\"transformer.h.{layer}.ln_1.bias\", None, None),\n",
    "        (f\"transformer.h.{layer}.ln_1.weight\", None, None),\n",
    "    ])\n",
    "transforms.extend([\n",
    "    (\"lm_head.bias\", unshard, 1),\n",
    "    (\"lm_head.weight\", unshard, 2),\n",
    "    (\"transformer.ln_f.bias\", None, None),\n",
    "    (\"transformer.ln_f.weight\", None, None),\n",
    "])\n",
    "\n",
    "part = 0\n",
    "element = 0\n",
    "while len(transforms) > 0:\n",
    "    print(f\"loading shards for part {part}\")\n",
    "    shards = list(map(read_shard, [f\"{ckpt_dir}shard_{i}/\" for i in range(total_shards)]))\n",
    "    print(f\"read from checkpoint\")\n",
    "\n",
    "    unsharded = []\n",
    "\n",
    "    for all_shards in zip(*shards):\n",
    "        x = np.stack(all_shards)\n",
    "        # No idea why this is V2...?\n",
    "        if x.dtype == np.dtype('V2'):\n",
    "            x.dtype = jnp.bfloat16\n",
    "        x = x.astype(np.float32)\n",
    "        unsharded.append(x)\n",
    "        #print(f\"unsharded: {x.shape}\")\n",
    "\n",
    "    while len(transforms) > 0 and len(unsharded) > 0:\n",
    "        transform = transforms.pop(0)\n",
    "        params = unsharded.pop(0)\n",
    "        if transform[2] is not None:\n",
    "            old_shape = (1,) + get_old_shape(params, transform[2])\n",
    "        else:\n",
    "            old_shape = (params.shape[1],)\n",
    "        print(f\"< {params.shape} to {old_shape}\")\n",
    "        params = reshard(params, old_shape).squeeze(0).T\n",
    "        params = torch.tensor(params.copy()).half()\n",
    "        if params.isnan().any() or params.isinf().any():\n",
    "            raise ValueError(f\"fp16 over/underflow at {part} {element}\")\n",
    "        checkpoint[transform[0]] = params\n",
    "        print(f\"> {transform[0]} {params.shape}\")\n",
    "        element += 1\n",
    "    part += 1\n",
    "\n",
    "checkpoint['transformer.wte.weight'] = (checkpoint['transformer.wte.weight'].T + checkpoint['transformer.wte.bias'])\n",
    "del checkpoint['transformer.wte.bias']\n",
    "\n",
    "print(f\"left over: {unsharded}\")\n",
    "print(\"saving\")\n",
    "torch.save(checkpoint, \"./gpt-j-hf/pytorch_model.bin\") # load as in: https://github.com/finetuneanon/misc/blob/main/SizeTest.ipynb\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "90f8bd0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "from transformers import GPTNeoForCausalLM, AutoConfig, GPT2Tokenizer\n",
    "import torch\n",
    "import transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b8118113",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-3-4fcea1767d85>:4: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3, and in 3.9 it will stop working\n",
      "  class Checkpoint(collections.MutableMapping):\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded\n",
      "ok\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "GPTNeoForCausalLM(\n",
       "  (transformer): GPTNeoModel(\n",
       "    (wte): Embedding(50400, 4096)\n",
       "    (drop): Dropout(p=0, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (7): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (8): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (9): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (10): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (11): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (12): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (13): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (14): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (15): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (16): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (17): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (18): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (19): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (20): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (21): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (22): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (23): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (24): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (25): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (26): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (27): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=50400, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import collections\n",
    "import os\n",
    "import json\n",
    "\n",
    "class Checkpoint(collections.MutableMapping):\n",
    "    def __init__(self):\n",
    "        self.checkpoint = torch.load(\"pytorch_model.bin\")\n",
    "        print(\"Loaded\")\n",
    "    def __len__(self):\n",
    "        return len(self.checkpoint)\n",
    "    def __getitem__(self, key):\n",
    "        return torch.load(self.checkpoint[key])\n",
    "    def __setitem__(self, key, value):\n",
    "        return\n",
    "    def __delitem__(self, key, value):\n",
    "        return\n",
    "    def keys(self):\n",
    "        return self.checkpoint.keys()\n",
    "    def __iter__(self):\n",
    "        for key in self.checkpoint:\n",
    "            yield (key, self.__getitem__(key))\n",
    "    def __copy__(self):\n",
    "        return self.__dict__\n",
    "    def copy(self):\n",
    "        return self.__dict__\n",
    "\n",
    "print(\"load\", flush=True)\n",
    "with open('config.json', 'r') as f:\n",
    "    config = json.load(f)\n",
    "\n",
    "model = GPTNeoForCausalLM.from_pretrained(pretrained_model_name_or_path=None, config=config, state_dict=Checkpoint())\n",
    "print(\"ok\")\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a6475e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GPTNeoForCausalLM.from_pretrained(\"./gpt-j-hf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f75f0a4",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0f731069",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = transformers.GPT2Tokenizer.from_pretrained('gpt2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76da2f0c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTNeoForCausalLM(\n",
       "  (transformer): GPTNeoModel(\n",
       "    (wte): Embedding(50400, 4096)\n",
       "    (drop): Dropout(p=0, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (7): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (8): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (9): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (10): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (11): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (12): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (13): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (14): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (15): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (16): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (17): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (18): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (19): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (20): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (21): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (22): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (23): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (24): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (25): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (26): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (27): GPTNeoBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTNeoAttention(\n",
       "          (attention): GPTNeoSelfAttention(\n",
       "            (attn_dropout): Dropout(p=0, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0, inplace=False)\n",
       "            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "            (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          )\n",
       "        )\n",
       "        (mlp): GPTNeoMLP(\n",
       "          (c_fc): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (c_proj): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=50400, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.half().cuda() # This should take about 12GB of Graphics RAM, if you have a larger than 16GB gpu you don't need the half()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "030d4615",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "input_text = \"Hello my name is Paul and\"\n",
    "input_ids = tokenizer.encode(str(input_text), return_tensors='pt').cuda()\n",
    "\n",
    "output = model.generate(\n",
    "    input_ids,\n",
    "    do_sample=True,\n",
    "    max_length=20,\n",
    "    top_p=0.7,\n",
    "    top_k=0,\n",
    "    temperature=1.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f783672d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello my name is Paul and\n",
      "this is a Quick Tutorial on how to\n",
      "make a Mac and\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode(output[0], skip_special_tokens=True))"
   ]
  }
 ],
 "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
}