notebooks/gpt2/saving_and_loading.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "372a8b51",
   "metadata": {
    "papermill": {
     "duration": 0.00261,
     "end_time": "2024-03-24T00:36:44.761056",
     "exception": false,
     "start_time": "2024-03-24T00:36:44.758446",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Saving and loading\n",
    "The models created with *transformer_heads* generally integrate well with huggingface and will work with automatic saving/checkpointing during training using for example the *Trainer* class. However, during loading it has to be ensured that all heads are attached correctly and that their parameters (and qlora parameters) are loaded correctly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f81dce",
   "metadata": {},
   "source": [
    "**GPU Requirements:** For running with GPT-2 you may be fine with just 8GB of GPU RAM. With about 24GB you should be able to run any 7B or 13B model. With 80GB (A100) GPU you may be able to run a 70B model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "546d774f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:44.766569Z",
     "iopub.status.busy": "2024-03-24T00:36:44.766150Z",
     "iopub.status.idle": "2024-03-24T00:36:47.979987Z",
     "shell.execute_reply": "2024-03-24T00:36:47.979136Z"
    },
    "papermill": {
     "duration": 3.218439,
     "end_time": "2024-03-24T00:36:47.981765",
     "exception": false,
     "start_time": "2024-03-24T00:36:44.763326",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformer_heads import (\n",
    "    create_headed_qlora,\n",
    "    load_lora_with_heads,\n",
    "    HeadConfig,\n",
    "    load_headed,\n",
    "    get_multi_head_transformer,\n",
    ")\n",
    "from transformer_heads.util.helpers import get_model_params\n",
    "from transformers import BitsAndBytesConfig\n",
    "from peft import LoraConfig\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "97de8fa2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:47.987738Z",
     "iopub.status.busy": "2024-03-24T00:36:47.987341Z",
     "iopub.status.idle": "2024-03-24T00:36:47.990409Z",
     "shell.execute_reply": "2024-03-24T00:36:47.989880Z"
    },
    "papermill": {
     "duration": 0.00746,
     "end_time": "2024-03-24T00:36:47.991761",
     "exception": false,
     "start_time": "2024-03-24T00:36:47.984301",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# GPT2 is the fastest and requires fewest memory. However, this works just the same with any Llama or Mistral model. Just change model_path to its huggingface path.\n",
    "model_path = \"gpt2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1d5824a8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:48.005923Z",
     "iopub.status.busy": "2024-03-24T00:36:48.005555Z",
     "iopub.status.idle": "2024-03-24T00:36:48.225408Z",
     "shell.execute_reply": "2024-03-24T00:36:48.224845Z"
    },
    "papermill": {
     "duration": 0.223607,
     "end_time": "2024-03-24T00:36:48.226593",
     "exception": false,
     "start_time": "2024-03-24T00:36:48.002986",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'vocab_size': 50257, 'n_positions': 1024, 'n_embd': 768, 'n_layer': 12, 'n_head': 12, 'n_inner': None, 'activation_function': 'gelu_new', 'resid_pdrop': 0.1, 'embd_pdrop': 0.1, 'attn_pdrop': 0.1, 'layer_norm_epsilon': 1e-05, 'initializer_range': 0.02, 'summary_type': 'cls_index', 'summary_use_proj': True, 'summary_activation': None, 'summary_first_dropout': 0.1, 'summary_proj_to_labels': True, 'scale_attn_weights': True, 'use_cache': True, 'scale_attn_by_inverse_layer_idx': False, 'reorder_and_upcast_attn': False, 'bos_token_id': 50256, 'eos_token_id': 50256, 'return_dict': True, 'output_hidden_states': False, 'output_attentions': False, 'torchscript': False, 'torch_dtype': None, 'use_bfloat16': False, 'tf_legacy_loss': False, 'pruned_heads': {}, 'tie_word_embeddings': True, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'is_decoder': False, 'cross_attention_hidden_size': None, 'add_cross_attention': False, 'tie_encoder_decoder': False, 'max_length': 20, 'min_length': 0, 'do_sample': False, 'early_stopping': False, 'num_beams': 1, 'num_beam_groups': 1, 'diversity_penalty': 0.0, 'temperature': 1.0, 'top_k': 50, 'top_p': 1.0, 'typical_p': 1.0, 'repetition_penalty': 1.0, 'length_penalty': 1.0, 'no_repeat_ngram_size': 0, 'encoder_no_repeat_ngram_size': 0, 'bad_words_ids': None, 'num_return_sequences': 1, 'output_scores': False, 'return_dict_in_generate': False, 'forced_bos_token_id': None, 'forced_eos_token_id': None, 'remove_invalid_values': False, 'exponential_decay_length_penalty': None, 'suppress_tokens': None, 'begin_suppress_tokens': None, 'architectures': ['GPT2LMHeadModel'], 'finetuning_task': None, 'id2label': {0: 'LABEL_0', 1: 'LABEL_1'}, 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, 'tokenizer_class': None, 'prefix': None, 'pad_token_id': None, 'sep_token_id': None, 'decoder_start_token_id': None, 'task_specific_params': {'text-generation': {'do_sample': True, 'max_length': 50}}, 'problem_type': None, '_name_or_path': 'gpt2', 'transformers_version': '4.37.2', 'model_type': 'gpt2', 'n_ctx': 1024, 'model_class': <class 'transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel'>, 'hidden_size': 768}\n"
     ]
    }
   ],
   "source": [
    "model_params = get_model_params(model_path)\n",
    "model_class = model_params[\"model_class\"]\n",
    "hidden_size = model_params[\"hidden_size\"]\n",
    "vocab_size = model_params[\"vocab_size\"]\n",
    "print(model_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d138075",
   "metadata": {
    "papermill": {
     "duration": 0.002228,
     "end_time": "2024-03-24T00:36:48.231137",
     "exception": false,
     "start_time": "2024-03-24T00:36:48.228909",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's define some random head configs for later use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2e4e70e8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:48.236221Z",
     "iopub.status.busy": "2024-03-24T00:36:48.236024Z",
     "iopub.status.idle": "2024-03-24T00:36:48.240241Z",
     "shell.execute_reply": "2024-03-24T00:36:48.239673Z"
    },
    "papermill": {
     "duration": 0.008246,
     "end_time": "2024-03-24T00:36:48.241514",
     "exception": false,
     "start_time": "2024-03-24T00:36:48.233268",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "heads = [\n",
    "    HeadConfig(\n",
    "        name=\"lm_head\",\n",
    "        layer_hook=-1,\n",
    "        in_size=hidden_size,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=True,\n",
    "        loss_fct=\"cross_entropy\",\n",
    "        num_outputs=vocab_size,\n",
    "    ),\n",
    "    HeadConfig(\n",
    "        name=\"classification_hook\",\n",
    "        layer_hook=-4,\n",
    "        in_size=hidden_size,\n",
    "        hidden_size=1024,\n",
    "        num_layers=2,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=False,\n",
    "        loss_fct=\"cross_entropy\",\n",
    "        num_outputs=2,\n",
    "    ),\n",
    "    HeadConfig(\n",
    "        name=\"regression_hook\",\n",
    "        layer_hook=-6,\n",
    "        in_size=4096,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=False,\n",
    "        loss_fct=\"mse\",\n",
    "        num_outputs=1,\n",
    "        is_regression=True,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76eba1a3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:48.246845Z",
     "iopub.status.busy": "2024-03-24T00:36:48.246510Z",
     "iopub.status.idle": "2024-03-24T00:36:48.250408Z",
     "shell.execute_reply": "2024-03-24T00:36:48.249947Z"
    },
    "papermill": {
     "duration": 0.007756,
     "end_time": "2024-03-24T00:36:48.251511",
     "exception": false,
     "start_time": "2024-03-24T00:36:48.243755",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "quantization_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    load_in_8bit=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4322df51",
   "metadata": {
    "papermill": {
     "duration": 0.002314,
     "end_time": "2024-03-24T00:36:48.256053",
     "exception": false,
     "start_time": "2024-03-24T00:36:48.253739",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Saving and loading a transformer with attached linear probes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e33b6b76",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:48.261313Z",
     "iopub.status.busy": "2024-03-24T00:36:48.261132Z",
     "iopub.status.idle": "2024-03-24T00:36:50.639077Z",
     "shell.execute_reply": "2024-03-24T00:36:50.638548Z"
    },
    "papermill": {
     "duration": 2.382144,
     "end_time": "2024-03-24T00:36:50.640462",
     "exception": false,
     "start_time": "2024-03-24T00:36:48.258318",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of TransformerWithHeads were not initialized from the model checkpoint at gpt2 and are newly initialized: ['heads.classification_hook.lins.0.bias', 'heads.classification_hook.lins.0.weight', 'heads.classification_hook.lins.1.weight', 'heads.regression_hook.lins.0.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of TransformerWithHeads were not initialized from the model checkpoint at gpt2 and are newly initialized: ['heads.classification_hook.lins.0.bias', 'heads.classification_hook.lins.0.weight', 'heads.classification_hook.lins.1.weight', 'heads.regression_hook.lins.0.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# Create a quantized model with multiple heads\n",
    "model = load_headed(\n",
    "    model_class,\n",
    "    model_path,\n",
    "    heads,\n",
    "    device_map=\"cuda\",\n",
    "    quantization_config=quantization_config,\n",
    ")\n",
    "# Now you would do some training ...\n",
    "# Save the model now\n",
    "model.save_pretrained(\"test_model\")\n",
    "# Model is saved, delete it\n",
    "del model\n",
    "\n",
    "# With load_headed we can load the quantized model with the heads\n",
    "model = load_headed(\n",
    "    model_class,\n",
    "    model_path,\n",
    "    head_folder_path=\"test_model\",\n",
    "    device_map=\"cuda\",\n",
    "    quantization_config=quantization_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c98a563",
   "metadata": {
    "papermill": {
     "duration": 0.002421,
     "end_time": "2024-03-24T00:36:50.645786",
     "exception": false,
     "start_time": "2024-03-24T00:36:50.643365",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Saving and loading a model finetuned with qlora with extra heads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "049029fc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-24T00:36:50.651697Z",
     "iopub.status.busy": "2024-03-24T00:36:50.651446Z",
     "iopub.status.idle": "2024-03-24T00:36:53.162458Z",
     "shell.execute_reply": "2024-03-24T00:36:53.161901Z"
    },
    "papermill": {
     "duration": 2.515864,
     "end_time": "2024-03-24T00:36:53.164065",
     "exception": false,
     "start_time": "2024-03-24T00:36:50.648201",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of TransformerWithHeads were not initialized from the model checkpoint at gpt2 and are newly initialized: ['heads.classification_hook.lins.0.bias', 'heads.classification_hook.lins.0.weight', 'heads.classification_hook.lins.1.weight', 'heads.regression_hook.lins.0.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# Some simple LoRA config. target_modules=None will result in all linear layers being adapted with LoRA\n",
    "lora_config = LoraConfig(\n",
    "    r=64,\n",
    "    lora_alpha=16,\n",
    "    target_modules=None,\n",
    ")\n",
    "# create_headed_qlora is the way to go for models with LoRA and newly initialized heads\n",
    "model = create_headed_qlora(\n",
    "    base_model_class=model_class,\n",
    "    model_name=model_path,\n",
    "    quantization_config=quantization_config,\n",
    "    lora_config=lora_config,\n",
    "    head_configs=heads,\n",
    "    fully_trained_heads=True,\n",
    "    device_map={\"\": torch.cuda.current_device()},\n",
    ")\n",
    "# Now you would do some training ...\n",
    "# Saving is still easy using the huggingface api\n",
    "model.save_pretrained(\"test_model_qlora\")\n",
    "del model\n",
    "\n",
    "# Load the qlora model with it's heads. We only need the base model class and the save location. Loading quantized is fully optional here.\n",
    "model = load_lora_with_heads(\n",
    "    model_class,\n",
    "    \"test_model_qlora\",\n",
    "    quantization_config,\n",
    "    device_map={\"\": torch.cuda.current_device()},\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sh_finetuning",
   "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.13"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 13.248571,
   "end_time": "2024-03-24T00:36:56.931019",
   "environment_variables": {},
   "exception": null,
   "input_path": "notebooks/base/saving_and_loading.ipynb",
   "output_path": "notebooks/gpt2/saving_and_loading.ipynb",
   "parameters": {
    "model_path": "gpt2"
   },
   "start_time": "2024-03-24T00:36:43.682448",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}