MoE_Interprebility.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyN7akHQ0r8oJRR8HYiZfgGC",
      "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/mnida/mech-interp-r1/blob/main/MoE_Interprebility.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": "Imuyk0gsW7Ta"
      },
      "outputs": [],
      "source": [
        "# fix numpy in colab\n",
        "import numpy\n",
        "\n",
        "# fix triton in colab\n",
        "!export LC_ALL=\"en_US.UTF-8\"\n",
        "!export LD_LIBRARY_PATH=\"/usr/lib64-nvidia\"\n",
        "!export LIBRARY_PATH=\"/usr/local/cuda/lib64/stubs\"\n",
        "!ldconfig /usr/lib64-nvidia\n",
        "\n",
        "!git clone https://github.com/dvmazur/mixtral-offloading.git --quiet\n",
        "!cd mixtral-offloading && pip install -q -r requirements.txt\n",
        "!huggingface-cli download lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo --quiet --local-dir Mixtral-8x7B-Instruct-v0.1-offloading-demo\n",
        "\n",
        "from IPython.display import clear_output\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "v7HuQzY5XMuT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = \"mistralai/Mixtral-8x7B-Instruct-v0.1\"\n",
        "quantized_model_name = \"lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo\"\n",
        "state_path = \"Mixtral-8x7B-Instruct-v0.1-offloading-demo\"\n",
        "\n",
        "config = AutoConfig.from_pretrained(quantized_model_name)\n",
        "device = torch.device(\"cuda:0\")\n",
        "\n",
        "##### Change this to 5 if you have only 12 GB of GPU VRAM #####\n",
        "offload_per_layer = 4\n",
        "# offload_per_layer = 5\n",
        "###############################################################\n",
        "\n",
        "num_experts = config.num_local_experts\n",
        "\n",
        "offload_config = OffloadConfig(\n",
        "    main_size=config.num_hidden_layers * (num_experts - offload_per_layer),\n",
        "    offload_size=config.num_hidden_layers * offload_per_layer,\n",
        "    buffer_size=4,\n",
        "    offload_per_layer=offload_per_layer,\n",
        ")\n",
        "\n",
        "\n",
        "attn_config = BaseQuantizeConfig(\n",
        "    nbits=4,\n",
        "    group_size=64,\n",
        "    quant_zero=True,\n",
        "    quant_scale=True,\n",
        ")\n",
        "attn_config[\"scale_quant_params\"][\"group_size\"] = 256\n",
        "\n",
        "\n",
        "ffn_config = BaseQuantizeConfig(\n",
        "    nbits=2,\n",
        "    group_size=16,\n",
        "    quant_zero=True,\n",
        "    quant_scale=True,\n",
        ")\n",
        "quant_config = QuantConfig(ffn_config=ffn_config, attn_config=attn_config)\n",
        "\n",
        "\n",
        "model = build_model(\n",
        "    device=device,\n",
        "    quant_config=quant_config,\n",
        "    offload_config=offload_config,\n",
        "    state_path=state_path,              # Mixtral Weights\n",
        ")"
      ],
      "metadata": {
        "id": "1WJD11YZXFMw"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}