diff --git a/model-inference.py b/model-inference.py new file mode 100644 index 0000000..cb67702 --- /dev/null +++ b/model-inference.py @@ -0,0 +1,79 @@ +import numpy as np +import sys +import torch +from torch.nn import functional as F +from transformers import AutoConfig, AutoTokenizer +from huggingface_hub import snapshot_download +from tqdm.auto import trange +from src.build_model import build_model, OffloadConfig +import os + +def initialize_model(): + model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1" + config = AutoConfig.from_pretrained(model_name) + state_path = snapshot_download("lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo") + device = torch.device("cuda:0") + + offload_config = OffloadConfig( + main_size=config.num_local_experts * config.num_hidden_layers * 4 // 8, + offload_size=config.num_local_experts * config.num_hidden_layers * 4 // 8, + buffer_size=4, + offload_per_layer=4, + ) + + model = build_model(device=device, offload_config=offload_config, state_path=state_path) + return model, model_name, device + +def get_user_input(): + return input("Enter your query (or type 'exit' to quit): ") + +def generate_text(model, model_name, device, user_query): + tokenizer = AutoTokenizer.from_pretrained(model_name) + id_to_token = {v: k for k, v in tokenizer.get_vocab().items()} + + input_ids = tokenizer.apply_chat_template( + [dict(role="user", content=user_query)], + return_tensors='pt', + ).to(device) + + inputs = dict(input_ids=input_ids, attention_mask=torch.ones_like(input_ids)) + + generated_tokens = [] + past_key_values = None + + with torch.no_grad(): + for _ in trange(100): + out = model(**inputs, past_key_values=past_key_values, output_hidden_states=True) + past_key_values = out.past_key_values + + probs = F.softmax(out.logits[0, -1] / 0.9, dim=-1) + token_id = torch.multinomial(probs, 1) + token = id_to_token[token_id.item()] + + if token[0] == '▁': + generated_tokens.append(' ') + token = token[1:] + generated_tokens.append(token) + + inp = token_id.reshape(1, 1) + inputs = dict(input_ids=inp) + + return ''.join(generated_tokens) + +# Ensure environment variables are set +os.environ["LC_ALL"] = "en_US.UTF-8" +os.environ["LD_LIBRARY_PATH"] = "/usr/lib64-nvidia" +os.environ["LIBRARY_PATH"] = "/usr/local/cuda/lib64/stubs" + +# Ensure the repository is cloned and added to sys.path +# Clone it manually or handle it in the script as needed +# sys.path.append("mixtral-offloading") + +model, model_name, device = initialize_model() + +while True: + user_query = get_user_input() + if user_query.lower() == 'exit': + break + response = generate_text(model, model_name, device, user_query) + print("Response:", response) diff --git a/notebooks/8_7bMixtral.ipynb b/notebooks/8_7bMixtral.ipynb new file mode 100644 index 0000000..d1b8ceb --- /dev/null +++ b/notebooks/8_7bMixtral.ipynb @@ -0,0 +1,524 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU", + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "3c1784f9ce07432197f3604e599f62fc": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HBoxModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HBoxModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HBoxView", + "box_style": "", + "children": [ + "IPY_MODEL_41ba401244854dd0bbf7879a6014e9f0", + "IPY_MODEL_7dce077c46644569bf1025ee395384d3", + "IPY_MODEL_97ea82d833e444ef94b8f13ca828de00" + ], + "layout": "IPY_MODEL_8be8cbe72de441ee80f980dc23858bcb" + } + }, + "41ba401244854dd0bbf7879a6014e9f0": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HTMLModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_1a066991e85e4b5d9ca517069877f3f2", + "placeholder": "​", + "style": "IPY_MODEL_651b506ff2564520b2dad773fe38faae", + "value": "Fetching 265 files: 100%" + } + }, + "7dce077c46644569bf1025ee395384d3": { + "model_module": "@jupyter-widgets/controls", + "model_name": "FloatProgressModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "FloatProgressModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "ProgressView", + "bar_style": "success", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_dffebc421f1f4d21a69d3c33ce251af4", + "max": 265, + "min": 0, + "orientation": "horizontal", + "style": "IPY_MODEL_674806c5d0c34855a7f89b2d3c5a14a9", + "value": 265 + } + }, + "97ea82d833e444ef94b8f13ca828de00": { + "model_module": "@jupyter-widgets/controls", + "model_name": "HTMLModel", + "model_module_version": "1.5.0", + "state": { + "_dom_classes": [], + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "HTMLModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/controls", + "_view_module_version": "1.5.0", + "_view_name": "HTMLView", + "description": "", + "description_tooltip": null, + "layout": "IPY_MODEL_a1f11b3a9af548148c845756a2a7bb9a", + "placeholder": "​", + "style": "IPY_MODEL_bcf1dbb928c14ff180fea41cc20aa03c", + "value": " 265/265 [00:00<00:00, 4153.54it/s]" + } + }, + "8be8cbe72de441ee80f980dc23858bcb": { + "model_module": "@jupyter-widgets/base", + "model_name": "LayoutModel", + "model_module_version": "1.2.0", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "1a066991e85e4b5d9ca517069877f3f2": { + "model_module": "@jupyter-widgets/base", + "model_name": "LayoutModel", + "model_module_version": "1.2.0", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "651b506ff2564520b2dad773fe38faae": { + "model_module": "@jupyter-widgets/controls", + "model_name": "DescriptionStyleModel", + "model_module_version": "1.5.0", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + }, + "dffebc421f1f4d21a69d3c33ce251af4": { + "model_module": "@jupyter-widgets/base", + "model_name": "LayoutModel", + "model_module_version": "1.2.0", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "674806c5d0c34855a7f89b2d3c5a14a9": { + "model_module": "@jupyter-widgets/controls", + "model_name": "ProgressStyleModel", + "model_module_version": "1.5.0", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "ProgressStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "bar_color": null, + "description_width": "" + } + }, + "a1f11b3a9af548148c845756a2a7bb9a": { + "model_module": "@jupyter-widgets/base", + "model_name": "LayoutModel", + "model_module_version": "1.2.0", + "state": { + "_model_module": "@jupyter-widgets/base", + "_model_module_version": "1.2.0", + "_model_name": "LayoutModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "LayoutView", + "align_content": null, + "align_items": null, + "align_self": null, + "border": null, + "bottom": null, + "display": null, + "flex": null, + "flex_flow": null, + "grid_area": null, + "grid_auto_columns": null, + "grid_auto_flow": null, + "grid_auto_rows": null, + "grid_column": null, + "grid_gap": null, + "grid_row": null, + "grid_template_areas": null, + "grid_template_columns": null, + "grid_template_rows": null, + "height": null, + "justify_content": null, + "justify_items": null, + "left": null, + "margin": null, + "max_height": null, + "max_width": null, + "min_height": null, + "min_width": null, + "object_fit": null, + "object_position": null, + "order": null, + "overflow": null, + "overflow_x": null, + "overflow_y": null, + "padding": null, + "right": null, + "top": null, + "visibility": null, + "width": null + } + }, + "bcf1dbb928c14ff180fea41cc20aa03c": { + "model_module": "@jupyter-widgets/controls", + "model_name": "DescriptionStyleModel", + "model_module_version": "1.5.0", + "state": { + "_model_module": "@jupyter-widgets/controls", + "_model_module_version": "1.5.0", + "_model_name": "DescriptionStyleModel", + "_view_count": null, + "_view_module": "@jupyter-widgets/base", + "_view_module_version": "1.2.0", + "_view_name": "StyleView", + "description_width": "" + } + } + } + } + }, + "cells": [ + { + "cell_type": "code", + "source": [ + "import numpy\n", + "from IPython.display import clear_output\n", + "import sys\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", + "%pip install torch==2.1.0 transformers==4.36.1 hivemind==1.1.9 --quiet\n", + "%pip install git+https://github.com/mobiusml/hqq.git@37502bea31f2969c6680c0c4a88ca74b3bb234a5 --quiet\n", + "%pip install numpy==1.23.5 --quiet\n", + "!git clone https://github.com/dvmazur/mixtral-offloading.git\n", + "\n", + "sys.path.append(\"mixtral-offloading\")\n", + "clear_output()\n" + ], + "metadata": { + "id": "kuaEzzRM1oyD" + }, + "execution_count": 1, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "import numpy as np\n", + "import sys\n", + "import torch\n", + "from torch.nn import functional as F\n", + "from transformers import AutoConfig, AutoTokenizer\n", + "from huggingface_hub import snapshot_download\n", + "from tqdm.auto import trange\n", + "from src.build_model import build_model, OffloadConfig\n", + "import os\n", + "\n", + "def initialize_model():\n", + " model_name = \"mistralai/Mixtral-8x7B-Instruct-v0.1\"\n", + " config = AutoConfig.from_pretrained(model_name)\n", + " state_path = snapshot_download(\"lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo\")\n", + " device = torch.device(\"cuda:0\")\n", + "\n", + " offload_config = OffloadConfig(\n", + " main_size=config.num_local_experts * config.num_hidden_layers * 4 // 8,\n", + " offload_size=config.num_local_experts * config.num_hidden_layers * 4 // 8,\n", + " buffer_size=4,\n", + " offload_per_layer=4,\n", + " )\n", + "\n", + " model = build_model(device=device, offload_config=offload_config, state_path=state_path)\n", + " return model, model_name, device\n", + "\n", + "def get_user_input():\n", + " return input(\"Enter your query (or type 'exit' to quit): \")\n", + "\n", + "def generate_text(model, model_name, device, user_query):\n", + " tokenizer = AutoTokenizer.from_pretrained(model_name)\n", + " id_to_token = {v: k for k, v in tokenizer.get_vocab().items()}\n", + "\n", + " input_ids = tokenizer.apply_chat_template(\n", + " [dict(role=\"user\", content=user_query)],\n", + " return_tensors='pt',\n", + " ).to(device)\n", + "\n", + " inputs = dict(input_ids=input_ids, attention_mask=torch.ones_like(input_ids))\n", + "\n", + " generated_tokens = []\n", + " past_key_values = None\n", + "\n", + " with torch.no_grad():\n", + " for _ in trange(100):\n", + " out = model(**inputs, past_key_values=past_key_values, output_hidden_states=True)\n", + " past_key_values = out.past_key_values\n", + "\n", + " probs = F.softmax(out.logits[0, -1] / 0.9, dim=-1)\n", + " token_id = torch.multinomial(probs, 1)\n", + " token = id_to_token[token_id.item()]\n", + "\n", + " if token[0] == '▁':\n", + " generated_tokens.append(' ')\n", + " token = token[1:]\n", + " generated_tokens.append(token)\n", + "\n", + " inp = token_id.reshape(1, 1)\n", + " inputs = dict(input_ids=inp)\n", + "\n", + " return ''.join(generated_tokens)\n", + "\n", + "model, model_name, device = initialize_model()\n", + "\n", + "while True:\n", + " user_query = get_user_input()\n", + " if user_query.lower() == 'exit':\n", + " break\n", + " response = generate_text(model, model_name, device, user_query)\n", + " print(\"Response:\", response)\n", + "\n" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 468, + "referenced_widgets": [ + "3c1784f9ce07432197f3604e599f62fc", + "41ba401244854dd0bbf7879a6014e9f0", + "7dce077c46644569bf1025ee395384d3", + "97ea82d833e444ef94b8f13ca828de00", + "8be8cbe72de441ee80f980dc23858bcb", + "1a066991e85e4b5d9ca517069877f3f2", + "651b506ff2564520b2dad773fe38faae", + "dffebc421f1f4d21a69d3c33ce251af4", + "674806c5d0c34855a7f89b2d3c5a14a9", + "a1f11b3a9af548148c845756a2a7bb9a", + "bcf1dbb928c14ff180fea41cc20aa03c" + ] + }, + "id": "pZyrk_DO1wN1", + "outputId": "d1b47736-35d6-4e01-dda2-75df8a3fe79c" + }, + "execution_count": 6, + "outputs": [ + { + "output_type": "display_data", + "data": { + "text/plain": [ + "Fetching 265 files: 0%| | 0/265 [00:00\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 61\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0;34m''\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgenerated_tokens\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 62\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 63\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minitialize_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 64\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 65\u001b[0m \u001b[0;32mwhile\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;32m\u001b[0m in \u001b[0;36minitialize_model\u001b[0;34m()\u001b[0m\n\u001b[1;32m 22\u001b[0m )\n\u001b[1;32m 23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbuild_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moffload_config\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moffload_config\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstate_path\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstate_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 25\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;32m/content/mixtral-offloading/src/build_model.py\u001b[0m in \u001b[0;36mbuild_model\u001b[0;34m(device, offload_config, state_path)\u001b[0m\n\u001b[1;32m 178\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 179\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwith_default_dtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat16\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 180\u001b[0;31m model = MixtralForCausalLM(\n\u001b[0m\u001b[1;32m 181\u001b[0m AutoConfig.from_pretrained(\n\u001b[1;32m 182\u001b[0m \u001b[0mmodel_name\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/transformers/models/mixtral/modeling_mixtral.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, config)\u001b[0m\n\u001b[1;32m 1130\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1131\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1132\u001b[0;31m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mMixtralModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1133\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab_size\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1134\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlm_head\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhidden_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/transformers/models/mixtral/modeling_mixtral.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, config)\u001b[0m\n\u001b[1;32m 953\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab_size\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 954\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 955\u001b[0;31m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0membed_tokens\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mEmbedding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhidden_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpadding_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 956\u001b[0m self.layers = nn.ModuleList(\n\u001b[1;32m 957\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mMixtralDecoderLayer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlayer_idx\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mlayer_idx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_hidden_layers\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/sparse.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, num_embeddings, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, _weight, _freeze, device, dtype)\u001b[0m\n\u001b[1;32m 140\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale_grad_by_freq\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mscale_grad_by_freq\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 141\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m_weight\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 142\u001b[0;31m self.weight = Parameter(torch.empty((num_embeddings, embedding_dim), **factory_kwargs),\n\u001b[0m\u001b[1;32m 143\u001b[0m requires_grad=not _freeze)\n\u001b[1;32m 144\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreset_parameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/utils/_device.py\u001b[0m in \u001b[0;36m__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m 75\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mfunc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0m_device_constructors\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'device'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 76\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'device'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 77\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 78\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 79\u001b[0m \u001b[0;31m# NB: This is directly called from C++ in torch/csrc/Device.cpp\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 250.00 MiB. GPU 0 has a total capacty of 14.75 GiB of which 17.06 MiB is free. Process 2899 has 14.71 GiB memory in use. Of the allocated memory 14.17 GiB is allocated by PyTorch, and 428.63 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF" + ] + } + ] + } + ] +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index d6977f9..6f2d1e3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,7 @@ +hivemind==1.1.9 torch>=2.1.0 transformers==4.36.1 git+https://github.com/mobiusml/hqq.git@37502bea31f2969c6680c0c4a88ca74b3bb234a5 numpy==1.24.4 tqdm==4.66.1 + diff --git a/src/custom_layers.py b/src/custom_layers.py index a20b6ee..6ed5399 100644 --- a/src/custom_layers.py +++ b/src/custom_layers.py @@ -1,16 +1,20 @@ import copy -import functools -from transformers.models.mixtral.configuration_mixtral import MixtralConfig -from transformers.activations import ACT2FN -from typing import Dict, Any -from hqq.core.quantize import HQQLinear, Quantizer - import torch +import functools from torch import nn from torch.nn import functional as F +from transformers.models.mixtral.configuration_mixtral import MixtralConfig +from transformers.activations import ACT2FN -from .packing import pack_4bit_u8_common, pack_2bit_u8_common, unpack_4bit_u8_common, unpack_2bit_u8_common -from .triton_kernels import triton_matmul4_transpose, triton_matmul3_transpose, triton_matmul2_transpose +from hqq.core.quantize import HQQLinear, Quantizer +from .packing import ( + pack_4bit_u8_common, pack_2bit_u8_common, + unpack_4bit_u8_common, unpack_2bit_u8_common +) +from .triton_kernels import ( + triton_matmul4_transpose, triton_matmul3_transpose, + triton_matmul2_transpose +) class HQQLinearTritonSavable(HQQLinear):