05_finetune/01_fine_tune_dolly_llama2_huggingface.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a9f1ed1-db76-4ecb-a96b-c425c68db9c9",
   "metadata": {},
   "source": [
    "# Requires p4d/p4de"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "07ad2e56-bc79-45a9-8017-272272e1eaeb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Collecting transformers==4.31.0\n",
      "  Obtaining dependency information for transformers==4.31.0 from https://files.pythonhosted.org/packages/21/02/ae8e595f45b6c8edee07913892b3b41f5f5f273962ad98851dc6a564bbb9/transformers-4.31.0-py3-none-any.whl.metadata\n",
      "  Downloading transformers-4.31.0-py3-none-any.whl.metadata (116 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.9/116.9 kB\u001b[0m \u001b[31m74.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from transformers==4.31.0) (3.6.0)\n",
      "Collecting huggingface-hub<1.0,>=0.14.1 (from transformers==4.31.0)\n",
      "  Obtaining dependency information for huggingface-hub<1.0,>=0.14.1 from https://files.pythonhosted.org/packages/7f/c4/adcbe9a696c135578cabcbdd7331332daad4d49b7c43688bc2d36b3a47d2/huggingface_hub-0.16.4-py3-none-any.whl.metadata\n",
      "  Downloading huggingface_hub-0.16.4-py3-none-any.whl.metadata (12 kB)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from transformers==4.31.0) (1.25.1)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from transformers==4.31.0) (21.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.10/site-packages/PyYAML-6.0-py3.10-linux-x86_64.egg (from transformers==4.31.0) (6.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.10/site-packages (from transformers==4.31.0) (2022.7.9)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from transformers==4.31.0) (2.31.0)\n",
      "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1 (from transformers==4.31.0)\n",
      "  Downloading tokenizers-0.13.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m149.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting safetensors>=0.3.1 (from transformers==4.31.0)\n",
      "  Obtaining dependency information for safetensors>=0.3.1 from https://files.pythonhosted.org/packages/6c/f0/c17bbdb1e5f9dab29d44cade445135789f75f8f08ea2728d04493ea8412b/safetensors-0.3.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading safetensors-0.3.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.7 kB)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/lib/python3.10/site-packages (from transformers==4.31.0) (4.64.1)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.14.1->transformers==4.31.0) (2022.7.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.14.1->transformers==4.31.0) (4.3.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->transformers==4.31.0) (3.0.9)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.31.0) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.31.0) (3.3)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.31.0) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->transformers==4.31.0) (2023.7.22)\n",
      "Downloading transformers-4.31.0-py3-none-any.whl (7.4 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.4/7.4 MB\u001b[0m \u001b[31m250.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading huggingface_hub-0.16.4-py3-none-any.whl (268 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m447.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading safetensors-0.3.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m459.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: tokenizers, safetensors, huggingface-hub, transformers\n",
      "Successfully installed huggingface-hub-0.16.4 safetensors-0.3.3 tokenizers-0.13.3 transformers-4.31.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Collecting peft==0.4.0\n",
      "  Obtaining dependency information for peft==0.4.0 from https://files.pythonhosted.org/packages/88/a0/6e1c23293a922a9c9e9bd8d56a60cd78ecf531fdabe45ac975e142bfbe86/peft-0.4.0-py3-none-any.whl.metadata\n",
      "  Downloading peft-0.4.0-py3-none-any.whl.metadata (21 kB)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from peft==0.4.0) (1.25.1)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from peft==0.4.0) (21.3)\n",
      "Requirement already satisfied: psutil in /opt/conda/lib/python3.10/site-packages (from peft==0.4.0) (5.9.0)\n",
      "Requirement already satisfied: pyyaml in /opt/conda/lib/python3.10/site-packages/PyYAML-6.0-py3.10-linux-x86_64.egg (from peft==0.4.0) (6.0)\n",
      "Collecting torch>=1.13.0 (from peft==0.4.0)\n",
      "  Downloading torch-2.0.1-cp310-cp310-manylinux1_x86_64.whl (619.9 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m619.9/619.9 MB\u001b[0m \u001b[31m151.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: transformers in /opt/conda/lib/python3.10/site-packages (from peft==0.4.0) (4.31.0)\n",
      "Collecting accelerate (from peft==0.4.0)\n",
      "  Obtaining dependency information for accelerate from https://files.pythonhosted.org/packages/4d/a7/05c67003d659a0035f2b3a8cf389c1d9645865aee84a73ce99ddab16682f/accelerate-0.22.0-py3-none-any.whl.metadata\n",
      "  Downloading accelerate-0.22.0-py3-none-any.whl.metadata (17 kB)\n",
      "Requirement already satisfied: safetensors in /opt/conda/lib/python3.10/site-packages (from peft==0.4.0) (0.3.3)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->peft==0.4.0) (3.0.9)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft==0.4.0) (3.6.0)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft==0.4.0) (4.3.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft==0.4.0) (1.10.1)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft==0.4.0) (2.8.4)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch>=1.13.0->peft==0.4.0) (3.1.2)\n",
      "Collecting nvidia-cuda-nvrtc-cu11==11.7.99 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cuda_nvrtc_cu11-11.7.99-2-py3-none-manylinux1_x86_64.whl (21.0 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.0/21.0 MB\u001b[0m \u001b[31m269.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cuda-runtime-cu11==11.7.99 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cuda_runtime_cu11-11.7.99-py3-none-manylinux1_x86_64.whl (849 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m849.3/849.3 kB\u001b[0m \u001b[31m500.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cuda-cupti-cu11==11.7.101 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cuda_cupti_cu11-11.7.101-py3-none-manylinux1_x86_64.whl (11.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.8/11.8 MB\u001b[0m \u001b[31m139.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cudnn-cu11==8.5.0.96 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cudnn_cu11-8.5.0.96-2-py3-none-manylinux1_x86_64.whl (557.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m557.1/557.1 MB\u001b[0m \u001b[31m158.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cublas-cu11==11.10.3.66 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cublas_cu11-11.10.3.66-py3-none-manylinux1_x86_64.whl (317.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m317.1/317.1 MB\u001b[0m \u001b[31m201.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cufft-cu11==10.9.0.58 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cufft_cu11-10.9.0.58-py3-none-manylinux1_x86_64.whl (168.4 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m168.4/168.4 MB\u001b[0m \u001b[31m217.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-curand-cu11==10.2.10.91 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_curand_cu11-10.2.10.91-py3-none-manylinux1_x86_64.whl (54.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.6/54.6 MB\u001b[0m \u001b[31m259.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cusolver-cu11==11.4.0.1 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cusolver_cu11-11.4.0.1-2-py3-none-manylinux1_x86_64.whl (102.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m102.6/102.6 MB\u001b[0m \u001b[31m283.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cusparse-cu11==11.7.4.91 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_cusparse_cu11-11.7.4.91-py3-none-manylinux1_x86_64.whl (173.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m173.2/173.2 MB\u001b[0m \u001b[31m181.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-nccl-cu11==2.14.3 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_nccl_cu11-2.14.3-py3-none-manylinux1_x86_64.whl (177.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m177.1/177.1 MB\u001b[0m \u001b[31m207.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-nvtx-cu11==11.7.91 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading nvidia_nvtx_cu11-11.7.91-py3-none-manylinux1_x86_64.whl (98 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m98.6/98.6 kB\u001b[0m \u001b[31m147.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting triton==2.0.0 (from torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (63.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m63.3/63.3 MB\u001b[0m \u001b[31m234.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: setuptools in /opt/conda/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.13.0->peft==0.4.0) (68.0.0)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.13.0->peft==0.4.0) (0.41.0)\n",
      "Collecting cmake (from triton==2.0.0->torch>=1.13.0->peft==0.4.0)\n",
      "  Obtaining dependency information for cmake from https://files.pythonhosted.org/packages/2e/51/3a4672a819b4532a378bfefad8f886cfe71057556e0d4eefb64523fd370a/cmake-3.27.2-py2.py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata\n",
      "  Downloading cmake-3.27.2-py2.py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (6.7 kB)\n",
      "Collecting lit (from triton==2.0.0->torch>=1.13.0->peft==0.4.0)\n",
      "  Downloading lit-16.0.6.tar.gz (153 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m153.7/153.7 kB\u001b[0m \u001b[31m413.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Installing backend dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: huggingface-hub<1.0,>=0.14.1 in /opt/conda/lib/python3.10/site-packages (from transformers->peft==0.4.0) (0.16.4)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.10/site-packages (from transformers->peft==0.4.0) (2022.7.9)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from transformers->peft==0.4.0) (2.31.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /opt/conda/lib/python3.10/site-packages (from transformers->peft==0.4.0) (0.13.3)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/lib/python3.10/site-packages (from transformers->peft==0.4.0) (4.64.1)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.14.1->transformers->peft==0.4.0) (2022.7.1)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch>=1.13.0->peft==0.4.0) (2.1.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->transformers->peft==0.4.0) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->transformers->peft==0.4.0) (3.3)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->transformers->peft==0.4.0) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->transformers->peft==0.4.0) (2023.7.22)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch>=1.13.0->peft==0.4.0) (1.3.0)\n",
      "Downloading peft-0.4.0-py3-none-any.whl (72 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m72.9/72.9 kB\u001b[0m \u001b[31m314.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading accelerate-0.22.0-py3-none-any.whl (251 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m251.2/251.2 kB\u001b[0m \u001b[31m199.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading cmake-3.27.2-py2.py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (26.1 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m26.1/26.1 MB\u001b[0m \u001b[31m205.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hBuilding wheels for collected packages: lit\n",
      "  Building wheel for lit (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for lit: filename=lit-16.0.6-py3-none-any.whl size=93584 sha256=46396c9563f479caa85d93eac94ddec407f5fa8958f4ff2e56e195eac769de1a\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-8izvp2cy/wheels/14/f9/07/bb2308587bc2f57158f905a2325f6a89a2befa7437b2d7e137\n",
      "Successfully built lit\n",
      "Installing collected packages: lit, cmake, nvidia-nvtx-cu11, nvidia-nccl-cu11, nvidia-cusparse-cu11, nvidia-curand-cu11, nvidia-cufft-cu11, nvidia-cuda-runtime-cu11, nvidia-cuda-nvrtc-cu11, nvidia-cuda-cupti-cu11, nvidia-cublas-cu11, nvidia-cusolver-cu11, nvidia-cudnn-cu11, triton, torch, accelerate, peft\n",
      "Successfully installed accelerate-0.22.0 cmake-3.27.2 lit-16.0.6 nvidia-cublas-cu11-11.10.3.66 nvidia-cuda-cupti-cu11-11.7.101 nvidia-cuda-nvrtc-cu11-11.7.99 nvidia-cuda-runtime-cu11-11.7.99 nvidia-cudnn-cu11-8.5.0.96 nvidia-cufft-cu11-10.9.0.58 nvidia-curand-cu11-10.2.10.91 nvidia-cusolver-cu11-11.4.0.1 nvidia-cusparse-cu11-11.7.4.91 nvidia-nccl-cu11-2.14.3 nvidia-nvtx-cu11-11.7.91 peft-0.4.0 torch-2.0.1 triton-2.0.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Collecting accelerate==0.21.0\n",
      "  Obtaining dependency information for accelerate==0.21.0 from https://files.pythonhosted.org/packages/70/f9/c381bcdd0c3829d723aa14eec8e75c6c377b4ca61ec68b8093d9f35fc7a7/accelerate-0.21.0-py3-none-any.whl.metadata\n",
      "  Downloading accelerate-0.21.0-py3-none-any.whl.metadata (17 kB)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from accelerate==0.21.0) (1.25.1)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from accelerate==0.21.0) (21.3)\n",
      "Requirement already satisfied: psutil in /opt/conda/lib/python3.10/site-packages (from accelerate==0.21.0) (5.9.0)\n",
      "Requirement already satisfied: pyyaml in /opt/conda/lib/python3.10/site-packages/PyYAML-6.0-py3.10-linux-x86_64.egg (from accelerate==0.21.0) (6.0)\n",
      "Requirement already satisfied: torch>=1.10.0 in /opt/conda/lib/python3.10/site-packages (from accelerate==0.21.0) (2.0.1)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->accelerate==0.21.0) (3.0.9)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (3.6.0)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (4.3.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (1.10.1)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (2.8.4)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (3.1.2)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu11==11.7.101 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.7.101)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (8.5.0.96)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-cufft-cu11==10.9.0.58 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (10.9.0.58)\n",
      "Requirement already satisfied: nvidia-curand-cu11==10.2.10.91 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (10.2.10.91)\n",
      "Requirement already satisfied: nvidia-cusolver-cu11==11.4.0.1 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.4.0.1)\n",
      "Requirement already satisfied: nvidia-cusparse-cu11==11.7.4.91 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.7.4.91)\n",
      "Requirement already satisfied: nvidia-nccl-cu11==2.14.3 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (2.14.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu11==11.7.91 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (11.7.91)\n",
      "Requirement already satisfied: triton==2.0.0 in /opt/conda/lib/python3.10/site-packages (from torch>=1.10.0->accelerate==0.21.0) (2.0.0)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.10.0->accelerate==0.21.0) (68.0.0)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.10.0->accelerate==0.21.0) (0.41.0)\n",
      "Requirement already satisfied: cmake in /opt/conda/lib/python3.10/site-packages (from triton==2.0.0->torch>=1.10.0->accelerate==0.21.0) (3.27.2)\n",
      "Requirement already satisfied: lit in /opt/conda/lib/python3.10/site-packages (from triton==2.0.0->torch>=1.10.0->accelerate==0.21.0) (16.0.6)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch>=1.10.0->accelerate==0.21.0) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch>=1.10.0->accelerate==0.21.0) (1.3.0)\n",
      "Downloading accelerate-0.21.0-py3-none-any.whl (244 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m244.2/244.2 kB\u001b[0m \u001b[31m90.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: accelerate\n",
      "  Attempting uninstall: accelerate\n",
      "    Found existing installation: accelerate 0.22.0\n",
      "    Uninstalling accelerate-0.22.0:\n",
      "      Successfully uninstalled accelerate-0.22.0\n",
      "Successfully installed accelerate-0.21.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Collecting datasets==2.14.1\n",
      "  Obtaining dependency information for datasets==2.14.1 from https://files.pythonhosted.org/packages/d1/7a/fff542b2958a664e2c78dd3cb0431f2e03a56e92826f97bcccf958f60324/datasets-2.14.1-py3-none-any.whl.metadata\n",
      "  Downloading datasets-2.14.1-py3-none-any.whl.metadata (19 kB)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (1.25.1)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (12.0.1)\n",
      "Requirement already satisfied: dill<0.3.8,>=0.3.0 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (0.3.7)\n",
      "Requirement already satisfied: pandas in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (1.4.4)\n",
      "Requirement already satisfied: requests>=2.19.0 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.62.1 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (4.64.1)\n",
      "Collecting xxhash (from datasets==2.14.1)\n",
      "  Obtaining dependency information for xxhash from https://files.pythonhosted.org/packages/13/c3/e942893f4864a424514c81640f114980cfd5aff7e7414d1e0255f4571111/xxhash-3.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading xxhash-3.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (12 kB)\n",
      "Requirement already satisfied: multiprocess in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (0.70.15)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (2022.7.1)\n",
      "Collecting aiohttp (from datasets==2.14.1)\n",
      "  Obtaining dependency information for aiohttp from https://files.pythonhosted.org/packages/3e/f6/fcda07dd1e72260989f0b22dde999ecfe80daa744f23ca167083683399bc/aiohttp-3.8.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading aiohttp-3.8.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.7 kB)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.14.0 in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (0.16.4)\n",
      "Requirement already satisfied: packaging in /opt/conda/lib/python3.10/site-packages (from datasets==2.14.1) (21.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.10/site-packages/PyYAML-6.0-py3.10-linux-x86_64.egg (from datasets==2.14.1) (6.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0.0,>=0.14.0->datasets==2.14.1) (3.6.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub<1.0.0,>=0.14.0->datasets==2.14.1) (4.3.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging->datasets==2.14.1) (3.0.9)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests>=2.19.0->datasets==2.14.1) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests>=2.19.0->datasets==2.14.1) (3.3)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests>=2.19.0->datasets==2.14.1) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests>=2.19.0->datasets==2.14.1) (2023.7.22)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp->datasets==2.14.1) (23.1.0)\n",
      "Collecting multidict<7.0,>=4.5 (from aiohttp->datasets==2.14.1)\n",
      "  Downloading multidict-6.0.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (114 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m114.5/114.5 kB\u001b[0m \u001b[31m68.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting async-timeout<5.0,>=4.0.0a3 (from aiohttp->datasets==2.14.1)\n",
      "  Obtaining dependency information for async-timeout<5.0,>=4.0.0a3 from https://files.pythonhosted.org/packages/a7/fa/e01228c2938de91d47b307831c62ab9e4001e747789d0b05baf779a6488c/async_timeout-4.0.3-py3-none-any.whl.metadata\n",
      "  Downloading async_timeout-4.0.3-py3-none-any.whl.metadata (4.2 kB)\n",
      "Collecting yarl<2.0,>=1.0 (from aiohttp->datasets==2.14.1)\n",
      "  Downloading yarl-1.9.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (268 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m426.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting frozenlist>=1.1.1 (from aiohttp->datasets==2.14.1)\n",
      "  Obtaining dependency information for frozenlist>=1.1.1 from https://files.pythonhosted.org/packages/1e/28/74b8b6451c89c070d34e753d8b65a1e4ce508a6808b18529f36e8c0e2184/frozenlist-1.4.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading frozenlist-1.4.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (5.2 kB)\n",
      "Collecting aiosignal>=1.1.2 (from aiohttp->datasets==2.14.1)\n",
      "  Downloading aiosignal-1.3.1-py3-none-any.whl (7.6 kB)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /opt/conda/lib/python3.10/site-packages (from pandas->datasets==2.14.1) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/conda/lib/python3.10/site-packages (from pandas->datasets==2.14.1) (2022.1)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.8.1->pandas->datasets==2.14.1) (1.16.0)\n",
      "Downloading datasets-2.14.1-py3-none-any.whl (492 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m492.4/492.4 kB\u001b[0m \u001b[31m439.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading aiohttp-3.8.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m471.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading xxhash-3.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (194 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.1/194.1 kB\u001b[0m \u001b[31m442.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading async_timeout-4.0.3-py3-none-any.whl (5.7 kB)\n",
      "Downloading frozenlist-1.4.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (225 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m225.7/225.7 kB\u001b[0m \u001b[31m450.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: xxhash, multidict, frozenlist, async-timeout, yarl, aiosignal, aiohttp, datasets\n",
      "Successfully installed aiohttp-3.8.5 aiosignal-1.3.1 async-timeout-4.0.3 datasets-2.14.1 frozenlist-1.4.0 multidict-6.0.4 xxhash-3.3.0 yarl-1.9.2\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install transformers==4.31.0\n",
    "%pip install peft==0.4.0\n",
    "%pip install accelerate==0.21.0\n",
    "#%pip install bitsandbytes==0.40.2\n",
    "#%pip install safetensors==0.3.1\n",
    "#%pip install tokenizers==0.13.3\n",
    "%pip install datasets==2.14.1\n",
    "\n",
    "# %pip install -U transformers==4.39.0\n",
    "# %pip install -U peft==0.5.0\n",
    "# %pip install -U accelerate==0.26.0\n",
    "# #%pip install bitsandbytes #==0.40.2\n",
    "# #%pip install safetensors==0.3.1\"\n",
    "# #%pip install tokenizers==0.13.3\n",
    "# %pip install -U datasets==2.17.0\n",
    "# %pip install --no-cache https://developer.download.nvidia.com/compute/redist/jp/v60dp/pytorch/torch-2.3.0a0+ebedce2.nv24.02-cp310-cp310-linux_aarch64.whl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "861f4453-1601-4420-9fd2-4d6f518ec327",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:155: UserWarning: A NumPy version >=1.18.5 and <1.25.0 is required for this version of SciPy (detected version 1.25.1\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import argparse\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    set_seed,\n",
    "    default_data_collator,\n",
    "#    BitsAndBytesConfig,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    ")\n",
    "from datasets import load_dataset\n",
    "import torch\n",
    "\n",
    "#import bitsandbytes as bnb\n",
    "#from huggingface_hub import login, HfFolder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f021069-dfec-4355-80dd-f8c18af2d791",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Fine-Tune LLaMA 7B in Amazon SageMaker Studio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ca663aa7-7c3f-4123-b775-e6dae28b4f71",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import argparse\n",
    "parser = argparse.ArgumentParser()\n",
    "\n",
    "# add model id and dataset path argument\n",
    "parser.add_argument(\n",
    "    \"--model_id\",\n",
    "    type=str,\n",
    "    default=\"NousResearch/Llama-2-7b-hf\", # not gated\n",
    "    help=\"Model id to use for training.\",\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--dataset_path\", \n",
    "    type=str, \n",
    "    default=\"lm_dataset\", \n",
    "    help=\"Path to dataset.\"\n",
    ")\n",
    "# parser.add_argument(\n",
    "#     \"--hf_token\", \n",
    "#     type=str, \n",
    "#     default=HfFolder.get_token(), \n",
    "#     help=\"Path to dataset.\"\n",
    "# )\n",
    "# add training hyperparameters for epochs, batch size, learning rate, and seed\n",
    "parser.add_argument(\n",
    "    \"--epochs\", \n",
    "    type=int, \n",
    "    default=1, \n",
    "    help=\"Number of epochs to train for.\"\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--per_device_train_batch_size\",\n",
    "    type=int,\n",
    "    default=1,\n",
    "    help=\"Batch size to use for training.\",\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--lr\", \n",
    "    type=float, \n",
    "    default=5e-5, \n",
    "    help=\"Learning rate to use for training.\"\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--seed\", \n",
    "    type=int, \n",
    "    default=42, \n",
    "    help=\"Seed to use for training.\"\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--gradient_checkpointing\",\n",
    "    type=bool,\n",
    "    default=True,\n",
    "    help=\"Path to deepspeed config file.\",\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--bf16\",\n",
    "    type=bool,\n",
    "    default=True if torch.cuda.get_device_capability()[0] >= 8 else False,\n",
    "    help=\"Whether to use bf16.\",\n",
    ")\n",
    "# parser.add_argument(\n",
    "#     \"--merge_weights\",\n",
    "#     type=bool,\n",
    "#     default=True,\n",
    "#     help=\"Whether to merge LoRA weights with base model.\",\n",
    "# )\n",
    "args, _ = parser.parse_known_args()\n",
    "\n",
    "# if args.hf_token:\n",
    "#     print(f\"Logging into the Hugging Face Hub with token {args.hf_token[:10]}...\")\n",
    "#     login(token=args.hf_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f47e9fa5-446d-4fb0-973d-c270a30db03e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # COPIED FROM https://github.com/artidoro/qlora/blob/main/qlora.py\n",
    "# def print_trainable_parameters(model, use_4bit=False):\n",
    "#     \"\"\"\n",
    "#     Prints the number of trainable parameters in the model.\n",
    "#     \"\"\"\n",
    "#     trainable_params = 0\n",
    "#     all_param = 0\n",
    "#     for _, param in model.named_parameters():\n",
    "#         num_params = param.numel()\n",
    "#         # if using DS Zero 3 and the weights are initialized empty        \n",
    "#         if num_params == 0 and hasattr(param, \"ds_numel\"):\n",
    "#             num_params = param.ds_numel\n",
    "\n",
    "#         all_param += num_params\n",
    "#         if param.requires_grad:\n",
    "#             trainable_params += num_params\n",
    "#     # if use_4bit:\n",
    "#     #     trainable_params /= 2\n",
    "#     print(\n",
    "#         f\"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}\"\n",
    "#     )\n",
    "\n",
    "\n",
    "# # # COPIED FROM https://github.com/artidoro/qlora/blob/main/qlora.py\n",
    "# # def find_all_linear_names(model):\n",
    "# #     lora_module_names = set()\n",
    "# # #    for name, module in model.named_modules():\n",
    "# #         # if isinstance(module, bnb.nn.Linear4bit):\n",
    "# #         #     names = name.split(\".\")\n",
    "# #         #     lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n",
    "\n",
    "# #     if \"lm_head\" in lora_module_names:  # needed for 16-bit\n",
    "# #         lora_module_names.remove(\"lm_head\")\n",
    "# #     return list(lora_module_names)\n",
    "\n",
    "\n",
    "# def create_peft_model(model, gradient_checkpointing=True, bf16=True):\n",
    "#     from peft import (\n",
    "#         get_peft_model,\n",
    "#         LoraConfig,\n",
    "#         TaskType,\n",
    "# #        prepare_model_for_kbit_training,\n",
    "#     )\n",
    "#     from peft.tuners.lora import LoraLayer\n",
    "\n",
    "#     # # prepare int-4 model for training\n",
    "#     # model = prepare_model_for_kbit_training(\n",
    "#     #     model, use_gradient_checkpointing=gradient_checkpointing\n",
    "#     # )\n",
    "#     if gradient_checkpointing:\n",
    "#         model.gradient_checkpointing_enable()\n",
    "\n",
    "#     # # get lora target modules\n",
    "#     # modules = find_all_linear_names(model)    \n",
    "    \n",
    "#     #If only targeting attention blocks of the model\n",
    "#     #modules = [\"q_proj\", \"v_proj\"]\n",
    "\n",
    "#     #If targeting all linear layers\n",
    "#     modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','down_proj','up_proj'] #,'lm_head']\n",
    "\n",
    "#     print(f\"Found {len(modules)} modules to quantize: {modules}\")\n",
    "\n",
    "#     peft_config = LoraConfig(\n",
    "#         r=64,\n",
    "#         lora_alpha=16,\n",
    "#         target_modules=modules,\n",
    "#         lora_dropout=0.1,\n",
    "#         bias=\"none\",\n",
    "#         task_type=TaskType.CAUSAL_LM,\n",
    "#     )\n",
    "\n",
    "#     model = get_peft_model(model, peft_config)\n",
    "\n",
    "#     # pre-process the model by upcasting the layer norms in float 32 for\n",
    "#     # for name, module in model.named_modules():\n",
    "#     #     if isinstance(module, LoraLayer):\n",
    "#     #         if bf16:\n",
    "#     #             module = module.to(torch.bfloat16)\n",
    "#     #     if \"norm\" in name:\n",
    "#     #         module = module.to(torch.float32)\n",
    "#     #     if \"lm_head\" in name or \"embed_tokens\" in name:\n",
    "#     #         if hasattr(module, \"weight\"):\n",
    "#     #             if bf16 and module.weight.dtype == torch.float32:\n",
    "#     #                 module = module.to(torch.bfloat16)\n",
    "\n",
    "#     model.print_trainable_parameters()\n",
    "#     return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b0c2612-8951-4f4c-b5c4-1277fca469cb",
   "metadata": {},
   "source": [
    "## Load and prepare the dataset\n",
    "\n",
    "we will use the [dolly](https://huggingface.co/datasets/databricks/databricks-dolly-15k) an open source dataset of instruction-following records generated by thousands of Databricks employees in several of the behavioral categories outlined in the [InstructGPT paper](https://arxiv.org/abs/2203.02155), including brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization.\n",
    "\n",
    "```python\n",
    "{\n",
    "  \"instruction\": \"What is world of warcraft\",\n",
    "  \"context\": \"\",\n",
    "  \"response\": \"World of warcraft is a massive online multi player role playing game. It was released in 2004 by bizarre entertainment\"\n",
    "}\n",
    "```\n",
    "\n",
    "To load the `samsum` dataset, we use the `load_dataset()` method from the 🤗 Datasets library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "83bace11-a0fb-433c-8d2a-c1e9bc8864f0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset size: 1000\n",
      "{'instruction': 'Give me a list of some of the most popular song from 70s Japanese Pop singer, Mariya Takeuchi', 'context': '', 'response': '1. Plastic Love\\n2. Stay with Me\\n3. September\\n4. Miracle Love\\n5. Yume No Tsuzuki', 'category': 'brainstorming'}\n"
     ]
    }
   ],
   "source": [
    "# set seed\n",
    "set_seed(args.seed)\n",
    "\n",
    "from datasets import load_dataset\n",
    "from random import randrange\n",
    "\n",
    "# Load dataset from the hub\n",
    "dataset = load_dataset(\"databricks/databricks-dolly-15k\", split=\"train\")\n",
    "dataset = dataset.select(range(1000))\n",
    "\n",
    "print(f\"dataset size: {len(dataset)}\")\n",
    "print(dataset[randrange(len(dataset))])\n",
    "# dataset size: 15011"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89f07a4c-8035-45ef-a498-93acd8ba0d33",
   "metadata": {},
   "source": [
    "To instruct tune our model we need to convert our structured examples into a collection of tasks described via instructions. We define a `formatting_function` that takes a sample and returns a string with our format instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4e938006-7605-4e4c-9cda-dc80f194cc57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def format_dolly(sample):\n",
    "    instruction = f\"### Instruction\\n{sample['instruction']}\"\n",
    "    context = f\"### Context\\n{sample['context']}\" if len(sample[\"context\"]) > 0 else None\n",
    "    response = f\"### Answer\\n{sample['response']}\"\n",
    "    # join all the parts together\n",
    "    prompt = \"\\n\\n\".join([i for i in [instruction, context, response] if i is not None])\n",
    "    return prompt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "483e7629-595f-49c3-a2c1-e19f65eac49a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Instruction\n",
      "What is the Research Collaboratory for Structural Bioinformatics Protein Data Bank (RCSB PDB)?\n",
      "\n",
      "### Answer\n",
      "The Research Collaboratory for Structural Bioinformatics Protein Data Bank (RCSB PDB) is a database that provides a wealth of information about the 3D structures of proteins, nucleic acids, and other macromolecules. The database contains experimentally determined atomic coordinates for a large number of macromolecules, which can be used to study their structures, functions, and interactions. The RCSB PDB is widely used in genomics research and drug discovery, as it provides a valuable resource for understanding the structural basis of many biological processes and for designing new drugs that target specific macromolecules.\n",
      "\n",
      "In addition to the atomic coordinates, the RCSB PDB contains a wealth of additional information about each macromolecule, including experimental methods used for structure determination, citations to relevant scientific literature, and information about biological function, sequence, and homology. The database also provides a variety of tools and resources for visualizing, analyzing, and downloading structural data, including a web-based viewer that allows users to interactively explore the 3D structures of macromolecules.\n"
     ]
    }
   ],
   "source": [
    "from random import randrange\n",
    "\n",
    "print(format_dolly(dataset[randrange(len(dataset))]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "48c0c16f-9471-4e84-8e0f-a2b4ca476b25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "#model_id = \"meta-llama/Llama-2-13b-hf\" # sharded weights, gated\n",
    "model_id = \"NousResearch/Llama-2-7b-hf\" # not gated\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ecb7e1ca-9c07-432b-ae73-c06776a4897b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Instruction\n",
      "Extract the owner of Lamborghini and a listing of the different types of Huracan cars that Lamborghini has produced for its Motorsport division.\n",
      "\n",
      "### Context\n",
      "Automobili Lamborghini S.p.A. (Italian pronunciation: [autoˈmɔːbili lamborˈɡiːni]) is an Italian manufacturer of luxury sports cars and SUVs based in Sant'Agata Bolognese. The company is owned by the Volkswagen Group through its subsidiary Audi.\n",
      "\n",
      "Ferruccio Lamborghini (1916–1993), an Italian manufacturing magnate, founded Automobili Ferruccio Lamborghini S.p.A. in 1963 to compete with Ferrari. The company was noted for using a rear mid-engine, rear-wheel drive layout. Lamborghini grew rapidly during its first decade, but sales plunged in the wake of the 1973 worldwide financial downturn and the oil crisis. The firm's ownership changed three times after 1973, including a bankruptcy in 1978. American Chrysler Corporation took control of Lamborghini in 1987 and sold it to Malaysian investment group Mycom Setdco and Indonesian group V'Power Corporation in 1994. In 1998, Mycom Setdco and V'Power sold Lamborghini to the Volkswagen Group where it was placed under the control of the group's Audi division.\n",
      "\n",
      "New products and model lines were introduced to the brand's portfolio and brought to the market and saw an increased productivity for the brand. In the late 2000s, during the worldwide financial crisis and the subsequent economic crisis, Lamborghini's sales saw a drop of nearly 50 per cent.\n",
      "\n",
      "Lamborghini currently produces the V12-powered Aventador and the V10-powered Huracán, along with the Urus SUV powered by a twin-turbo V8 engine. In addition, the company produces V12 engines for offshore powerboat racing.\n",
      "\n",
      "Lamborghini Trattori, founded in 1948 by Ferruccio Lamborghini, is headquartered in Pieve di Cento, Italy and continues to produce tractors. Since 1973, Lamborghini Trattori has been a separate entity from the automobile manufacturer.\n",
      "\n",
      "History\n",
      "Main article: History of Lamborghini\n",
      "\n",
      "Ferruccio Lamborghini with a Jarama and a tractor of his brand\n",
      "Manufacturing magnate Italian Ferruccio Lamborghini founded the company in 1963 with the objective of producing a refined grand touring car to compete with offerings from established marques such as Ferrari. The company's first models, such as the 350 GT, were released in the mid-1960s. Lamborghini was noted for the 1966 Miura sports coupé, which used a rear mid-engine, rear-wheel drive layout.\n",
      "\n",
      "Lamborghini grew rapidly during its first ten years, but sales fell in the wake of the 1973 worldwide financial downturn and the oil crisis. Ferruccio Lamborghini sold the company to Georges-Henri Rossetti and René Leimer and retired in 1974. The company went bankrupt in 1978, and was placed in the receivership of brothers Jean-Claude and Patrick Mimran in 1980. The Mimrans purchased the company out of receivership by 1984 and invested heavily in its expansion. Under the Mimrans' management, Lamborghini's model line was expanded from the Countach to include the Jalpa sports car and the LM002 high-performance off-road vehicle.\n",
      "\n",
      "The Mimrans sold Lamborghini to the Chrysler Corporation in 1987. After replacing the Countach with the Diablo and discontinuing the Jalpa and the LM002, Chrysler sold Lamborghini to Malaysian investment group Mycom Setdco and Indonesian group V'Power Corporation in 1994. In 1998, Mycom Setdco and V'Power sold Lamborghini to the Volkswagen Group where it was placed under the control of the group's Audi division. New products and model lines were introduced to the brand's portfolio and brought to the market and saw an increased productivity for the brand Lamborghini. In the late 2000s, during the worldwide financial crisis and the subsequent economic crisis, Lamborghini's sales saw a drop of nearly 50 per cent.\n",
      "\n",
      "In 2021, the CEO of Lamborghini said that by 2024 all its models will be hybrid.\n",
      "\n",
      "Automobiles produced\n",
      "Lamborghini Motorsport Division Squadra Corse produces GT3 cars and cars for their Super Trofeo events based on the Gallardo and Huracán. Apart from them, the Squadra Corse builds cars upon customer request.\n",
      "\n",
      "GT3 and Super Trofeo Cars\n",
      "Gallardo LP 570-4 Super Trofeo\n",
      "Gallardo LP 560-4 Super Trofeo\n",
      "Huracán LP 620-2 Super Trofeo EVO\n",
      "Huracán LP 620-2 Super Trofeo EVO2\n",
      "Huracán Super Trofeo GT2\n",
      "Huracán GT3\n",
      "Huracán GT3 Evo\n",
      "Huracán GT3 Evo 2\n",
      "Special cars\n",
      "These cars were built by Squadra Corse upon customer request.\n",
      "\n",
      "Essenza SCV12\n",
      "SC18 Alston\n",
      "SC20\n",
      "\n",
      "### Answer\n",
      "Vokswagen Group owns Lamborghini through its subsidiary, Audi.\n",
      "\n",
      "Lamborghini has produced the following Huracan motorsport models:\n",
      "- Huracán LP 620-2 Super Trofeo EVO\n",
      "- Huracán LP 620-2 Super Trofeo EVO2\n",
      "- Huracán Super Trofeo GT2\n",
      "- Huracán GT3\n",
      "- Huracán GT3 Evo\n",
      "- Huracán GT3 Evo 2</s>\n",
      "Total number of samples: 105\n"
     ]
    }
   ],
   "source": [
    "from random import randint\n",
    "from itertools import chain\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "# template dataset to add prompt to each sample\n",
    "def template_dataset(sample):\n",
    "    sample[\"text\"] = f\"{format_dolly(sample)}{tokenizer.eos_token}\"\n",
    "    return sample\n",
    "\n",
    "\n",
    "# apply prompt template per sample\n",
    "dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))\n",
    "# print random sample\n",
    "print(dataset[randint(0, len(dataset))][\"text\"])\n",
    "\n",
    "# empty list to save remainder from batches to use in next batch\n",
    "remainder = {\"input_ids\": [], \"attention_mask\": [], \"token_type_ids\": []}\n",
    "\n",
    "def chunk(sample, chunk_length=2048):\n",
    "    # define global remainder variable to save remainder from batches to use in next batch\n",
    "    global remainder\n",
    "    # Concatenate all texts and add remainder from previous batch\n",
    "    concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}\n",
    "    concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}\n",
    "    # get total number of tokens for batch\n",
    "    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])\n",
    "\n",
    "    # get max number of chunks for batch\n",
    "    if batch_total_length >= chunk_length:\n",
    "        batch_chunk_length = (batch_total_length // chunk_length) * chunk_length\n",
    "\n",
    "    # Split by chunks of max_len.\n",
    "    result = {\n",
    "        k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    # add remainder to global variable for next batch\n",
    "    remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}\n",
    "    # prepare labels\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result\n",
    "\n",
    "\n",
    "# tokenize and chunk dataset\n",
    "lm_dataset = dataset.map(\n",
    "    lambda sample: tokenizer(sample[\"text\"]), batched=True, remove_columns=list(dataset.features)\n",
    ").map(\n",
    "    partial(chunk, chunk_length=2048),\n",
    "    batched=True,\n",
    ")\n",
    "\n",
    "# Print total number of samples\n",
    "print(f\"Total number of samples: {len(lm_dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "71d95685-4573-4943-8ecb-9d42167a5808",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['input_ids', 'attention_mask', 'labels'],\n",
      "    num_rows: 105\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "# The chunking above will reduce the number of rows\n",
    "print(lm_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6eafc7bb-cd67-4ddc-8930-90dc649d5626",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "078227766fdb4d0185f34cac51b2404c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# load model from the hub with a bnb config\n",
    "# bnb_config = BitsAndBytesConfig(\n",
    "#     load_in_4bit=True,\n",
    "#     bnb_4bit_use_double_quant=True,\n",
    "#     bnb_4bit_quant_type=\"nf4\",\n",
    "#     bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "# )\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    args.model_id,\n",
    "    use_cache=False\n",
    "    if args.gradient_checkpointing\n",
    "    else True,  # this is needed for gradient checkpointing\n",
    "    device_map=\"auto\",\n",
    "#    quantization_config=bnb_config,\n",
    ")\n",
    "\n",
    "# create peft config\n",
    "# model = create_peft_model(\n",
    "#     model, gradient_checkpointing=args.gradient_checkpointing, bf16=args.bf16\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "37c19f83-da6a-4f7b-9b0c-e37704b9bd5c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='105' max='105' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [105/105 02:00, Epoch 1/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>1.655900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>1.575000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>1.565500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>1.500000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>1.607200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>1.491800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>70</td>\n",
       "      <td>1.463000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>1.599700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>90</td>\n",
       "      <td>1.642400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>1.532600</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=105, training_loss=1.5594511486235119, metrics={'train_runtime': 127.5971, 'train_samples_per_second': 0.823, 'train_steps_per_second': 0.823, 'total_flos': 8525059027107840.0, 'train_loss': 1.5594511486235119, 'epoch': 1.0})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define training args\n",
    "output_dir = \"./tmp/llama2\"\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    per_device_train_batch_size=args.per_device_train_batch_size,\n",
    "    bf16=args.bf16,  # Use BF16 if available\n",
    "    learning_rate=args.lr,\n",
    "    num_train_epochs=args.epochs,\n",
    "    gradient_checkpointing=args.gradient_checkpointing,\n",
    "    # logging strategies\n",
    "    logging_dir=f\"{output_dir}/logs\",\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"no\",\n",
    ")\n",
    "\n",
    "# Create Trainer instance\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=lm_dataset,\n",
    "    data_collator=default_data_collator,\n",
    ")\n",
    "\n",
    "# Start training\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2a831a1c-adcc-4674-8bd5-5cfc4986aa34",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./llama2_dolly/tokenizer_config.json',\n",
       " './llama2_dolly/special_tokens_map.json',\n",
       " './llama2_dolly/tokenizer.json')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sagemaker_save_dir=\"./llama2_dolly\"\n",
    "# if args.merge_weights:\n",
    "#     # merge adapter weights with base model and save\n",
    "#     # save int 4 model\n",
    "#     trainer.model.save_pretrained(output_dir, safe_serialization=False)\n",
    "#     # clear memory\n",
    "#     del model\n",
    "#     del trainer\n",
    "#     torch.cuda.empty_cache()\n",
    "\n",
    "#     from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "#     # load PEFT model in fp16\n",
    "#     model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "#         output_dir,\n",
    "#         low_cpu_mem_usage=True,\n",
    "#         torch_dtype=torch.bfloat16,\n",
    "#     )  \n",
    "#     # Merge LoRA and base model and save\n",
    "#     model = model.merge_and_unload()        \n",
    "#     model.save_pretrained(\n",
    "#         sagemaker_save_dir, safe_serialization=True, max_shard_size=\"2GB\"\n",
    "#     )\n",
    "# else:\n",
    "\n",
    "trainer.model.save_pretrained(\n",
    "    sagemaker_save_dir, safe_serialization=True\n",
    ")\n",
    "\n",
    "# save tokenizer for easy inference\n",
    "tokenizer = AutoTokenizer.from_pretrained(args.model_id)\n",
    "tokenizer.save_pretrained(sagemaker_save_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a81b21d-cc86-45a4-9f06-151d5f55153f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   }
  ],
  "instance_type": "ml.t3.medium",
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}