qlora-minimal.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e7bec93-c233-4493-b5d7-23d06f02d218",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U accelerate bitsandbytes datasets peft transformers"
   ]
  },
  {
   "cell_type": "raw",
   "id": "e64aabb0-25ef-455f-9494-b6f4ca3ecfc9",
   "metadata": {},
   "source": [
    "1. Load Open Assistant dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fe883fe8-868c-4cc8-92e5-ed9889143ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"OpenAssistant/oasst_top1_2023-08-25\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7aab5773-415c-4c90-904a-f5c0a755abfa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text'],\n",
       "        num_rows: 12947\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['text'],\n",
       "        num_rows: 690\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f2b666c7-108b-4a64-aaec-64dfb1b10078",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>user\n",
      "Consigliami 5 nomi per il mio cucciolo di dobberman<|im_end|>\n",
      "<|im_start|>assistant\n",
      "Ecco 5 nomi per il tuo cucciolo di dobermann:\n",
      "\n",
      "- Zeus\n",
      "- Apollo\n",
      "- Thor\n",
      "- Athena\n",
      "- Odin<|im_end|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(dataset[\"train\"][0][\"text\"])"
   ]
  },
  {
   "cell_type": "raw",
   "id": "c851ab41-4a7b-4df1-8140-48623daeae99",
   "metadata": {},
   "source": [
    "2. Load and prepare model and tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a8ef7f26-f87b-4c54-924f-c9661bc1bf2f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "86846f6983d24ab4bcdd61644011bdd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, BitsAndBytesConfig\n",
    "from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model\n",
    "\n",
    "modelpath=\"models/Mistral-7B-v0.1\"\n",
    "\n",
    "# Load 4-bit quantized model\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    modelpath,    \n",
    "    device_map=\"auto\",\n",
    "    quantization_config=BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "    ),\n",
    "    torch_dtype=torch.bfloat16,\n",
    ")\n",
    "\n",
    "# Load (slow) Tokenizer, fast tokenizer sometimes ignores added tokens\n",
    "tokenizer = AutoTokenizer.from_pretrained(modelpath, use_fast=False)   \n",
    "\n",
    "# Add tokens <|im_start|> and <|im_end|>, latter is special eos token \n",
    "tokenizer.pad_token = \"</s>\"\n",
    "tokenizer.add_tokens([\"<|im_start|>\"])\n",
    "tokenizer.add_special_tokens(dict(eos_token=\"<|im_end|>\"))\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "model.config.eos_token_id = tokenizer.eos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "948e0d32-60ee-4e0e-9cf5-f17970e4e110",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add LoRA adapters to model\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "config = LoraConfig(\n",
    "    r=64, \n",
    "    lora_alpha=16, \n",
    "    target_modules = ['q_proj', 'k_proj', 'down_proj', 'v_proj', 'gate_proj', 'o_proj', 'up_proj'],\n",
    "    lora_dropout=0.1, \n",
    "    bias=\"none\", \n",
    "    modules_to_save = [\"lm_head\", \"embed_tokens\"],\t\t# needed because we added new tokens to tokenizer/model\n",
    "    task_type=\"CAUSAL_LM\"\n",
    ")\n",
    "model = get_peft_model(model, config)\n",
    "model.config.use_cache = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e7211c68-07a3-4d24-a7af-a3691063a758",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): MistralForCausalLM(\n",
       "      (model): MistralModel(\n",
       "        (embed_tokens): ModulesToSaveWrapper(\n",
       "          (original_module): Embedding(32002, 4096)\n",
       "          (modules_to_save): ModuleDict(\n",
       "            (default): Embedding(32002, 4096)\n",
       "          )\n",
       "        )\n",
       "        (layers): ModuleList(\n",
       "          (0-31): 32 x MistralDecoderLayer(\n",
       "            (self_attn): MistralAttention(\n",
       "              (q_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "              )\n",
       "              (k_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=1024, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "              )\n",
       "              (v_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=1024, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "              )\n",
       "              (o_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "              )\n",
       "              (rotary_emb): MistralRotaryEmbedding()\n",
       "            )\n",
       "            (mlp): MistralMLP(\n",
       "              (gate_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=14336, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "              )\n",
       "              (up_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=14336, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "              )\n",
       "              (down_proj): Linear4bit(\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=14336, out_features=64, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "                (base_layer): Linear4bit(in_features=14336, out_features=4096, bias=False)\n",
       "              )\n",
       "              (act_fn): SiLUActivation()\n",
       "            )\n",
       "            (input_layernorm): MistralRMSNorm()\n",
       "            (post_attention_layernorm): MistralRMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): MistralRMSNorm()\n",
       "      )\n",
       "      (lm_head): ModulesToSaveWrapper(\n",
       "        (original_module): Linear(in_features=4096, out_features=32002, bias=False)\n",
       "        (modules_to_save): ModuleDict(\n",
       "          (default): Linear(in_features=4096, out_features=32002, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "raw",
   "id": "9913bc18-8a26-4a1b-8abc-3bc8e672f191",
   "metadata": {},
   "source": [
    "3. Prepare data for Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "92a941dc-9e9e-4bbd-9c2a-a54f2fd72071",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "408be2721e3b4b4d82aaeec5a1d7e415",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=32):   0%|          | 0/12947 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dca702d381104f0fa97ea52645ab4d29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=32):   0%|          | 0/690 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os \n",
    "\n",
    "def tokenize(element):\n",
    "    return tokenizer(\n",
    "        element[\"text\"],\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "\n",
    "dataset_tokenized = dataset.map(\n",
    "    tokenize, \n",
    "    batched=True, \n",
    "    num_proc=os.cpu_count(),    # multithreaded\n",
    "    remove_columns=[\"text\"]     # don't need this anymore, we have tokens from here on\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c427e534-2214-4bc6-8c73-4a81c4984db5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask'],\n",
       "        num_rows: 12947\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask'],\n",
       "        num_rows: 690\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_tokenized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "aaae4749-0fd9-4db1-a8b7-33ad164348b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define collate function - transform list of dictionaries [ {input_ids: [123, ..]}, {.. ] to single batch dictionary { input_ids: [..], labels: [..], attention_mask: [..] }\n",
    "def collate(elements):\n",
    "    tokenlist=[e[\"input_ids\"] for e in elements]\n",
    "    tokens_maxlen=max([len(t) for t in tokenlist])\n",
    "\n",
    "    input_ids,labels,attention_masks = [],[],[]\n",
    "    for tokens in tokenlist:\n",
    "        pad_len=tokens_maxlen-len(tokens)\n",
    "\n",
    "        # pad input_ids with pad_token, labels with ignore_index (-100) and set attention_mask 1 where content otherwise 0\n",
    "        input_ids.append( tokens + [tokenizer.pad_token_id]*pad_len )   \n",
    "        labels.append( tokens + [-100]*pad_len )    \n",
    "        attention_masks.append( [1]*len(tokens) + [0]*pad_len ) \n",
    "\n",
    "    batch={\n",
    "        \"input_ids\": torch.tensor(input_ids),\n",
    "        \"labels\": torch.tensor(labels),\n",
    "        \"attention_mask\": torch.tensor(attention_masks)\n",
    "    }\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "raw",
   "id": "27287f0e-56f0-458f-8c2e-9124a9739d48",
   "metadata": {},
   "source": [
    " Training Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e36b9ad9-d0c0-4dc7-a4c8-03765891dac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "bs=1        # batch size\n",
    "ga_steps=1  # gradient acc. steps\n",
    "epochs=5\n",
    "steps_per_epoch=len(dataset_tokenized[\"train\"])//(bs*ga_steps)\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"out\",\n",
    "    per_device_train_batch_size=bs,\n",
    "    per_device_eval_batch_size=bs,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    logging_steps=1,\n",
    "    eval_steps=steps_per_epoch,\t\t# eval and save once per epoch  \t\n",
    "    save_steps=steps_per_epoch,\n",
    "    gradient_accumulation_steps=ga_steps,\n",
    "    num_train_epochs=epochs,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    optim=\"paged_adamw_32bit\",\n",
    "    learning_rate=0.0002,\n",
    "    group_by_length=True,\n",
    "    fp16=True,\n",
    "    ddp_find_unused_parameters=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8e5d8b96-bcfe-490d-8462-dfb44d575432",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mg-ronimo\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.15.12 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.3"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/g/jupyter/wandb/run-20231105_185849-lrfiioan</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/g-ronimo/huggingface/runs/lrfiioan' target=\"_blank\">sunny-bush-127</a></strong> to <a href='https://wandb.ai/g-ronimo/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/g-ronimo/huggingface' target=\"_blank\">https://wandb.ai/g-ronimo/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/g-ronimo/huggingface/runs/lrfiioan' target=\"_blank\">https://wandb.ai/g-ronimo/huggingface/runs/lrfiioan</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/g/.local/lib/python3.10/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='5' max='8095' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [   5/8095 00:18 < 14:10:39, 0.16 it/s, Epoch 0.00/5]\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",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[14], line 10\u001b[0m\n\u001b[1;32m      1\u001b[0m trainer \u001b[38;5;241m=\u001b[39m Trainer(\n\u001b[1;32m      2\u001b[0m     model\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[1;32m      3\u001b[0m     tokenizer\u001b[38;5;241m=\u001b[39mtokenizer,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m      7\u001b[0m     args\u001b[38;5;241m=\u001b[39margs,\n\u001b[1;32m      8\u001b[0m )\n\u001b[0;32m---> 10\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/transformers/trainer.py:1555\u001b[0m, in \u001b[0;36mTrainer.train\u001b[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[1;32m   1553\u001b[0m         hf_hub_utils\u001b[38;5;241m.\u001b[39menable_progress_bars()\n\u001b[1;32m   1554\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1555\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1556\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1557\u001b[0m \u001b[43m        \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1558\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtrial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1559\u001b[0m \u001b[43m        \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1560\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/transformers/trainer.py:1860\u001b[0m, in \u001b[0;36mTrainer._inner_training_loop\u001b[0;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[1;32m   1857\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontrol \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcallback_handler\u001b[38;5;241m.\u001b[39mon_step_begin(args, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontrol)\n\u001b[1;32m   1859\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maccelerator\u001b[38;5;241m.\u001b[39maccumulate(model):\n\u001b[0;32m-> 1860\u001b[0m     tr_loss_step \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1862\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m   1863\u001b[0m     args\u001b[38;5;241m.\u001b[39mlogging_nan_inf_filter\n\u001b[1;32m   1864\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_torch_tpu_available()\n\u001b[1;32m   1865\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m (torch\u001b[38;5;241m.\u001b[39misnan(tr_loss_step) \u001b[38;5;129;01mor\u001b[39;00m torch\u001b[38;5;241m.\u001b[39misinf(tr_loss_step))\n\u001b[1;32m   1866\u001b[0m ):\n\u001b[1;32m   1867\u001b[0m     \u001b[38;5;66;03m# if loss is nan or inf simply add the average of previous logged losses\u001b[39;00m\n\u001b[1;32m   1868\u001b[0m     tr_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m tr_loss \u001b[38;5;241m/\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mglobal_step \u001b[38;5;241m-\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_globalstep_last_logged)\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/transformers/trainer.py:2734\u001b[0m, in \u001b[0;36mTrainer.training_step\u001b[0;34m(self, model, inputs)\u001b[0m\n\u001b[1;32m   2732\u001b[0m         scaled_loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[1;32m   2733\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 2734\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maccelerator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2736\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m loss\u001b[38;5;241m.\u001b[39mdetach() \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mgradient_accumulation_steps\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/accelerate/accelerator.py:1987\u001b[0m, in \u001b[0;36mAccelerator.backward\u001b[0;34m(self, loss, **kwargs)\u001b[0m\n\u001b[1;32m   1985\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m   1986\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mscaler \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m-> 1987\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscaler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscale\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1988\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1989\u001b[0m     loss\u001b[38;5;241m.\u001b[39mbackward(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/_tensor.py:492\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    482\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    483\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    484\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    485\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    490\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    491\u001b[0m     )\n\u001b[0;32m--> 492\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    493\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    494\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/autograd/__init__.py:251\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    246\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    248\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    249\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    250\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 251\u001b[0m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    252\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    253\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    254\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    255\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    256\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    257\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    258\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    259\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error in callback <bound method _WandbInit._pause_backend of <wandb.sdk.wandb_init._WandbInit object at 0x7fb68349b280>> (for post_run_cell), with arguments args (<ExecutionResult object at 7fb458fd3d00, execution_count=14 error_before_exec=None error_in_exec= info=<ExecutionInfo object at 7fb458fd3f70, raw_cell=\"trainer = Trainer(\n",
      "    model=model,\n",
      "    tokenizer=..\" store_history=True silent=False shell_futures=True cell_id=8e5d8b96-bcfe-490d-8462-dfb44d575432> result=None>,),kwargs {}:\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "_WandbInit._pause_backend() takes 1 positional argument but 2 were given",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;31mTypeError\u001b[0m: _WandbInit._pause_backend() takes 1 positional argument but 2 were given"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=collate,\n",
    "    train_dataset=dataset_tokenized[\"train\"],\n",
    "    eval_dataset=dataset_tokenized[\"test\"],\n",
    "    args=args,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}