Add mistral 7b support (#443)

* add mistral model name and alias

* add code for converting mistral config to hooked transformer config

* add function for converting mistral weights

* add GroupedQueryAttention

* add abstract base class for attention

* adapt keyvaluecache if grouped query attention is used

* fix fold_value_biases when using grouped query attention

* Add unit test for grouped query attention

* Add demo notebook for Mistral

* fix formatting

* add documentation for grouped query attention

* update lock file

* use Union instead of | for union types

* hardcode mistral config so building docs works with older versions of transformers

* don't set final_rms in Mistral config

* make Mistral-7b's alias name consistent with other models

* update main demo notebook

* require transformers>=3.34

* improve docstrings and clarify test name for grouped query attention

* remove Mistral demo

* fix docstring format

Author: Felhof

demos/Main_Demo.ipynb

@@ -45,7 +45,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 292,
+            "execution_count": 11,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -80,7 +80,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 293,
+            "execution_count": 12,
             "metadata": {},
             "outputs": [
                 {
@@ -103,32 +103,28 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 294,
+            "execution_count": 13,
             "metadata": {},
             "outputs": [
                 {
                     "data": {
                         "text/html": [
-                            "<div id=\"circuits-vis-7e4c8a75-1335\" style=\"margin: 15px 0;\"/>\n",
+                            "<div id=\"circuits-vis-1f2a8687-9cd7\" style=\"margin: 15px 0;\"/>\n",
                             "    <script crossorigin type=\"module\">\n",
-                            "    import { render, Hello } from \"https://unpkg.com/circuitsvis@1.43.0/dist/cdn/esm.js\";\n",
+                            "    import { render, Hello } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
                             "    render(\n",
-                            "      \"circuits-vis-7e4c8a75-1335\",\n",
+                            "      \"circuits-vis-1f2a8687-9cd7\",\n",
                             "      Hello,\n",
                             "      {\"name\": \"Neel\"}\n",
                             "    )\n",
                             "    </script>"
                         ],
                         "text/plain": [
-                            "<circuitsvis.utils.render.RenderedHTML at 0xffff10cc9f10>"
+                            "<circuitsvis.utils.render.RenderedHTML at 0x7f21437f1c30>"
                         ]
                     },
-                    "execution_count": 294,
-                    "metadata": {
-                        "text/html": {
-                            "Content-Type": "text/html"
-                        }
-                    },
+                    "execution_count": 13,
+                    "metadata": {},
                     "output_type": "execute_result"
                 }
             ],
@@ -140,7 +136,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 295,
+            "execution_count": 14,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -158,7 +154,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 296,
+            "execution_count": 15,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -179,16 +175,16 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 297,
+            "execution_count": 16,
             "metadata": {},
             "outputs": [
                 {
                     "data": {
                         "text/plain": [
-                            "<torch.autograd.grad_mode.set_grad_enabled at 0xffff425948e0>"
+                            "<torch.autograd.grad_mode.set_grad_enabled at 0x7f213de735e0>"
                         ]
                     },
-                    "execution_count": 297,
+                    "execution_count": 16,
                     "metadata": {},
                     "output_type": "execute_result"
                 }
@@ -254,7 +250,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 299,
+            "execution_count": 17,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -263,7 +259,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 300,
+            "execution_count": 18,
             "metadata": {},
             "outputs": [
                 {
@@ -1210,21 +1206,21 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 314,
+            "execution_count": 19,
             "metadata": {},
             "outputs": [
                 {
                     "name": "stdout",
                     "output_type": "stream",
                     "text": [
                         "blocks.0.attn.W_Q torch.Size([12, 768, 64])\n",
-                        "blocks.0.attn.W_K torch.Size([12, 768, 64])\n",
-                        "blocks.0.attn.W_V torch.Size([12, 768, 64])\n",
                         "blocks.0.attn.W_O torch.Size([12, 64, 768])\n",
                         "blocks.0.attn.b_Q torch.Size([12, 64])\n",
+                        "blocks.0.attn.b_O torch.Size([768])\n",
+                        "blocks.0.attn.W_K torch.Size([12, 768, 64])\n",
+                        "blocks.0.attn.W_V torch.Size([12, 768, 64])\n",
                         "blocks.0.attn.b_K torch.Size([12, 64])\n",
                         "blocks.0.attn.b_V torch.Size([12, 64])\n",
-                        "blocks.0.attn.b_O torch.Size([768])\n",
                         "blocks.0.mlp.W_in torch.Size([768, 3072])\n",
                         "blocks.0.mlp.b_in torch.Size([3072])\n",
                         "blocks.0.mlp.W_out torch.Size([3072, 768])\n",
@@ -1247,7 +1243,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 315,
+            "execution_count": 20,
             "metadata": {},
             "outputs": [
                 {