Mixtral in colab

HiroRolli · Feb 10, 2024 · aec8950 · aec8950
1 parent a693a35
commit aec8950
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/notebooks/colab_example.ipynb b/notebooks/colab_example.ipynb
@@ -20,7 +20,8 @@
       },
       "source": [
         "**Install the `aqlm` library**\n",
-        "- the only extra dependency to run AQLM models."
+        "- `aqlm` is the only extra dependency to run AQLM models.\n",
+        "- Install latest `accelerate` to pull the latest bugfixes."
       ]
     },
     {
@@ -32,7 +33,8 @@
       "outputs": [],
       "source": [
         "%%capture\n",
-        "!pip install aqlm[gpu]==1.0.0"
+        "!pip install aqlm[gpu]==1.0.0\n",
+        "!pip install git+https://github.com/huggingface/accelerate.git@main"
       ]
     },
     {
@@ -44,10 +46,11 @@
         "**Load the model as usual**\n",
         "\n",
         "Just don't forget to add:\n",
-        " - `trust_remote_code=True` to pull the inference code\n",
+        " - `trust_remote_code=True` to pull the inference code.\n",
         " - `torch_dtype=\"auto\"` to load the model in it's native dtype.\n",
+        " - `device_map=\"cuda\"` to load the model on GPU straight away, saving RAM.\n",
         "\n",
-        "The tokenizer is just a normal `Llama 2` tokenizer."
+        "The tokenizer is just a normal `Mixtral` tokenizer."
       ]
     },
     {
@@ -167,10 +170,10 @@
         "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
         "\n",
         "quantized_model = AutoModelForCausalLM.from_pretrained(\n",
-        "    \"BlackSamorez/Llama-2-7b-AQLM-2Bit-1x16-hf\",\n",
-        "    trust_remote_code=True, torch_dtype=\"auto\"\n",
+        "    \"BlackSamorez/Mixtral-8x7b-AQLM-2Bit-1x16-hf\",\n",
+        "    trust_remote_code=True, torch_dtype=\"auto\", device_map=\"cuda\"\n",
         ").cuda()\n",
-        "tokenizer = AutoTokenizer.from_pretrained(\"BlackSamorez/Llama-2-7b-AQLM-2Bit-1x16-hf\")"
+        "tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mixtral-8x7B-v0.1\")"
       ]
     },
     {
@@ -243,7 +246,7 @@
         "id": "nvShqlguccep"
       },
       "source": [
-        "**Check that the output is what one would expect from Llama-2-7b**"
+        "**Check that the output is what one would expect from Mixtral**"
       ]
     },
     {