merge_lora_weights_into_cerebras.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c49855ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "CUDA SETUP: CUDA runtime path found: /home/lxe/miniconda3/envs/llama-lora/lib/libcudart.so\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 113\n",
      "CUDA SETUP: Loading binary /home/lxe/miniconda3/envs/llama-lora/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda113.so...\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import transformers\n",
    "from peft import PeftModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "342fe4bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_MODEL = 'cerebras/Cerebras-GPT-2.7B'\n",
    "\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(BASE_MODEL)\n",
    "\n",
    "base_model = transformers.AutoModelForCausalLM.from_pretrained(\n",
    "    BASE_MODEL,\n",
    "    load_in_8bit=False,\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map={\"\": \"cpu\"},\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "43cc6987",
   "metadata": {},
   "outputs": [],
   "source": [
    "# base_model.sd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b0e9d29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_weight = base_model.transformer.h[0].attn.c_attn.weight\n",
    "first_weight_old = first_weight.clone()\n",
    "\n",
    "lora_model = PeftModel.from_pretrained(\n",
    "    base_model,\n",
    "    \"lora-cerebras-gpt2.7b-alpaca\",\n",
    "    device_map={\"\": \"cpu\"},\n",
    "    torch_dtype=torch.float16,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "497da48e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_weight = lora_model.base_model.transformer.h[0].attn.c_attn.weight\n",
    "\n",
    "assert torch.allclose(first_weight_old, first_weight)\n",
    "\n",
    "# merge weights\n",
    "for layer in lora_model.base_model.transformer.h:\n",
    "    layer.attn.c_attn.merge_weights = True\n",
    "\n",
    "lora_model.train(False)\n",
    "\n",
    "# did we do anything?\n",
    "assert not torch.allclose(first_weight_old, first_weight)\n",
    "\n",
    "lora_model_sd = lora_model.state_dict()\n",
    "deloreanized_sd = {\n",
    "    k.replace(\"base_model.model.\", \"\"): v\n",
    "    for k, v in lora_model_sd.items()\n",
    "    if \"lora\" not in k\n",
    "}\n",
    "\n",
    "base_model.save_pretrained(\"./hf_ckpt\", state_dict=deloreanized_sd, max_shard_size=\"12GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4563bc5f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}