alpaca-convert-colab-fixed.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPNs1ApgbyjyR5X83qlFhUJ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/taiyou2000/alpaca-convert-colab/blob/main/alpaca-convert-colab-fixed.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Arn1_0Y5B-s8"
      },
      "outputs": [],
      "source": [
        "#Install dependencies. Connecting to drive is optional.\n",
        "\n",
        "#from google.colab import drive\n",
        "#drive.mount(\"/content/drive\")\n",
        "\n",
        "!pip install -q datasets loralib sentencepiece\n",
        "!pip uninstall transformers\n",
        "!pip install -q git+https://github.com/zphang/transformers@c3dc391\n",
        "!pip install bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Download Alpaca-native model and convert .bin files into a .pth file\n",
        "import os\n",
        "import json\n",
        "\n",
        "import torch\n",
        "from transformers import LLaMATokenizer, LLaMAForCausalLM\n",
        "\n",
        "tokenizer = LLaMATokenizer.from_pretrained(\"decapoda-research/llama-7b-hf\")\n",
        "\n",
        "base_model = LLaMAForCausalLM.from_pretrained(\n",
        "    \"chavinlo/alpaca-native\",\n",
        "    load_in_8bit=False,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map={\"\": \"cpu\"},\n",
        ")\n",
        "\n",
        "base_model_sd = base_model.state_dict()\n",
        "\n",
        "params = {\n",
        "    \"dim\": 4096,\n",
        "    \"multiple_of\": 256,\n",
        "    \"n_heads\": 32,\n",
        "    \"n_layers\": 32,\n",
        "    \"norm_eps\": 1e-06,\n",
        "    \"vocab_size\": -1,\n",
        "}\n",
        "n_layers = params[\"n_layers\"]\n",
        "n_heads = params[\"n_heads\"]\n",
        "dim = params[\"dim\"]\n",
        "dims_per_head = dim // n_heads\n",
        "base = 10000.0\n",
        "inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))\n",
        "\n",
        "\n",
        "def permute(w):\n",
        "    return (\n",
        "        w.view(n_heads, dim // n_heads // 2, 2, dim).transpose(1, 2).reshape(dim, dim)\n",
        "    )\n",
        "\n",
        "\n",
        "def unpermute(w):\n",
        "    return (\n",
        "        w.view(n_heads, 2, dim // n_heads // 2, dim).transpose(1, 2).reshape(dim, dim)\n",
        "    )\n",
        "\n",
        "\n",
        "def translate_state_dict_key(k):\n",
        "    k = k.replace(\"base_model.model.\", \"\")\n",
        "    if k == \"model.embed_tokens.weight\":\n",
        "        return \"tok_embeddings.weight\"\n",
        "    elif k == \"model.norm.weight\":\n",
        "        return \"norm.weight\"\n",
        "    elif k == \"lm_head.weight\":\n",
        "        return \"output.weight\"\n",
        "    elif k.startswith(\"model.layers.\"):\n",
        "        layer = k.split(\".\")[2]\n",
        "        if k.endswith(\".self_attn.q_proj.weight\"):\n",
        "            return f\"layers.{layer}.attention.wq.weight\"\n",
        "        elif k.endswith(\".self_attn.k_proj.weight\"):\n",
        "            return f\"layers.{layer}.attention.wk.weight\"\n",
        "        elif k.endswith(\".self_attn.v_proj.weight\"):\n",
        "            return f\"layers.{layer}.attention.wv.weight\"\n",
        "        elif k.endswith(\".self_attn.o_proj.weight\"):\n",
        "            return f\"layers.{layer}.attention.wo.weight\"\n",
        "        elif k.endswith(\".mlp.gate_proj.weight\"):\n",
        "            return f\"layers.{layer}.feed_forward.w1.weight\"\n",
        "        elif k.endswith(\".mlp.down_proj.weight\"):\n",
        "            return f\"layers.{layer}.feed_forward.w2.weight\"\n",
        "        elif k.endswith(\".mlp.up_proj.weight\"):\n",
        "            return f\"layers.{layer}.feed_forward.w3.weight\"\n",
        "        elif k.endswith(\".input_layernorm.weight\"):\n",
        "            return f\"layers.{layer}.attention_norm.weight\"\n",
        "        elif k.endswith(\".post_attention_layernorm.weight\"):\n",
        "            return f\"layers.{layer}.ffn_norm.weight\"\n",
        "        elif k.endswith(\"rotary_emb.inv_freq\") or \"lora\" in k:\n",
        "            return None\n",
        "        else:\n",
        "            print(layer, k)\n",
        "            raise NotImplementedError\n",
        "    else:\n",
        "        print(k)\n",
        "        raise NotImplementedError\n",
        "\n",
        "\n",
        "new_state_dict = {}\n",
        "for k, v in base_model_sd.items():\n",
        "    new_k = translate_state_dict_key(k)\n",
        "    if new_k is not None:\n",
        "        if \"wq\" in new_k or \"wk\" in new_k:\n",
        "            new_state_dict[new_k] = unpermute(v)\n",
        "        else:\n",
        "            new_state_dict[new_k] = v\n",
        "\n",
        "torch.save(new_state_dict, \"consolidated.00.pth\")\n",
        "\n",
        "with open(\"params.json\", \"w\") as f:\n",
        "    json.dump(params, f)"
      ],
      "metadata": {
        "id": "h_HANiwTCGX8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#Resize tensors\n",
        "model = torch.load(\"consolidated.00.pth\", map_location=torch.device('cpu'))\n",
        "x = model[\"tok_embeddings.weight\"]\n",
        "y = model[\"output.weight\"]\n",
        "row_exclude = 32000\n",
        "x = x[:row_exclude]\n",
        "y = y[:row_exclude]\n",
        "model[\"tok_embeddings.weight\"] = x\n",
        "model[\"output.weight\"] = y\n",
        "torch.save(model, \"consolidated.01.pth\")\n",
        "#Delete consolidated.00.pth and rename consolidated.01.pth into consolidated.00.pth"
      ],
      "metadata": {
        "id": "bYgfDtxRCMhq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#Clone alpaca.cpp repository to quantize the model.\n",
        "!git clone https://github.com/antimatter15/alpaca.cpp"
      ],
      "metadata": {
        "id": "nmfMjlhQCPnf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cd llama.cpp"
      ],
      "metadata": {
        "id": "bDedHHfxCRhp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir models"
      ],
      "metadata": {
        "id": "TBpcI0K4CTzb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cd models"
      ],
      "metadata": {
        "id": "fgNUKgGVCWCu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir 7B\n",
        "#Move consolidated.00.pth(.pth file you converted into above) and params.json to llama.cpp/models/7B folder"
      ],
      "metadata": {
        "id": "PF6CwThcCY3X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cd ..\n",
        "cd .."
      ],
      "metadata": {
        "id": "sn7z21DpCbj7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!make\n",
        "!python3 convert-pth-to-ggml.py models/7B/ 1\n",
        "!./quantize models/7B/ggml-model-f16.bin ggml-alpaca-7b-native-q4.bin 2"
      ],
      "metadata": {
        "id": "u9LyDh0tCgaV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "./chat -8 -n 512"
      ],
      "metadata": {
        "id": "mgk4Ia54CjO_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}