train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TODOs\n",
    "\n",
    "- Figure out why resumption is not working\n",
    "- Verify that loss function is implemented right \n",
    "- record / plot loss better and easy background training\n",
    "- [X] Research more QLoRA\n",
    "  - [ ] make sure that QLoRA is actually happening..\n",
    "  - RTFP\n",
    "  - [X] Find good solid example code with transformers\n",
    "  - Determine QLoRA parameters to experiment with and how\n",
    "    - datatype\n",
    "    - what bits\n",
    "    - r\n",
    "    - dropout\n",
    "    - target_modules\n",
    "    - lora alpha\n",
    "- Determine training parameters to experiment with\n",
    "  - AdamW RTFP\n",
    "  - tune adamW starting learning rate?\n",
    "  - Batch size. Should it vary over time?\n",
    "  - \n",
    "- get complicated and weird\n",
    "  - explore mining hard negatives\n",
    "  - explore synthesizing hard negatives\n",
    "  - synthesize more varied data\n",
    "  - generated data off of a target domain (easy-on fine tuning)\n",
    "- benchmark\n",
    "  - add inferrence code path\n",
    "  - figure out how to run model against an mteb benchmark\n",
    "- update code to be able to run it distributed in the cloud?\n",
    "  - related to deepspeed and accelerate\n",
    "    - https://arxiv.org/abs/2104.07857\n",
    "    - https://huggingface.co/docs/accelerate/usage_guides/deepspeed\n",
    "- [ ] get a bigger graphics card"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/freezer/dev/tinyllama-embed/.venv/lib/python3.10/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": [
    "from transformers import LlamaModel, AutoTokenizer, BitsAndBytesConfig\n",
    "from datasets import load_dataset\n",
    "import torch\n",
    "\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 = LlamaModel.from_pretrained(\n",
    "    \"TinyLlama/TinyLlama-1.1B-intermediate-step-1195k-token-2.5T\",\n",
    "    device_map=\"auto\",\n",
    "    quantization_config=bnb_config,\n",
    "    # load_in_4bit=True,\n",
    "    torch_dtype=torch.float16,\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    \"TinyLlama/TinyLlama-1.1B-intermediate-step-1195k-token-2.5T\"\n",
    ")\n",
    "\n",
    "base_dataset = load_dataset(\n",
    "    \"andersonbcdefg/synthetic_tuples_gpt35_turbo\", split=\"train\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import TaskType\n",
    "\n",
    "\n",
    "# consider and experiment withadding a specific pad token\n",
    "# tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "tokenizer.pad_token = tokenizer.unk_token\n",
    "tokenizer.pad_token_id = tokenizer.unk_token_id\n",
    "\n",
    "# Tokenize the dataset\n",
    "def tokenize_function(examples):\n",
    "    max_len = 128 # TODO - reconsider this\n",
    "    tokenized_query = tokenizer(examples[\"query\"], padding=True, truncation=True, max_length=max_len, return_tensors=\"pt\")\n",
    "    tokenized_pos = tokenizer(examples[\"pos\"], padding=True,  truncation=True, max_length=max_len, return_tensors=\"pt\")\n",
    "    tokenized_neg = tokenizer(examples[\"neg\"], padding=True, truncation=True, max_length=max_len, return_tensors=\"pt\")\n",
    "    return {\n",
    "        \"input_ids_query\": tokenized_query[\"input_ids\"],\n",
    "        \"attention_mask_query\": tokenized_query[\"attention_mask\"],\n",
    "        \"input_ids_pos\": tokenized_pos[\"input_ids\"],\n",
    "        \"attention_mask_pos\": tokenized_pos[\"attention_mask\"],\n",
    "        \"input_ids_neg\": tokenized_neg[\"input_ids\"],\n",
    "        \"attention_mask_neg\": tokenized_neg[\"attention_mask\"],\n",
    "    }\n",
    "\n",
    "\n",
    "dataset = base_dataset.map(tokenize_function, batched=True, cache_file_name=\"./cache/tokenized_datasets\")\n",
    "\n",
    "train_test_split = dataset.train_test_split(test_size=0.2, seed=seed)\n",
    "\n",
    "# Access the new train and test datasets\n",
    "train_dataset = train_test_split[\"train\"]\n",
    "eval_dataset = train_test_split[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    task_type=TaskType.FEATURE_EXTRACTION,\n",
    "    r=8,\n",
    "    lora_alpha=64,\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=[\n",
    "        \"q_proj\",\n",
    "        \"v_proj\",\n",
    "        \"o_proj\",\n",
    "        \"down_proj\",\n",
    "        \"up_proj\",\n",
    "        \"gate_proj\",\n",
    "    ],\n",
    "    inference_mode=False,\n",
    ")\n",
    "model = prepare_model_for_kbit_training(model, use_gradient_checkpointing = True)\n",
    "model = get_peft_model(model, peft_config)\n",
    "# model.add_adapter(peft_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments, IntervalStrategy\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from datasets import Dataset\n",
    "\n",
    "\n",
    "class TinyEmbedTrainer(Trainer):\n",
    "    \"\"\"\n",
    "    Subclasses hugging face trainer to do all the needed things:\n",
    "    - embedding training needs multiple inputs to train on (query, positive, negative).\n",
    "    -   collate function to reshape the dict and pad the batch\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        args: TrainingArguments,\n",
    "        train_dataset: Dataset,\n",
    "        eval_dataset: Dataset,\n",
    "        tokenizer: AutoTokenizer,\n",
    "    ):\n",
    "        # Consider reworking the model's signature to conform to training expectations\n",
    "        args.remove_unused_columns = False\n",
    "        super().__init__(\n",
    "            model,\n",
    "            args,\n",
    "            train_dataset=train_dataset,\n",
    "            eval_dataset=eval_dataset,\n",
    "            tokenizer=tokenizer,\n",
    "            data_collator=self._map_collate_fn,\n",
    "        )\n",
    "        # self.train_dataset = train_dataset\n",
    "        # self.eval_dataset = eval_dataset\n",
    "        # self.tokenizer = tokenizer\n",
    "\n",
    "    def _map_collate_fn(self, batch):\n",
    "        \"\"\"\n",
    "        pads and truncates the batch for each of the three inputs\n",
    "        and returns a single dict, with each key \n",
    "        shape the the List[Dict[str, torch.Tensor]] into a Dict[str, torch.Tensor]\n",
    "        \n",
    "        \"\"\"\n",
    "        def get_field(field: str):\n",
    "            attention = [item[\"attention_mask_\" + field] for item in batch]\n",
    "            inputs = [item[\"input_ids_\" + field] for item in batch]\n",
    "            # probably something wrong at the dataset level... but munge things into the right shape here,\n",
    "            # truncate the batch to the max length, but also 0 pad those that are less than the max length\n",
    "            # for both input_ids and attention_mask\n",
    "            max_len = max([sum(x) for x in attention])\n",
    "            attention_mask = []\n",
    "            input_ids = []\n",
    "            for i in range(len(attention)):\n",
    "                attention_i = attention[i]\n",
    "                input_ids_i = inputs[i]\n",
    "                attention_trunc = attention_i[:max_len]\n",
    "                input_ids_trunc = input_ids_i[:max_len]\n",
    "                attention_pad = attention_trunc + [0] * (max_len - len(attention_trunc))\n",
    "                input_ids_pad = input_ids_trunc + [self.tokenizer.pad_token_id] * (\n",
    "                    max_len - len(input_ids_trunc)\n",
    "                )\n",
    "                attention_mask.append(attention_pad)\n",
    "                input_ids.append(input_ids_pad)\n",
    "\n",
    "            return {\n",
    "                (\"attention_mask_\" + field): torch.tensor(attention_mask),\n",
    "                (\"input_ids_\" + field): torch.tensor(input_ids),\n",
    "            }\n",
    "        return {\n",
    "            **get_field(\"query\"), **get_field(\"pos\"), **get_field(\"neg\")\n",
    "        }\n",
    "\n",
    "    def _get_embeddings(self, model, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:\n",
    "        outputs = model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)\n",
    "        hidden_states = outputs.hidden_states[-1]  # Get the last hidden states\n",
    "        # Compute the indices of the last non-padding tokens\n",
    "        sequence_lengths = attention_mask.sum(dim=1)\n",
    "        last_token_indices = sequence_lengths - 1\n",
    "        embeddings = hidden_states[\n",
    "            torch.arange(hidden_states.size(0)), last_token_indices, :\n",
    "        ]\n",
    "        # Normalize the embeddings\n",
    "        embeddings = F.normalize(embeddings, p=2, dim=1)\n",
    "        return embeddings\n",
    "    \n",
    "    def compute_loss(self, model, inputs, return_outputs=False):\n",
    "        query_embeddings = self._get_embeddings(model, inputs[\"input_ids_query\"], inputs[\"attention_mask_query\"])\n",
    "        pos_embeddings = self._get_embeddings(model, inputs[\"input_ids_pos\"], inputs[\"attention_mask_pos\"])\n",
    "        neg_embeddings = self._get_embeddings(model, inputs[\"input_ids_neg\"], inputs[\"attention_mask_neg\"])\n",
    "\n",
    "        # Compute InfoNCE loss\n",
    "        pos_similarity = torch.sum(query_embeddings * pos_embeddings, dim=1)\n",
    "        neg_similarity = torch.sum(query_embeddings * neg_embeddings, dim=1)\n",
    "        losses = -torch.log(\n",
    "            torch.exp(pos_similarity)\n",
    "            / (torch.exp(pos_similarity) + torch.exp(neg_similarity))\n",
    "        )\n",
    "        return losses.mean()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading model from ./results/checkpoint-1300.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running training *****\n",
      "  Num examples = 163,636\n",
      "  Num Epochs = 3\n",
      "  Instantaneous batch size per device = 16\n",
      "  Training with DataParallel so batch size has been adjusted to: 1\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 1\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 490,908\n",
      "  Number of trainable parameters = 5,902,336\n",
      "  Continuing training from checkpoint, will skip to saved global_step\n",
      "  Continuing training from epoch 0\n",
      "  Continuing training from global step 1300\n",
      "  Will skip the first 0 epochs then the first 1300 batches in the first epoch.\n",
      "/mnt/freezer/dev/tinyllama-embed/.venv/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='1407' max='490908' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [  1407/490908 01:31 < 118:32:31, 1.15 it/s, Epoch 0.01/3]\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>1350</td>\n",
       "      <td>0.280100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>0.349800</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to ./results/tmp-checkpoint-1400\n",
      "tokenizer config file saved in ./results/tmp-checkpoint-1400/tokenizer_config.json\n",
      "Special tokens file saved in ./results/tmp-checkpoint-1400/special_tokens_map.json\n",
      "/mnt/freezer/dev/tinyllama-embed/.venv/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"
     ]
    }
   ],
   "source": [
    "\n",
    "# Training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./results\",\n",
    "    per_device_train_batch_size=16,\n",
    "    num_train_epochs=3,\n",
    "    logging_dir=\"./logs\",\n",
    "    log_level=\"info\",\n",
    "    logging_strategy=IntervalStrategy.STEPS,\n",
    "    gradient_checkpointing=True,\n",
    "    logging_steps=5,\n",
    "    save_steps=100\n",
    ")\n",
    "\n",
    "trainer = TinyEmbedTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "trainer.train(resume_from_checkpoint=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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": 2
}