Mamba_training_steps.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "A100",
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "#1. Introduction to Mamba\n",
        "Mamba is a new architecture for LLM that can handle long sequences more efficiently than traditional models such as Transformers. It utilizes a Selective State Space Model (SSM) to dynamically filter and process information based on content, allowing the model to selectively remember or ignore parts of the input. Mamba offers significant improvements in processing speed and scaling capabilities, especially with longer sequences.\n",
        "\n",
        "But what really sets Mamba apart? Let’s test it out with an in-depth interactive experience with Mamba.\n",
        "\n",
        "#2. Mamba model chat\n",
        "While the current base implementation provides the familiar from_pretrained method and generated base parameters, some functionality (such as repetition_chamine) is not available. Also, we cannot use text-generation-webui like text-generation-webui( https://github.com/oobabooga/text-generation-webui) such a tool. So, in order to use Mamba, we will use Python code for inference. I've made the code as simple as possible.\n",
        "\n"
      ],
      "metadata": {
        "id": "IA9QKmxZMEzd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, let's load the model."
      ],
      "metadata": {
        "id": "WcHl46RPMLpl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install causal-conv1d==1.0.0\n",
        "!pip install mamba-ssm==1.0.1\n",
        "!pip install transformers\n",
        "!pip install datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cyilTEHhQSvA",
        "outputId": "e5dd3618-cf26-4c0e-d289-191aff26cfd8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting causal-conv1d==1.0.0\n",
            "  Downloading causal_conv1d-1.0.0.tar.gz (6.4 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from causal-conv1d==1.0.0) (2.1.0+cu121)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from causal-conv1d==1.0.0) (23.2)\n",
            "Collecting ninja (from causal-conv1d==1.0.0)\n",
            "  Downloading ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl (307 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.2/307.2 kB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch->causal-conv1d==1.0.0) (2.1.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->causal-conv1d==1.0.0) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->causal-conv1d==1.0.0) (1.3.0)\n",
            "Building wheels for collected packages: causal-conv1d\n",
            "  Building wheel for causal-conv1d (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for causal-conv1d: filename=causal_conv1d-1.0.0-cp310-cp310-linux_x86_64.whl size=9116761 sha256=4bbd2c2672ecd02c1e43f8e52552de593099619abc6dda18b2ac650e08110124\n",
            "  Stored in directory: /root/.cache/pip/wheels/9a/48/f5/eb0c6d6d8e00131eaa57927b537a23832b37e2f01b801d9c5d\n",
            "Successfully built causal-conv1d\n",
            "Installing collected packages: ninja, causal-conv1d\n",
            "Successfully installed causal-conv1d-1.0.0 ninja-1.11.1.1\n",
            "Collecting mamba-ssm==1.0.1\n",
            "  Downloading mamba_ssm-1.0.1.tar.gz (28 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from mamba-ssm==1.0.1) (2.1.0+cu121)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from mamba-ssm==1.0.1) (23.2)\n",
            "Requirement already satisfied: ninja in /usr/local/lib/python3.10/dist-packages (from mamba-ssm==1.0.1) (1.11.1.1)\n",
            "Collecting einops (from mamba-ssm==1.0.1)\n",
            "  Downloading einops-0.7.0-py3-none-any.whl (44 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: triton in /usr/local/lib/python3.10/dist-packages (from mamba-ssm==1.0.1) (2.1.0)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (from mamba-ssm==1.0.1) (4.35.2)\n",
            "Requirement already satisfied: causal_conv1d in /usr/local/lib/python3.10/dist-packages (from mamba-ssm==1.0.1) (1.0.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm==1.0.1) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm==1.0.1) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm==1.0.1) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm==1.0.1) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm==1.0.1) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm==1.0.1) (2023.6.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (0.20.2)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (1.23.5)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (2023.6.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (0.15.0)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (0.4.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm==1.0.1) (4.66.1)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->mamba-ssm==1.0.1) (2.1.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm==1.0.1) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm==1.0.1) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm==1.0.1) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm==1.0.1) (2023.11.17)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->mamba-ssm==1.0.1) (1.3.0)\n",
            "Building wheels for collected packages: mamba-ssm\n",
            "  Building wheel for mamba-ssm (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for mamba-ssm: filename=mamba_ssm-1.0.1-cp310-cp310-linux_x86_64.whl size=137567739 sha256=1775b610f76d6bc71ffaa72375df8c0afde52c1c14f1c788ad6afee4290adff2\n",
            "  Stored in directory: /root/.cache/pip/wheels/08/cf/65/cc589985f9689241fe2c154ce1c60738f58a24e76ce474cc20\n",
            "Successfully built mamba-ssm\n",
            "Installing collected packages: einops, mamba-ssm\n",
            "Successfully installed einops-0.7.0 mamba-ssm-1.0.1\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.35.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.13.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.20.2)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.23.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2023.6.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.15.0)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.66.1)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers) (2023.6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers) (4.5.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2023.11.17)\n",
            "Collecting datasets\n",
            "  Downloading datasets-2.16.1-py3-none-any.whl (507 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m507.1/507.1 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets) (3.13.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.23.5)\n",
            "Requirement already satisfied: pyarrow>=8.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (10.0.1)\n",
            "Requirement already satisfied: pyarrow-hotfix in /usr/local/lib/python3.10/dist-packages (from datasets) (0.6)\n",
            "Collecting dill<0.3.8,>=0.3.0 (from datasets)\n",
            "  Downloading dill-0.3.7-py3-none-any.whl (115 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m8.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (1.5.3)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.31.0)\n",
            "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.1)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets) (3.4.1)\n",
            "Collecting multiprocess (from datasets)\n",
            "  Downloading multiprocess-0.70.15-py310-none-any.whl (134 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m9.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: fsspec[http]<=2023.10.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (2023.6.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.9.1)\n",
            "Requirement already satisfied: huggingface-hub>=0.19.4 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.20.2)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (23.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (23.2.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.9.4)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.4.1)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.19.4->datasets) (4.5.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (2023.11.17)\n",
            "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2023.3.post1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.1->pandas->datasets) (1.16.0)\n",
            "Installing collected packages: dill, multiprocess, datasets\n",
            "Successfully installed datasets-2.16.1 dill-0.3.7 multiprocess-0.70.15\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Set the environment variable for LD_LIBRARY_PATH\n",
        "%env LD_LIBRARY_PATH=/usr/lib64-nvidia"
      ],
      "metadata": {
        "id": "C-HP-bvVMkOv",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "10c887d4-2af3-4f80-d010-f494246ef596"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "env: LD_LIBRARY_PATH=/usr/lib64-nvidia\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2o42Y7RR2iPy",
        "outputId": "989b5b1b-ffb3-42ea-9b48-df7784b73a2f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "xnKMI2AUw5d8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!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"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wk1VJpwxQWpZ",
        "outputId": "a1a13880-8f9f-4f32-f628-01a740bc79fd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_0.so.3 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbb.so.12 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc.so.2 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind.so.3 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc_proxy.so.2 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_5.so.3 is not a symbolic link\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "txg0LCnGMA0Z",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c1924638-c3c2-4903-fec2-f7fed5294bc3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "# Determine if a GPU is available and set the device accordingly\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Load the model\n",
        "model = MambaLMHeadModel.from_pretrained(\n",
        "    \"state-spaces/mamba-1.4b\",\n",
        "    device=device,  # Use the device variable here\n",
        "    dtype=torch.float16\n",
        ").to(device)  # Move the model to the specified device\n",
        "\n",
        "# Load the tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\")\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!nvcc --version"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aIa5bvOBuUEe",
        "outputId": "f200bab6-562e-49bc-fc7d-43c1477e0005"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "nvcc: NVIDIA (R) Cuda compiler driver\n",
            "Copyright (c) 2005-2023 NVIDIA Corporation\n",
            "Built on Tue_Aug_15_22:02:13_PDT_2023\n",
            "Cuda compilation tools, release 12.2, V12.2.140\n",
            "Build cuda_12.2.r12.2/compiler.33191640_0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Check NVIDIA GPU status\n",
        "!nvidia-smi"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HkplO00muT_j",
        "outputId": "c95f814d-dbe3-42c0-a3f3-237e65dda636"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mon Jan 22 10:58:00 2024       \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |\n",
            "|-----------------------------------------+----------------------+----------------------+\n",
            "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                                         |                      |               MIG M. |\n",
            "|=========================================+======================+======================|\n",
            "|   0  NVIDIA A100-SXM4-40GB          Off | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   31C    P0              51W / 400W |   3327MiB / 40960MiB |      0%      Default |\n",
            "|                                         |                      |             Disabled |\n",
            "+-----------------------------------------+----------------------+----------------------+\n",
            "                                                                                         \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| Processes:                                                                            |\n",
            "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
            "|        ID   ID                                                             Usage      |\n",
            "|=======================================================================================|\n",
            "+---------------------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "prompt=\\\n",
        "\"\"\"A conversation between a user and a smart AI assistant.\n",
        "​\n",
        "### User: Hello!\n",
        "### Assistant:\"\"\"\n",
        "\n",
        "prompt_tokenized=tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
        "\n",
        "# from https://github.com/state-spaces/mamba/blob/main/benchmarks/benchmark_generation_mamba_simple.py#L54\n",
        "output_tokenized = model.generate(\n",
        "    input_ids=prompt_tokenized[\"input_ids\"],\n",
        "    max_length=70,\n",
        "    cg=True,\n",
        "    output_scores=True,\n",
        "    enable_timing=False,\n",
        "    temperature=0.7,\n",
        "    top_k=40,\n",
        "    top_p=0.1,\n",
        "    )\n",
        "output=tokenizer.decode(output_tokenized[0])\n",
        "\n",
        "print(output)"
      ],
      "metadata": {
        "id": "kjIn_1VyMZ60",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "29fe097a-686b-4798-ba1a-19ae2b700738"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "A conversation between a user and a smart AI assistant.\n",
            "​\n",
            "### User: Hello!\n",
            "### Assistant: Hello!\n",
            "### User: I'm hungry!\n",
            "### Assistant: I'm hungry too!\n",
            "### User: I'm thirsty!\n",
            "### Assistant: I'm thirsty too!\n",
            "### User: I'm tired!\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Load the data set and tokenize it"
      ],
      "metadata": {
        "id": "FpKawuXsMnpG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset=load_dataset(\"OpenAssistant/oasst_top1_2023-08-25\")"
      ],
      "metadata": {
        "id": "DE6O69SjMpt9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "def tokenize(element):\n",
        "    return tokenizer(\n",
        "        element[\"text\"],\n",
        "        truncation=True,\n",
        "        max_length=1024,\n",
        "        add_special_tokens=False,\n",
        "    )\n",
        "\n",
        "\n",
        "dataset_tokenized = dataset.map(\n",
        "    tokenize,\n",
        "    batched=True,\n",
        "    num_proc=os.cpu_count(),    # multithreaded\n",
        "    remove_columns=[\"text\"]     # don't need this anymore, we have tokens from here on\n",
        ")"
      ],
      "metadata": {
        "id": "EZaDVo1qMprf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Define the collate function"
      ],
      "metadata": {
        "id": "Z5Xc2jZbM5qT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "# collate function - to transform list of dictionaries [ {input_ids: [123, ..]}, {.. ] to single batch dictionary { input_ids: [..], labels: [..], attention_mask: [..] }\n",
        "def collate(elements):\n",
        "    tokenlist=[e[\"input_ids\"] for e in elements]\n",
        "    tokens_maxlen=max([len(t) for t in tokenlist])\n",
        "\n",
        "    input_ids,labels = [],[]\n",
        "    for tokens in tokenlist:\n",
        "        pad_len=tokens_maxlen-len(tokens)\n",
        "\n",
        "        # pad input_ids with pad_token, labels with ignore_index (-100) and set attention_mask 1 where content otherwise 0\n",
        "        input_ids.append( tokens + [tokenizer.pad_token_id]*pad_len )\n",
        "        labels.append( tokens + [-100]*pad_len )\n",
        "\n",
        "    batch={\n",
        "        \"input_ids\": torch.tensor(input_ids),\n",
        "        \"labels\": torch.tensor(labels),\n",
        "    }\n",
        "    return batch"
      ],
      "metadata": {
        "id": "Cdi_o_2UMppG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Prepare Mamba Trainer"
      ],
      "metadata": {
        "id": "dJpTp_RONBzW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# monkey patch MambaLMHeadModel.forward\n",
        "def forward_with_loss(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, labels = None):\n",
        "    \"\"\"\n",
        "    \"position_ids\" is just to be compatible with Transformer generation. We don't use it.\n",
        "    num_last_tokens: if > 0, only return the logits for the last n tokens\n",
        "    \"\"\"\n",
        "    hidden_states = self.backbone(input_ids, inference_params=inference_params)\n",
        "    if num_last_tokens > 0:\n",
        "        hidden_states = hidden_states[:, -num_last_tokens:]\n",
        "    lm_logits = self.lm_head(hidden_states)\n",
        "\n",
        "    # Source: https://github.com/huggingface/transformers/blob/80377eb018c077dba434bc8e7912bcaed3a64d09/src/transformers/models/llama/modeling_llama.py#L1196\n",
        "    from torch.nn import CrossEntropyLoss\n",
        "    if labels is not None:\n",
        "        logits = lm_logits\n",
        "        # Shift so that tokens < n predict n\n",
        "        shift_logits = logits[..., :-1, :].contiguous()\n",
        "        shift_labels = labels[..., 1:].contiguous()\n",
        "        # Flatten the tokens\n",
        "        loss_fct = CrossEntropyLoss()\n",
        "        # shift_logits = shift_logits.view(-1, self.config.vocab_size)\n",
        "        shift_logits = shift_logits.view(-1, self.backbone.embedding.weight.size()[0])\n",
        "        shift_labels = shift_labels.view(-1)\n",
        "        # Enable model parallelism\n",
        "        shift_labels = shift_labels.to(shift_logits.device)\n",
        "        loss = loss_fct(shift_logits, shift_labels)\n",
        "        return (loss,)\n",
        "    else:\n",
        "        CausalLMOutput = namedtuple(\"CausalLMOutput\", [\"logits\"])\n",
        "        return CausalLMOutput(logits=lm_logits)\n",
        "MambaLMHeadModel.forward=forward_with_loss\n",
        "\n",
        "# patch MambaLMHeadModel\n",
        "MambaLMHeadModel.forward=forward_with_loss\n",
        "\n",
        "# (re)load model\n",
        "# model = MambaLMHeadModel.from_pretrained(\"state-spaces/mamba-1.4b\", device=\"cuda\", dtype=torch.float16)\n",
        "model = MambaLMHeadModel.from_pretrained(\"state-spaces/mamba-1.4b\", device=\"cuda\")\n"
      ],
      "metadata": {
        "id": "4q46M4QRMpms"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Training the Mamba model"
      ],
      "metadata": {
        "id": "2K-tF5W9NMRu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers[torch]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lTw0Tdx8rdwi",
        "outputId": "d0787509-2ff0-4faf-8311-aa7c824af7d3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: transformers[torch] in /usr/local/lib/python3.10/dist-packages (4.35.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (3.13.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (0.20.2)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (1.23.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (23.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (2023.6.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (0.15.0)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (0.4.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (4.66.1)\n",
            "Requirement already satisfied: torch!=1.12.0,>=1.10 in /usr/local/lib/python3.10/dist-packages (from transformers[torch]) (2.1.0+cu121)\n",
            "Collecting accelerate>=0.20.3 (from transformers[torch])\n",
            "  Downloading accelerate-0.26.1-py3-none-any.whl (270 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m270.9/270.9 kB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate>=0.20.3->transformers[torch]) (5.9.5)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers[torch]) (2023.6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers[torch]) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch!=1.12.0,>=1.10->transformers[torch]) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch!=1.12.0,>=1.10->transformers[torch]) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch!=1.12.0,>=1.10->transformers[torch]) (3.1.3)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch!=1.12.0,>=1.10->transformers[torch]) (2.1.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers[torch]) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers[torch]) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers[torch]) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers[torch]) (2023.11.17)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch!=1.12.0,>=1.10->transformers[torch]) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch!=1.12.0,>=1.10->transformers[torch]) (1.3.0)\n",
            "Installing collected packages: accelerate\n",
            "Successfully installed accelerate-0.26.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from transformers import Trainer, TrainingArguments\n",
        "\n",
        "# Create output directory\n",
        "output_dir = '/content/drive/MyDrive/mamba-1'\n",
        "if not os.path.exists(output_dir):\n",
        "    os.makedirs(output_dir)\n",
        "\n",
        "\n",
        "bs=4        # batch size\n",
        "ga_steps=1  # gradient acc. steps\n",
        "epochs=3\n",
        "steps_per_epoch=len(dataset_tokenized[\"train\"])//(bs*ga_steps)\n",
        "lr=0.0005\n",
        "\n",
        "args = TrainingArguments(\n",
        "    output_dir=output_dir,\n",
        "    per_device_train_batch_size=bs,\n",
        "    per_device_eval_batch_size=bs,\n",
        "    evaluation_strategy=\"steps\",\n",
        "    logging_steps=1,\n",
        "    eval_steps=steps_per_epoch,\n",
        "    save_steps=steps_per_epoch,\n",
        "    gradient_accumulation_steps=ga_steps,\n",
        "    num_train_epochs=epochs,\n",
        "    lr_scheduler_type=\"constant\",\n",
        "    learning_rate=lr,\n",
        "    group_by_length=True,\n",
        "    bf16=False,                  # mixed precision training\n",
        "    fp16=True,\n",
        "    save_safetensors=False,     # saving will fail without this\n",
        ")\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    args=args,\n",
        "    data_collator=collate,\n",
        "    train_dataset=dataset_tokenized[\"train\"],\n",
        "    eval_dataset=dataset_tokenized[\"test\"],\n",
        ")\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "py3cwjxUNKgq",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "outputId": "ccbc53af-27ba-4b5c-887f-07988b51c5d4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='9711' max='9711' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [9711/9711 53:24, Epoch 3/3]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>3236</td>\n",
              "      <td>2.314600</td>\n",
              "      <td>2.887654</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>6472</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>nan</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>9708</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>nan</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TrainOutput(global_step=9711, training_loss=1.7071560114827733, metrics={'train_runtime': 3208.7907, 'train_samples_per_second': 12.105, 'train_steps_per_second': 3.026, 'total_flos': 0.0, 'train_loss': 1.7071560114827733, 'epoch': 3.0})"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Next Steps and insights for Fine-Tuning and Evaluating the Mamba Model\n",
        "\n",
        "## 1. Adjust Learning Rate\n",
        "- Initial learning rate of 0.0005 was ineffective.\n",
        "- Reducing the learning rate to 0.00005 improved outcomes.\n",
        "\n",
        "## 2. Evaluate the Mamba Model\n",
        "- Evaluation is challenging due to subjective metrics.\n",
        "- Utilize benchmarks like [EleutherAI's lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness), Chatbot Arena, and artificial intelligence referee at [Chatbot Arena Leaderboard](https://huggingface.co/spaces/lmsys/chatbot-arena-leaderboard).\n",
        "\n",
        "## 3. Benchmarking and Skepticism\n",
        "- Used Mamba author's benchmarks as a starting point.\n",
        "- Skepticism around benchmark numbers without prior Mamba experience.\n",
        "\n",
        "## 4. Learning Rate Reconsideration\n",
        "- Original learning rate of 0.0005 potentially too high.\n",
        "- Lack of clarity on Mamba's pre-training learning rate.\n",
        "\n",
        "## 5. Further Fine-Tuning Experiments\n",
        "- Experimented with even lower learning rates (3x10e-5 and 2x10e-5).\n",
        "- Tested different training rounds and datasets, including the OA dataset and [HuggingFaceH4/ultrachat_200k](https://huggingface.co/datasets/ultrachat_200k).\n",
        "\n",
        "## 6. Comparison with TinyLlama\n",
        "- Noted significant speed advantage of Mamba over TinyLlama.\n",
        "- Mamba's lower VRAM usage and faster token generation rate.\n",
        "\n",
        "## 7. Long Context Capability of Mamba\n",
        "- Tested Mamba's ability to handle long prompts (up to 10k tokens).\n",
        "- Mamba struggles with very long texts (136K tokens), but performs better with shorter ones (1.54K tokens).\n",
        "- Example: Triathlon article [triathlon features](https://www.tri247.com/triathlon-features/interviews/lionel-sanders-championship-preview).\n",
        "\n",
        "## 8. Limitations in Generating High-Quality Content\n",
        "- Mamba's pre-training limited to 2048 tokens might hinder its ability to summarize large texts.\n",
        "- Suggestion to fine-tune smaller Mamba models for potential improvement.\n",
        "\n",
        "## 9. Summary\n",
        "- Mamba excels in speed and token handling capacity.\n",
        "- Fine-tuning Mamba is currently complex, with anticipation for future improvements.\n",
        "- TinyLlama generates better text, likely due to more extensive pre-training data.\n"
      ],
      "metadata": {
        "id": "bX5FBU8HNj2l"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KEQwZVTkNKeC"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}