mamba-tf-lora-merge-peft-test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2073977",
   "metadata": {},
   "source": [
    "### Load Pretrained Model\n",
    "Load a pretrained Mamba Model that is compatible with Transformers Library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db293bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/soma/opt/anaconda3/envs/shane/lib/python3.9/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",
      "/Users/soma/opt/anaconda3/envs/shane/lib/python3.9/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, I'm looking for a new job. I've been working at a company for about a\n"
     ]
    }
   ],
   "source": [
    "from modeling_mamba import MambaForCausalLM\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "model = MambaForCausalLM.from_pretrained('Q-bert/Mamba-130M')\n",
    "tokenizer = AutoTokenizer.from_pretrained('Q-bert/Mamba-130M')\n",
    "\n",
    "text = \"Hi\"\n",
    "\n",
    "input_ids = tokenizer.encode(text, return_tensors=\"pt\")\n",
    "\n",
    "output = model.generate(input_ids, max_length=20, num_beams=5, no_repeat_ngram_size=2)\n",
    "\n",
    "generated_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47713d90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('', <class 'modeling_mamba.MambaForCausalLM'>), ('model', <class 'modeling_mamba.MambaModel'>), ('model.embedding', <class 'torch.nn.modules.sparse.Embedding'>), ('model.layers', <class 'torch.nn.modules.container.ModuleList'>), ('model.layers.0', <class 'modeling_mamba.MambaBlock'>), ('model.layers.0.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.0.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.0.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.0.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.0.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.0.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.1', <class 'modeling_mamba.MambaBlock'>), ('model.layers.1.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.1.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.1.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.1.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.1.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.1.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.2', <class 'modeling_mamba.MambaBlock'>), ('model.layers.2.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.2.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.2.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.2.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.2.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.2.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.3', <class 'modeling_mamba.MambaBlock'>), ('model.layers.3.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.3.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.3.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.3.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.3.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.3.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.4', <class 'modeling_mamba.MambaBlock'>), ('model.layers.4.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.4.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.4.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.4.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.4.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.4.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.5', <class 'modeling_mamba.MambaBlock'>), ('model.layers.5.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.5.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.5.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.5.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.5.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.5.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.6', <class 'modeling_mamba.MambaBlock'>), ('model.layers.6.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.6.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.6.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.6.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.6.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.6.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.7', <class 'modeling_mamba.MambaBlock'>), ('model.layers.7.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.7.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.7.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.7.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.7.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.7.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.8', <class 'modeling_mamba.MambaBlock'>), ('model.layers.8.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.8.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.8.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.8.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.8.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.8.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.9', <class 'modeling_mamba.MambaBlock'>), ('model.layers.9.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.9.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.9.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.9.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.9.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.9.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.10', <class 'modeling_mamba.MambaBlock'>), ('model.layers.10.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.10.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.10.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.10.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.10.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.10.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.11', <class 'modeling_mamba.MambaBlock'>), ('model.layers.11.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.11.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.11.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.11.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.11.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.11.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.12', <class 'modeling_mamba.MambaBlock'>), ('model.layers.12.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.12.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.12.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.12.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.12.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.12.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.13', <class 'modeling_mamba.MambaBlock'>), ('model.layers.13.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.13.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.13.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.13.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.13.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.13.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.14', <class 'modeling_mamba.MambaBlock'>), ('model.layers.14.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.14.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.14.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.14.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.14.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.14.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.15', <class 'modeling_mamba.MambaBlock'>), ('model.layers.15.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.15.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.15.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.15.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.15.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.15.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.16', <class 'modeling_mamba.MambaBlock'>), ('model.layers.16.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.16.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.16.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.16.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.16.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.16.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.17', <class 'modeling_mamba.MambaBlock'>), ('model.layers.17.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.17.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.17.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.17.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.17.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.17.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.18', <class 'modeling_mamba.MambaBlock'>), ('model.layers.18.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.18.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.18.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.18.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.18.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.18.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.19', <class 'modeling_mamba.MambaBlock'>), ('model.layers.19.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.19.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.19.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.19.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.19.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.19.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.20', <class 'modeling_mamba.MambaBlock'>), ('model.layers.20.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.20.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.20.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.20.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.20.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.20.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.21', <class 'modeling_mamba.MambaBlock'>), ('model.layers.21.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.21.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.21.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.21.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.21.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.21.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.22', <class 'modeling_mamba.MambaBlock'>), ('model.layers.22.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.22.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.22.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.22.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.22.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.22.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.layers.23', <class 'modeling_mamba.MambaBlock'>), ('model.layers.23.in_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.23.conv1d', <class 'torch.nn.modules.conv.Conv1d'>), ('model.layers.23.x_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.23.dt_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.23.out_proj', <class 'torch.nn.modules.linear.Linear'>), ('model.layers.23.norm', <class 'modeling_mamba.MambaRMSNorm'>), ('model.norm_f', <class 'modeling_mamba.MambaRMSNorm'>), ('lm_head', <class 'torch.nn.modules.linear.Linear'>)]\n"
     ]
    }
   ],
   "source": [
    "print([(n, type(m)) for n, m in model.named_modules()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d639baf3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'modeling_mamba.MambaForCausalLM'>\n",
      "trainable params: 129135360 || all params: 129135360 || trainable%: 100.0\n",
      "plain None\n"
     ]
    }
   ],
   "source": [
    "print(type(model))\n",
    "\n",
    "def print_trainable_parameters(model):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\n",
    "    \"\"\"\n",
    "    trainable_params = 0\n",
    "    all_param = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_param += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_params += param.numel()\n",
    "    print(\n",
    "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
    "    )\n",
    "    \n",
    "print('plain',print_trainable_parameters(model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ba6b4fe",
   "metadata": {},
   "source": [
    "### Add LoRA adapters\n",
    "1. Identify a particular layer in the Mamba and add an LoRA layer there\n",
    "2. At this time, is only layer to verify if the code works\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf338c0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 12,928 || all params: 129,148,288 || trainable%: 0.010010198509174199\n"
     ]
    }
   ],
   "source": [
    "# adapter-1\n",
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "\n",
    "target_modules=[\"model.layers.3.x_proj\"]\n",
    "\n",
    "model = MambaForCausalLM.from_pretrained('Q-bert/Mamba-130M')\n",
    "\n",
    "config = LoraConfig(\n",
    "target_modules = target_modules,\n",
    "task_type=\"CAUSAL_LM\")\n",
    "m1 = get_peft_model(model, config)\n",
    "m1.print_trainable_parameters()\n",
    "m1.save_pretrained(\"./wts/ada-1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33dc0241",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 12,928 || all params: 129,148,288 || trainable%: 0.010010198509174199\n"
     ]
    }
   ],
   "source": [
    "# adapter-2\n",
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "\n",
    "target_modules=[\"model.layers.2.x_proj\"]\n",
    "model = MambaForCausalLM.from_pretrained('Q-bert/Mamba-130M')\n",
    "\n",
    "config = LoraConfig(\n",
    "target_modules = target_modules,\n",
    "task_type=\"CAUSAL_LM\")\n",
    "m2 = get_peft_model(model, config)\n",
    "m2.print_trainable_parameters()\n",
    "m2.save_pretrained(\"./wts/ada-2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bfeb5cd",
   "metadata": {},
   "source": [
    "### Merge the adpater into the Model\n",
    "merge the adapter back to the model, so the merged model will have exactly the same architecture\n",
    "except with the weights modified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95115979",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://stackoverflow.com/questions/77596271/i-want-to-merge-my-peft-adapter-model-with-the-base-model-and-make-a-fully-new-m\n",
    "\n",
    "#from peft import PeftConfig, PeftModel\n",
    "#base_model = MambaForCausalLM.from_pretrained('Q-bert/Mamba-130M')\n",
    "#adapter_path = './wts/ada-1'\n",
    "#m1 = PeftModel.from_pretrained(base_model, adapter_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "631483dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 129135360 || all params: 129135360 || trainable%: 100.0\n",
      "base mamba None\n",
      "trainable params: 12928 || all params: 129148288 || trainable%: 0.010010198509174199\n",
      "base mamba None\n",
      "trainable params: 25856 || all params: 129161216 || trainable%: 0.02001839313745699\n",
      "base mamba None\n",
      "Hi, I'm looking for a new job. I've been working at a company for about a\n"
     ]
    }
   ],
   "source": [
    "# https://github.com/huggingface/peft\n",
    "from peft import PeftMixedModel\n",
    "\n",
    "\n",
    "base_model = MambaForCausalLM.from_pretrained('Q-bert/Mamba-130M')\n",
    "print('base mamba',print_trainable_parameters(base_model))\n",
    "peft_model = PeftMixedModel.from_pretrained(base_model, adapter_path , \"adapter1\")\n",
    "print('base mamba',print_trainable_parameters(base_model))\n",
    "peft_model.load_adapter('./wts/ada-2', \"adapter2\")\n",
    "peft_model.set_adapter([\"adapter1\", \"adapter2\"])\n",
    "print('base mamba',print_trainable_parameters(base_model))\n",
    "\n",
    "output = peft_model.generate(input_ids, max_length=20, num_beams=5, no_repeat_ngram_size=2)\n",
    "generated_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cc066cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# createa combined fed model\n",
    "peft_model.merge_and_unload()\n",
    "peft_model.base_model.save_pretrained(\"./mbins/fed-lora/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74ae981d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "shane",
   "language": "python",
   "name": "shane"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}