llama-2-13b-on-tpu-training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "412fa4d8",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2023-11-10T09:59:36.730132Z",
     "iopub.status.busy": "2023-11-10T09:59:36.729791Z",
     "iopub.status.idle": "2023-11-10T10:01:39.732511Z",
     "shell.execute_reply": "2023-11-10T10:01:39.731406Z"
    },
    "papermill": {
     "duration": 123.011398,
     "end_time": "2023-11-10T10:01:39.735246",
     "exception": false,
     "start_time": "2023-11-10T09:59:36.723848",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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\r\n",
      "\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.1\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\r\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.\r\n",
      "torchvision 0.15.1 requires torch==2.0.0, but you have torch 2.1.0+cpu which is incompatible.\r\n",
      "torchtext 0.15.1 requires torch==2.0.0, but you have torch 2.1.0+cpu which is incompatible.\r\n",
      "torchdata 0.6.0 requires torch==2.0.0, but you have torch 2.1.0+cpu which is incompatible.\r\n",
      "torchaudio 2.0.0 requires torch==2.0.0, but you have torch 2.1.0+cpu which is incompatible.\u001b[0m\u001b[31m\r\n",
      "\u001b[0m\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\r\n",
      "\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.1\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\r\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.\r\n",
      "tensorboardx 2.6.1 requires protobuf>=4.22.3, but you have protobuf 3.20.3 which is incompatible.\u001b[0m\u001b[31m\r\n",
      "\u001b[0m\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\r\n",
      "\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.1\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\r\n",
      "Found existing installation: tensorflow 2.12.0\r\n",
      "Uninstalling tensorflow-2.12.0:\r\n",
      "  Successfully uninstalled tensorflow-2.12.0\r\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\r\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip3 install transformers datasets sentencepiece peft -q\n",
    "!pip install torch~=2.1.0 --index-url https://download.pytorch.org/whl/cpu -q # Updating torch since we need the latest version\n",
    "!pip install torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html -q\n",
    "!pip uninstall tensorflow -y # If we don't do this, TF may take over TPU and cause permission error for PT\n",
    "!cp /kaggle/input/utils-xla/spmd_util.py . # Copy script to working directory from this repo https://github.com/HeegyuKim/torch-xla-SPMD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "32348dfe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:01:39.744910Z",
     "iopub.status.busy": "2023-11-10T10:01:39.744625Z",
     "iopub.status.idle": "2023-11-10T10:01:54.063187Z",
     "shell.execute_reply": "2023-11-10T10:01:54.061546Z"
    },
    "papermill": {
     "duration": 14.326263,
     "end_time": "2023-11-10T10:01:54.065568",
     "exception": false,
     "start_time": "2023-11-10T10:01:39.739305",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import datasets\n",
    "import torch.optim as optim\n",
    "import torch_xla.debug.profiler as xp\n",
    "import torch_xla.core.xla_model as xm\n",
    "import torch_xla.distributed.xla_multiprocessing as xmp # We also import mp modules if we wanna use that for some reason\n",
    "import torch_xla.distributed.parallel_loader as pl\n",
    "import torch_xla.test.test_utils as test_utils\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import re\n",
    "import torch_xla.experimental.xla_sharding as xs\n",
    "import torch_xla.core.xla_model as xm\n",
    "from transformers import (\n",
    "    GPTNeoXConfig, T5Config, LlamaConfig, AutoTokenizer, AutoModelForCausalLM, DataCollatorWithPadding, AutoConfig, AutoModelForSequenceClassification\n",
    ") # You can use any of models with those configs (even flan T5 xxl!). Other models are not supported.\n",
    "\n",
    "from transformers import logging as hf_logging\n",
    "import torch.nn.functional as F\n",
    "import torch_xla.runtime as xr\n",
    "\n",
    "xr.use_spmd()\n",
    "\n",
    "import torch_xla.experimental.xla_sharding as xs # \"experimental\" prefix always means you're gonna have a good time LMAO\n",
    "from torch_xla.experimental.xla_sharded_tensor import XLAShardedTensor\n",
    "from torch_xla.experimental.xla_sharding import Mesh\n",
    "\n",
    "from peft import LoraConfig, TaskType, get_peft_model # If we wanna use peft. Quantazation requiers GPU though. You'll have to download already quantazed models\n",
    "from spmd_util import partition_module                # You could experiment with using already quantazed models like 4bit/Llama-2-7b-Chat-GPTQ if you're feeling funny\n",
    "from datasets import Dataset, load_dataset, concatenate_datasets\n",
    "from dataclasses import dataclass\n",
    "from tqdm import tqdm\n",
    "\n",
    "import transformers\n",
    "import datasets\n",
    "from datasets import Dataset\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "!export USE_TORCH=True #If we don't do this, transformers will seemingly bork the session upon import. Really weird error.\n",
    "os.environ[\"PJRT_DEVICE\"] = \"TPU\"\n",
    "os.environ.pop('TPU_PROCESS_ADDRESSES')\n",
    "os.environ.pop('CLOUD_TPU_TASK_ID')\n",
    "hf_logging.set_verbosity_error() # It can still display warnings which is a bit annoying but whatever\n",
    "\n",
    "\n",
    "MAX_INPUT=512\n",
    "MODEL = \"/kaggle/input/best-llm-starter-pack/nous_hermes/Nous-Hermes-Llama2-13b/snapshots/8f95aa9cd207db7b24179fc779c2b8973e71bee2\" # Provide your own model directory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e9d832f",
   "metadata": {},
   "source": [
    "Below is pandas processing done for Kaggle \"Detect AI\" competition. Provide your own dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1af0f562",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:01:54.075951Z",
     "iopub.status.busy": "2023-11-10T10:01:54.075214Z",
     "iopub.status.idle": "2023-11-10T10:01:58.214017Z",
     "shell.execute_reply": "2023-11-10T10:01:58.213070Z"
    },
    "papermill": {
     "duration": 4.147212,
     "end_time": "2023-11-10T10:01:58.216803",
     "exception": false,
     "start_time": "2023-11-10T10:01:54.069591",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_ai = pd.read_csv('/kaggle/input/daigt-external-dataset/daigt_external_dataset.csv')\n",
    "ai_samples_500 = pd.read_csv('/kaggle/input/llm-generated-essays/ai_generated_train_essays.csv')\n",
    "persuade_df = pd.read_csv('/kaggle/input/persaude-corpus-2/persuade_2.0_human_scores_demo_id_github.csv')\n",
    "df_llama = pd.read_csv('/kaggle/input/daigt-data-llama-70b-and-falcon180b/llama_70b_v1.csv')\n",
    "df_falcon = pd.read_csv('/kaggle/input/daigt-data-llama-70b-and-falcon180b/falcon_180b_v1.csv')\n",
    "proper_train = pd.read_csv('/kaggle/input/daigt-proper-train-dataset/train_drcat_01.csv')\n",
    "aug_df = pd.read_csv('/kaggle/input/argugpt/argugpt.csv')\n",
    "corrected_df = pd.read_csv('/kaggle/input/grammar-corrected-persuade-10k/corrected_persuade_10k.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0dfa3f6f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:01:58.226537Z",
     "iopub.status.busy": "2023-11-10T10:01:58.226242Z",
     "iopub.status.idle": "2023-11-10T10:01:58.778845Z",
     "shell.execute_reply": "2023-11-10T10:01:58.777977Z"
    },
    "papermill": {
     "duration": 0.560215,
     "end_time": "2023-11-10T10:01:58.781172",
     "exception": false,
     "start_time": "2023-11-10T10:01:58.220957",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1638, 4, 361.42858331946064)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aug_df = pd.DataFrame({'text': aug_df['text'], 'generated': 1})\n",
    "proper_train_ai = pd.DataFrame({'text': proper_train['text'][proper_train['label']==1], 'generated': 1}).reset_index(drop=True)\n",
    "df_nbroad = pd.concat([df_llama, df_falcon]).reset_index(drop=True)\n",
    "nbroad_ai = pd.DataFrame({'text': df_nbroad['generated_text'], 'generated': 1}).reset_index(drop=True)\n",
    "ai_ds = pd.concat([proper_train_ai, nbroad_ai, ai_samples_500.drop(columns=['prompt_id', 'id']), aug_df]).reset_index(drop=True)\n",
    "human_ds = pd.DataFrame({'text': corrected_df['corrected_text'], 'generated': 0}).reset_index(drop=True)\n",
    "final_df = pd.concat([ai_ds, human_ds]).sample(frac=1).reset_index(drop=True)\n",
    "final_df\n",
    "\n",
    "lengths = []\n",
    "for sample in final_df['text']:\n",
    "    lengths.append(len(sample.split(' ')))\n",
    "max(lengths), min(lengths), sum(lengths) / len(lengths) # Rough estimation of lengths of sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0d54107b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:01:58.790837Z",
     "iopub.status.busy": "2023-11-10T10:01:58.790481Z",
     "iopub.status.idle": "2023-11-10T10:01:59.376384Z",
     "shell.execute_reply": "2023-11-10T10:01:59.375089Z"
    },
    "papermill": {
     "duration": 0.593958,
     "end_time": "2023-11-10T10:01:59.379250",
     "exception": false,
     "start_time": "2023-11-10T10:01:58.785292",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
    "tokenizer.pad_token = tokenizer.eos_token # Honestly like is pad token ever set to something else than EOS?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "429e3f81",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:01:59.390348Z",
     "iopub.status.busy": "2023-11-10T10:01:59.389483Z",
     "iopub.status.idle": "2023-11-10T10:01:59.395951Z",
     "shell.execute_reply": "2023-11-10T10:01:59.395103Z"
    },
    "papermill": {
     "duration": 0.014234,
     "end_time": "2023-11-10T10:01:59.397962",
     "exception": false,
     "start_time": "2023-11-10T10:01:59.383728",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def preprocess_function(examples):\n",
    "    return tokenizer(examples['text'], max_length=MAX_INPUT, padding='max_length', truncation=True) # We could try `padding_to_multiple_of=128`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b2b49958",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:01:59.407331Z",
     "iopub.status.busy": "2023-11-10T10:01:59.407041Z",
     "iopub.status.idle": "2023-11-10T10:02:57.096193Z",
     "shell.execute_reply": "2023-11-10T10:02:57.095099Z"
    },
    "papermill": {
     "duration": 57.696515,
     "end_time": "2023-11-10T10:02:57.098461",
     "exception": false,
     "start_time": "2023-11-10T10:01:59.401946",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map (num_proc=96): 100%|██████████| 20423/20423 [00:28<00:00, 729.29 examples/s]\n",
      "/usr/local/lib/python3.8/site-packages/datasets/table.py:1395: FutureWarning: promote has been superseded by mode='default'.\n",
      "  block_group = [InMemoryTable(cls._concat_blocks(list(block_group), axis=axis))]\n",
      "/usr/local/lib/python3.8/site-packages/datasets/table.py:1421: FutureWarning: promote has been superseded by mode='default'.\n",
      "  table = cls._concat_blocks(blocks, axis=0)\n",
      "Map (num_proc=96): 100%|██████████| 3605/3605 [00:25<00:00, 141.27 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['label', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 20423\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['label', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 3605\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = Dataset.from_pandas(final_df).rename_column(\"generated\", \"label\").train_test_split(test_size=0.15)\n",
    "\n",
    "ds['train'] = ds['train'].map(preprocess_function, batched=False, num_proc=96, remove_columns=['text']) # Kaggle TPU session's CPUs are copiously powerfull\n",
    "ds['test'] = ds['test'].map(preprocess_function, batched=False, num_proc=96, remove_columns=['text'])\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "09f4c12c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:02:57.141286Z",
     "iopub.status.busy": "2023-11-10T10:02:57.140662Z",
     "iopub.status.idle": "2023-11-10T10:07:49.334991Z",
     "shell.execute_reply": "2023-11-10T10:07:49.334076Z"
    },
    "papermill": {
     "duration": 292.231553,
     "end_time": "2023-11-10T10:07:49.351204",
     "exception": false,
     "start_time": "2023-11-10T10:02:57.119651",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 3/3 [02:41<00:00, 53.85s/it]\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2, torch_dtype=torch.bfloat16) # Some may argue training in pure bfloat16 is bad but it's fine for most purposes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c43190f1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:07:49.390919Z",
     "iopub.status.busy": "2023-11-10T10:07:49.390605Z",
     "iopub.status.idle": "2023-11-10T10:07:49.395249Z",
     "shell.execute_reply": "2023-11-10T10:07:49.394480Z"
    },
    "papermill": {
     "duration": 0.026933,
     "end_time": "2023-11-10T10:07:49.397214",
     "exception": false,
     "start_time": "2023-11-10T10:07:49.370281",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "FLAGS = {'MAX_INPUT': MAX_INPUT,\n",
    "         'LOGGING_STEPS': 100,\n",
    "         'NUM_EPOCHS': 1,\n",
    "         'BATCH_SIZE': 8, #Making batch_size lower then 8 will result in slower training, but will allow for larger models\\context. Fortunately, we have 128GBs. Setting higher batch_size doesn't seem to improve time.\n",
    "          'NUM_STEPS': len(ds)} "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "965c985c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:07:49.436403Z",
     "iopub.status.busy": "2023-11-10T10:07:49.436167Z",
     "iopub.status.idle": "2023-11-10T10:07:49.441318Z",
     "shell.execute_reply": "2023-11-10T10:07:49.440454Z"
    },
    "papermill": {
     "duration": 0.02713,
     "end_time": "2023-11-10T10:07:49.443198",
     "exception": false,
     "start_time": "2023-11-10T10:07:49.416068",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "training_loader = torch.utils.data.DataLoader(ds['train'], batch_size=FLAGS['BATCH_SIZE'], collate_fn=DataCollatorWithPadding(tokenizer=tokenizer), shuffle=True) # We could reinvent the wheel by writing our own data collator\n",
    "testing_loader = torch.utils.data.DataLoader(ds['test'], batch_size=FLAGS['BATCH_SIZE'], collate_fn=DataCollatorWithPadding(tokenizer=tokenizer))\n",
    "\n",
    "device = xm.xla_device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a887cc0e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:07:49.482944Z",
     "iopub.status.busy": "2023-11-10T10:07:49.482667Z",
     "iopub.status.idle": "2023-11-10T10:08:21.292986Z",
     "shell.execute_reply": "2023-11-10T10:08:21.292080Z"
    },
    "papermill": {
     "duration": 31.833317,
     "end_time": "2023-11-10T10:08:21.295603",
     "exception": false,
     "start_time": "2023-11-10T10:07:49.462286",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "cnt = 0\n",
    "for param in model.parameters(): # Very direct way of freezing parameters\n",
    "    cnt += 1\n",
    "    param.requires_grad = False\n",
    "    if cnt > 260:\n",
    "        param.requires_grad = True\n",
    "\n",
    "config = AutoConfig.from_pretrained(MODEL)\n",
    "num_devices = xr.global_runtime_device_count()\n",
    "mesh_shape = (1, num_devices, 1)\n",
    "device_ids = np.array(range(num_devices))\n",
    "mesh = Mesh(device_ids, mesh_shape, ('dp', 'fsdp', 'mp'))\n",
    "partition_module(model, mesh) # After this, the model is sharded between TPU cores but still has the same API as if it was on single device. Neat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4a389a56",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:08:21.337747Z",
     "iopub.status.busy": "2023-11-10T10:08:21.337374Z",
     "iopub.status.idle": "2023-11-10T10:08:23.225335Z",
     "shell.execute_reply": "2023-11-10T10:08:23.223953Z"
    },
    "papermill": {
     "duration": 1.911711,
     "end_time": "2023-11-10T10:08:23.227770",
     "exception": false,
     "start_time": "2023-11-10T10:08:21.316059",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!export XLA_USE_BF16=1 # Idk if this does anything\n",
    "num_iterations = int(FLAGS['NUM_STEPS'] / FLAGS['BATCH_SIZE'])\n",
    "lr = 1e-5\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.1, total_iters=FLAGS['NUM_STEPS'] * FLAGS['BATCH_SIZE'])\n",
    "\n",
    "def training_loop(training_loader, epoch):\n",
    "    model.train()\n",
    "    print('Epoch {} train begin {}'.format(epoch, test_utils.now()))\n",
    "    for step, batch in enumerate(training_loader):\n",
    "        optimizer.zero_grad()\n",
    "        input_ids, attention_mask, labels = batch.input_ids.to(device), batch.attention_mask.to(device),  torch.unsqueeze(batch.labels.to(device), 1) # Here we inhale some copium\n",
    "        xs.mark_sharding(input_ids, mesh, (0, 1)) # Sharding inputs\n",
    "        xs.mark_sharding(attention_mask, mesh, (0, 1))\n",
    "        xs.mark_sharding(labels, mesh, (0, 1))\n",
    "        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)\n",
    "        loss = outputs.loss\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        xm.mark_step()\n",
    "        if step % FLAGS['LOGGING_STEPS'] == 0:\n",
    "            print(f'loss: {loss.item()}, time: {test_utils.now()}, step: {step}')\n",
    "        scheduler.step()\n",
    "    print('Epoch {} train end {}'.format(epoch, test_utils.now()))\n",
    "\n",
    "def testing_loop(testing_loader, epoch):\n",
    "    model.eval()\n",
    "    total_loss = 0.0\n",
    "    total_steps = 0\n",
    "    with torch.no_grad():\n",
    "        for step, batch in enumerate(testing_loader):\n",
    "            input_ids, attention_mask, labels = batch.input_ids.to(device), batch.attention_mask.to(device), torch.unsqueeze(batch.labels.to(device), 1) # Even more copium\n",
    "            xs.mark_sharding(input_ids, mesh, (0, 1))\n",
    "            xs.mark_sharding(attention_mask, mesh, (0, 1))\n",
    "            xs.mark_sharding(labels, mesh, (0, 1))\n",
    "            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)\n",
    "            loss = outputs.loss\n",
    "            total_loss += loss.item()\n",
    "            total_steps += 1\n",
    "    average_loss = total_loss / total_steps\n",
    "    print('Epoch {} test end {}, test loss={:.2f}'.format(epoch, test_utils.now(), average_loss))\n",
    "\n",
    "def train(FLAGS):\n",
    "    for epoch in range(1, FLAGS['NUM_EPOCHS'] + 1):\n",
    "        training_loop(training_loader, epoch)\n",
    "        testing_loop(testing_loader, epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f5933ccb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T10:08:23.270636Z",
     "iopub.status.busy": "2023-11-10T10:08:23.270284Z",
     "iopub.status.idle": "2023-11-10T10:44:46.017760Z",
     "shell.execute_reply": "2023-11-10T10:44:46.016290Z"
    },
    "papermill": {
     "duration": 2182.771275,
     "end_time": "2023-11-10T10:44:46.019612",
     "exception": true,
     "start_time": "2023-11-10T10:08:23.248337",
     "status": "failed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 train begin 10:08:23\n",
      "loss: 2.484375, time: 10:09:20, step: 0\n",
      "loss: 0.6015625, time: 10:11:06, step: 100\n",
      "loss: 0.63671875, time: 10:12:17, step: 200\n",
      "loss: 1.3046875, time: 10:13:27, step: 300\n",
      "loss: 0.28515625, time: 10:14:38, step: 400\n",
      "loss: 0.0947265625, time: 10:15:48, step: 500\n",
      "loss: 0.7265625, time: 10:16:58, step: 600\n",
      "loss: 0.375, time: 10:18:09, step: 700\n",
      "loss: 1.2421875, time: 10:19:19, step: 800\n",
      "loss: 0.244140625, time: 10:20:30, step: 900\n",
      "loss: 0.55859375, time: 10:21:40, step: 1000\n",
      "loss: 0.474609375, time: 10:22:51, step: 1100\n",
      "loss: 0.6875, time: 10:24:01, step: 1200\n",
      "loss: 0.302734375, time: 10:25:11, step: 1300\n",
      "loss: 0.138671875, time: 10:26:22, step: 1400\n",
      "loss: 0.56640625, time: 10:27:32, step: 1500\n",
      "loss: 0.5, time: 10:28:43, step: 1600\n",
      "loss: 0.62109375, time: 10:29:53, step: 1700\n",
      "loss: 1.1953125, time: 10:31:04, step: 1800\n",
      "loss: 0.349609375, time: 10:32:14, step: 1900\n",
      "loss: 0.53125, time: 10:33:24, step: 2000\n",
      "loss: 0.64453125, time: 10:34:35, step: 2100\n",
      "loss: 0.47265625, time: 10:35:45, step: 2200\n",
      "loss: 0.4140625, time: 10:36:56, step: 2300\n",
      "loss: 0.443359375, time: 10:38:06, step: 2400\n",
      "loss: 0.365234375, time: 10:39:17, step: 2500\n",
      "Epoch 1 train end 10:40:31\n",
      "Epoch 1 test end 10:44:45, test loss=0.45\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'average_loss' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mFLAGS\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[12], line 47\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(FLAGS)\u001b[0m\n\u001b[1;32m     45\u001b[0m     training_loop(training_loader, epoch)\n\u001b[1;32m     46\u001b[0m     testing_loop(testing_loader, epoch)\n\u001b[0;32m---> 47\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mTraining done\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(epoch, test_utils\u001b[38;5;241m.\u001b[39mnow(), \u001b[43maverage_loss\u001b[49m))\n",
      "\u001b[0;31mNameError\u001b[0m: name 'average_loss' is not defined"
     ]
    }
   ],
   "source": [
    "train(FLAGS) # Please ignore the error, it doesn't exist anymore. Just look at the training loss, it decreases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1295897e",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# We can just download the model directly but Kaggle hardware has little disk storage so I had to do this\n",
    "from kaggle_secrets import UserSecretsClient\n",
    "from huggingface_hub import login\n",
    "\n",
    "user_secrets = UserSecretsClient()\n",
    "hf_token = user_secrets.get_secret(\"hf_write\") # Provide your own HF API token with write access\n",
    "login(hf_token)\n",
    "\n",
    "model = model.cpu() # We have to move the model onto CPU to push it to HF.\n",
    "print('now saving the model')\n",
    "model.push_to_hub(\n",
    "    \"Defetya/example_repo_detectAI\", \n",
    "    tokenizer=tokenizer,\n",
    "    private=False,\n",
    "    create_pr=False,\n",
    "    max_shard_size=\"2GB\", # Sharding isn't as important as before since hardware is better now but who cares anyway\n",
    "    )# We have to push the model to HF since there is not enough memory on disk. Download weights from there"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.17"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 2718.743894,
   "end_time": "2023-11-10T10:44:52.886558",
   "environment_variables": {},
   "exception": true,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2023-11-10T09:59:34.142664",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}