runing_quantized_mixtral_7bx8.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EWOtpjJMDmYE"
      },
      "source": [
        "# Runing Falcon-180B on a single A100 80GB"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This file including three section:\n",
        "- [(Optional) Train the quantization parameters of Falcon-180B by yourself.](#optional-train-the-quantization-parameters-of-falcon-180b-by-yourself)\n",
        "- [Download the pre-quantized models](#download-the-pre-quantized-models)\n",
        "- [Let's Infer!](#lets-infer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## (Optional) Train the quantization parameters of Mixtral-7bx8 by yourself.\n",
        "\n",
        "This section provids how to train the quantization parameters of Mixtral-7bx8 by yourself. You can skip this section because we have provided the pre-built quantized models in [Download the pre-quantized models](#download-the-pre-quantized-models)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!CUDA_VISIBLE_DEVICES=0 python main.py \\\n",
        "--model /PATH/TO/Mixtral-8x7B-v0.1 \\\n",
        "--epochs 10 --output_dir ./log/mixtral-8x7b_w4a16g128 \\\n",
        "--wbits 4 --abits 16 --group_size 128 --lwc  \\\n",
        "--nsamples 128 --net mixtral-8x7b \\\n",
        "--real_quant --eval_ppl \\\n",
        "--save_dir ./checkpoint/mixtral-8x7b_w4a16g128"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Download the prebuilt quantized model:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FwsWd1WbDmYE"
      },
      "source": [
        "We have provide the prebuilt quantized model on Huggingface. In order to download the large weights, we'll have to use `git lfs`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V0GjINnMDmYF"
      },
      "outputs": [],
      "source": [
        "!conda install git git-lfs\n",
        "!git lfs install"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FSAe7Ew_DmYF"
      },
      "outputs": [],
      "source": [
        "!mkdir -p pre_quantized_models/\n",
        "\n",
        "# download mixtral-8x7b-v0.1 with w4a16g128 quantization\n",
        "!git clone https://huggingface.co/ChenMnZ/Mixtral-8x7B-v0.1-OmniQuantv1-w4a16g128 ./pre_quantized_models/Mixtral-8x7B-v0.1-OmniQuantv1-w4a16g128\n",
        "\n",
        "# download mixtral-8x7b-v0.1-Instruct with w4a16g128 quantization\n",
        "# !git clone https://huggingface.co/ChenMnZ/Mixtral-8x7B-Instruct-v0.1-OmniQuantv1-w4a16g128 ./pre_quantized_models/Mixtral-8x7B-Instruct-v0.1-OmniQuantv1-w4a16g128"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "76Ru5__tDmYF"
      },
      "source": [
        "## Let's Infer!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Constraint in one GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_in_model\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig\n",
        "import torch\n",
        "import auto_gptq.nn_modules.qlinear.qlinear_cuda as qlinear_cuda\n",
        "import auto_gptq.nn_modules.qlinear.qlinear_triton as qlinear_triton\n",
        "from tqdm import tqdm\n",
        "import gc   \n",
        "import time\n",
        "\n",
        "def get_named_linears(module):\n",
        "    return {name: m for name, m in module.named_modules() if isinstance(m, torch.nn.Linear) and not 'gate' in name}\n",
        "\n",
        "def set_op_by_name(layer, name, new_module):\n",
        "    levels = name.split('.')\n",
        "    if len(levels) > 1:\n",
        "        mod_ = layer\n",
        "        for l_idx in range(len(levels)-1):\n",
        "            if levels[l_idx].isdigit():\n",
        "                mod_ = mod_[int(levels[l_idx])]\n",
        "            else:\n",
        "                mod_ = getattr(mod_, levels[l_idx])\n",
        "        setattr(mod_, levels[-1], new_module)\n",
        "    else:\n",
        "        setattr(layer, name, new_module)\n",
        "\n",
        "# manually adjust the model_path, corresponding weight bit width, group size.\n",
        "model_path = './pre_quantized_models/Mixtral-8x7B-v0.1-OmniQuantv1-w4a16g128'\n",
        "wbits = 4\n",
        "group_size = 128\n",
        "\n",
        "\n",
        "\n",
        "config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)\n",
        "enc = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)\n",
        "with init_empty_weights():\n",
        "    model = AutoModelForCausalLM.from_config(config=config,torch_dtype=torch.float16, trust_remote_code=True)\n",
        "\n",
        "layers = model.model.layers\n",
        "for i in tqdm(range(len(layers))):\n",
        "    layer = layers[i]\n",
        "    named_linears = get_named_linears(layer)\n",
        "    for name, module in named_linears.items():\n",
        "        if wbits in [2,4]:\n",
        "            q_linear = qlinear_triton.QuantLinear(wbits, group_size, module.in_features,module.out_features,not module.bias is None,kernel_switch_threshold=128)\n",
        "        elif wbits == 3:\n",
        "            q_linear = qlinear_cuda.QuantLinear(wbits, group_size, module.in_features,module.out_features,not module.bias is None,kernel_switch_threshold=128)\n",
        "        else:\n",
        "            raise NotImplementedError(\"Only 2,3,4 bits are supported.\")\n",
        "        q_linear.to(next(layer.parameters()).device)\n",
        "        set_op_by_name(layer, name, q_linear)\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "model.tie_weights()\n",
        "device_map = infer_auto_device_map(model)\n",
        "print(\"Loading pre-computed quantized weights...\")\n",
        "load_checkpoint_in_model(model,checkpoint=model_path,device_map=device_map,offload_state_dict=True)\n",
        "print(\"Loading pre-computed quantized weights Successfully\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Start inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model.eval()\n",
        "prompt = \"Hello my name is\"\n",
        "input_ids = enc(prompt, return_tensors='pt').input_ids.cuda()\n",
        "model = model.cuda()\n",
        "start_time = time.time()\n",
        "output = model.generate(inputs=input_ids, max_new_tokens=128)\n",
        "end_time = time.time()\n",
        "speed = len(output[0])/(end_time-start_time)\n",
        "print(enc.decode(output[0], skip_special_tokens=True))\n",
        "print(f\"speed:{speed:.2f}token/s max memory: {torch.cuda.max_memory_allocated(model.device)/ 1024**2:.2f}M\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.16"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}