{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "collapsed_sections": [ "YHmIf6af4Z-U", "1UEi0qad4kmP", "J2HsgqNC4tPr", "3ichua0a46T2", "v1acCbfA5AAT" ], "gpuType": "T4" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU", "widgets": { "application/vnd.jupyter.widget-state+json": { "53741c147b034fc7b5956e6c7ece37bf": { "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_4838166db3da46d4a85569c7d1ccdfcb", "IPY_MODEL_ba92622c9bc44b11b7a997ad3a8764bd", "IPY_MODEL_b509fc35f04841719bf76a721e0714b7" ], "layout": "IPY_MODEL_84a8e6cfed5945f8bbb6d9b6560a8ddb" } }, "4838166db3da46d4a85569c7d1ccdfcb": { "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_bee6dd0794514e03b389d53efbb65a2b", "placeholder": "​", "style": "IPY_MODEL_55169138a94c45f187c281f9c848ba4f", "value": "Datasets loaded (3/3): " } }, "ba92622c9bc44b11b7a997ad3a8764bd": { "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_23bc1161526b4150bbbdf1a88e9054d2", "max": 1, "min": 0, "orientation": "horizontal", "style": "IPY_MODEL_7e5554cbf74a43b7a4a1465a6676c3e3", "value": 1 } }, "b509fc35f04841719bf76a721e0714b7": { "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_69cb1ea890e1434687d4fb6a01c36168", "placeholder": "​", "style": "IPY_MODEL_ae11348757874e39a58530601b3f12b0", "value": " 3/? [00:00<00:00, 11.30it/s]" } }, "84a8e6cfed5945f8bbb6d9b6560a8ddb": { "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 } }, "bee6dd0794514e03b389d53efbb65a2b": { "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 } }, "55169138a94c45f187c281f9c848ba4f": { "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": "" } }, "23bc1161526b4150bbbdf1a88e9054d2": { "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": "20px" } }, "7e5554cbf74a43b7a4a1465a6676c3e3": { "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": "" } }, "69cb1ea890e1434687d4fb6a01c36168": { "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 } }, "ae11348757874e39a58530601b3f12b0": { "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": "" } }, "cd355626cc6f46e18c1a7522dd357242": { "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_d18bf8c6210541d8b3f26140c4da844e", "IPY_MODEL_22c17aca1ac0411d94de5b264a02da4e", "IPY_MODEL_a47d690037bb4c5cb4144d4597dfe48f" ], "layout": "IPY_MODEL_5a55b1c282bc4fe3bd0828ed7195d3a1" } }, "d18bf8c6210541d8b3f26140c4da844e": { "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_5909524c735f41c2ae112d2f978df8a0", "placeholder": "​", "style": "IPY_MODEL_3d95234372c949cc94c95d58db94cab7", "value": "Epoch: 24:   2%" } }, "22c17aca1ac0411d94de5b264a02da4e": { "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": "danger", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_48be40e9979c41aca1922c2bf98a2076", "max": 1000, "min": 0, "orientation": "horizontal", "style": "IPY_MODEL_e0515b940ce640a58acae622ee201496", "value": 24 } }, "a47d690037bb4c5cb4144d4597dfe48f": { "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_4296cdfac6f84f6d873c68ab77dc6db3", "placeholder": "​", "style": "IPY_MODEL_3d7dd31be9ab40e0b6163494a2b38a59", "value": " 24/1000 [00:19<12:55,  1.26it/s]" } }, "5a55b1c282bc4fe3bd0828ed7195d3a1": { "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 } }, "5909524c735f41c2ae112d2f978df8a0": { "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 } }, "3d95234372c949cc94c95d58db94cab7": { "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": "" } }, "48be40e9979c41aca1922c2bf98a2076": { "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 } }, "e0515b940ce640a58acae622ee201496": { "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": "" } }, "4296cdfac6f84f6d873c68ab77dc6db3": { "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 } }, "3d7dd31be9ab40e0b6163494a2b38a59": { "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": "markdown", "source": [ "# Edge Cases in Expressivity\n", "\n", "In this notebook we show the results of the counting experiments on single layer S4D and Mamba models. The models have been trained on time series of length 200 for 1k epochs each. The checkpoints we use here are the `best` with respect to predicting the offset.\n", "\n", "Here we provide results for the following experiments:\n", "\n", "\n", "1. Offset Prediction\n", "2. Consecutive ISI Phase Alignment Test\n", "3. Double ISI Phase Alignment Test\n", "\n", "\n", "This notebook takes ~30 mins to run. ~15 mins to train each model.\n", "\n", "**Note**\n", "Sometimes the notebook crashes after training Mamba, just reconnect and run the cells for training S4D onwards. If you refresh/reconnect you will lose the trained checkpoints etc." ], "metadata": { "id": "bjANxGBFzNp9" } }, { "cell_type": "markdown", "source": [ "## Ensure Torch and Cuda are matching Mamba Wheel" ], "metadata": { "id": "YHmIf6af4Z-U" } }, { "cell_type": "code", "source": [ "# ensure that torch and cuda version match the required version for the mamba whl\n", "import torch\n", "\n", "# Check current version\n", "current_version = torch.__version__\n", "target_version = \"2.6.0+cu124\"\n", "\n", "print(f\"Current PyTorch: {current_version}\")\n", "print(f\"Target PyTorch: {target_version}\")\n", "\n", "if current_version != target_version:\n", " print(\"Upgrading PyTorch...\")\n", " !pip install --upgrade torch==2.6.0+cu124 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124\n", " print(\"Upgrade complete! Please restart runtime.\")\n", "else:\n", " print(\"PyTorch version is correct!\")\n", "\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "import matplotlib.pyplot as plt\n", "from pathlib import Path\n", "import imageio.v2 as imageio # make sure this version works on Colab" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "V-RVtcE8wa-Y", "outputId": "bd44b726-8183-4064-f4ab-82ed1ee76b27" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Current PyTorch: 2.6.0+cu124\n", "Target PyTorch: 2.6.0+cu124\n", "PyTorch version is correct!\n" ] } ] }, { "cell_type": "markdown", "source": [ "## Install pre-built Mamba Wheel" ], "metadata": { "id": "1UEi0qad4kmP" } }, { "cell_type": "code", "source": [ "# install the mamba wheel pre-built from repo\n", "!pip install https://github.com/jayesh365/mamba-wheels/releases/download/v2.2.2-cu124/mamba_ssm-2.2.2-cp311-cp311-linux_x86_64.whl" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "2IO4hlW2hU68", "outputId": "d968496d-05ba-4bbe-a8ff-04f310872693" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Collecting mamba-ssm==2.2.2\n", " Downloading https://github.com/jayesh365/mamba-wheels/releases/download/v2.2.2-cu124/mamba_ssm-2.2.2-cp311-cp311-linux_x86_64.whl (323.5 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m323.5/323.5 MB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hRequirement already satisfied: torch in /usr/local/lib/python3.11/dist-packages (from mamba-ssm==2.2.2) (2.6.0+cu124)\n", "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from mamba-ssm==2.2.2) (24.2)\n", "Collecting ninja (from mamba-ssm==2.2.2)\n", " Downloading ninja-1.11.1.4-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl.metadata (5.0 kB)\n", "Requirement already satisfied: einops in /usr/local/lib/python3.11/dist-packages (from mamba-ssm==2.2.2) (0.8.1)\n", "Requirement already satisfied: triton in /usr/local/lib/python3.11/dist-packages (from mamba-ssm==2.2.2) (3.2.0)\n", "Requirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (from mamba-ssm==2.2.2) (4.51.3)\n", "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (3.18.0)\n", "Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (4.13.2)\n", "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (3.4.2)\n", "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (3.1.6)\n", "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (2025.3.2)\n", "Collecting nvidia-cuda-nvrtc-cu12==12.4.127 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", "Collecting nvidia-cuda-runtime-cu12==12.4.127 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", "Collecting nvidia-cuda-cupti-cu12==12.4.127 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", "Collecting nvidia-cudnn-cu12==9.1.0.70 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", "Collecting nvidia-cublas-cu12==12.4.5.8 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", "Collecting nvidia-cufft-cu12==11.2.1.3 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", "Collecting nvidia-curand-cu12==10.3.5.147 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", "Collecting nvidia-cusolver-cu12==11.6.1.9 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", "Collecting nvidia-cusparse-cu12==12.3.1.170 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.2 in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (0.6.2)\n", "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (2.21.5)\n", "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (12.4.127)\n", "Collecting nvidia-nvjitlink-cu12==12.4.127 (from torch->mamba-ssm==2.2.2)\n", " Downloading nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch->mamba-ssm==2.2.2) (1.13.1)\n", "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch->mamba-ssm==2.2.2) (1.3.0)\n", "Requirement already satisfied: huggingface-hub<1.0,>=0.30.0 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (0.31.2)\n", "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (2.0.2)\n", "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (6.0.2)\n", "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (2024.11.6)\n", "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (2.32.3)\n", "Requirement already satisfied: tokenizers<0.22,>=0.21 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (0.21.1)\n", "Requirement already satisfied: safetensors>=0.4.3 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (0.5.3)\n", "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers->mamba-ssm==2.2.2) (4.67.1)\n", "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch->mamba-ssm==2.2.2) (3.0.2)\n", "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->transformers->mamba-ssm==2.2.2) (3.4.2)\n", "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests->transformers->mamba-ssm==2.2.2) (3.10)\n", "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->transformers->mamba-ssm==2.2.2) (2.4.0)\n", "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests->transformers->mamba-ssm==2.2.2) (2025.4.26)\n", "Downloading ninja-1.11.1.4-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (422 kB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m422.8/422.8 kB\u001b[0m \u001b[31m13.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl (363.4 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (13.8 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m66.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (24.6 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m36.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (883 kB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m53.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl (664.8 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl (211.5 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m5.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl (56.3 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m12.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl (127.9 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl (207.5 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m78.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hInstalling collected packages: nvidia-nvjitlink-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, ninja, nvidia-cusparse-cu12, nvidia-cudnn-cu12, nvidia-cusolver-cu12, mamba-ssm\n", " Attempting uninstall: nvidia-nvjitlink-cu12\n", " Found existing installation: nvidia-nvjitlink-cu12 12.5.82\n", " Uninstalling nvidia-nvjitlink-cu12-12.5.82:\n", " Successfully uninstalled nvidia-nvjitlink-cu12-12.5.82\n", " Attempting uninstall: nvidia-curand-cu12\n", " Found existing installation: nvidia-curand-cu12 10.3.6.82\n", " Uninstalling nvidia-curand-cu12-10.3.6.82:\n", " Successfully uninstalled nvidia-curand-cu12-10.3.6.82\n", " Attempting uninstall: nvidia-cufft-cu12\n", " Found existing installation: nvidia-cufft-cu12 11.2.3.61\n", " Uninstalling nvidia-cufft-cu12-11.2.3.61:\n", " Successfully uninstalled nvidia-cufft-cu12-11.2.3.61\n", " Attempting uninstall: nvidia-cuda-runtime-cu12\n", " Found existing installation: nvidia-cuda-runtime-cu12 12.5.82\n", " Uninstalling nvidia-cuda-runtime-cu12-12.5.82:\n", " Successfully uninstalled nvidia-cuda-runtime-cu12-12.5.82\n", " Attempting uninstall: nvidia-cuda-nvrtc-cu12\n", " Found existing installation: nvidia-cuda-nvrtc-cu12 12.5.82\n", " Uninstalling nvidia-cuda-nvrtc-cu12-12.5.82:\n", " Successfully uninstalled nvidia-cuda-nvrtc-cu12-12.5.82\n", " Attempting uninstall: nvidia-cuda-cupti-cu12\n", " Found existing installation: nvidia-cuda-cupti-cu12 12.5.82\n", " Uninstalling nvidia-cuda-cupti-cu12-12.5.82:\n", " Successfully uninstalled nvidia-cuda-cupti-cu12-12.5.82\n", " Attempting uninstall: nvidia-cublas-cu12\n", " Found existing installation: nvidia-cublas-cu12 12.5.3.2\n", " Uninstalling nvidia-cublas-cu12-12.5.3.2:\n", " Successfully uninstalled nvidia-cublas-cu12-12.5.3.2\n", " Attempting uninstall: nvidia-cusparse-cu12\n", " Found existing installation: nvidia-cusparse-cu12 12.5.1.3\n", " Uninstalling nvidia-cusparse-cu12-12.5.1.3:\n", " Successfully uninstalled nvidia-cusparse-cu12-12.5.1.3\n", " Attempting uninstall: nvidia-cudnn-cu12\n", " Found existing installation: nvidia-cudnn-cu12 9.3.0.75\n", " Uninstalling nvidia-cudnn-cu12-9.3.0.75:\n", " Successfully uninstalled nvidia-cudnn-cu12-9.3.0.75\n", " Attempting uninstall: nvidia-cusolver-cu12\n", " Found existing installation: nvidia-cusolver-cu12 11.6.3.83\n", " Uninstalling nvidia-cusolver-cu12-11.6.3.83:\n", " Successfully uninstalled nvidia-cusolver-cu12-11.6.3.83\n", "Successfully installed mamba-ssm-2.2.2 ninja-1.11.1.4 nvidia-cublas-cu12-12.4.5.8 nvidia-cuda-cupti-cu12-12.4.127 nvidia-cuda-nvrtc-cu12-12.4.127 nvidia-cuda-runtime-cu12-12.4.127 nvidia-cudnn-cu12-9.1.0.70 nvidia-cufft-cu12-11.2.1.3 nvidia-curand-cu12-10.3.5.147 nvidia-cusolver-cu12-11.6.1.9 nvidia-cusparse-cu12-12.3.1.170 nvidia-nvjitlink-cu12-12.4.127\n" ] } ] }, { "cell_type": "markdown", "source": [ "## Clone Checkpoints and other files" ], "metadata": { "id": "J2HsgqNC4tPr" } }, { "cell_type": "code", "source": [ "# get the model checkpoints from repo\n", "!git clone https://github.com/jayesh365/mamba-wheels.git\n", "# move dir\n", "%cd /content/mamba-wheels/" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "OKsRtiKbE8v_", "outputId": "473eb175-2d4e-4446-fa95-e7e5421eac8b" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Cloning into 'mamba-wheels'...\n", "remote: Enumerating objects: 159, done.\u001b[K\n", "remote: Counting objects: 100% (159/159), done.\u001b[K\n", "remote: Compressing objects: 100% (114/114), done.\u001b[K\n", "remote: Total 159 (delta 70), reused 130 (delta 41), pack-reused 0 (from 0)\u001b[K\n", "Receiving objects: 100% (159/159), 179.14 KiB | 2.71 MiB/s, done.\n", "Resolving deltas: 100% (70/70), done.\n", "/content/mamba-wheels\n" ] } ] }, { "cell_type": "markdown", "source": [ "## Verify Mamba is imported correctly" ], "metadata": { "id": "3ichua0a46T2" } }, { "cell_type": "code", "source": [ "# check mamba\n", "try:\n", " from mamba_bits import Mamba_model as Mamba\n", " print('Successfully imported Mamba!')\n", "except ImportError as e:\n", " print('Could not import Mamba!', e)\n" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "Sxmok3NKNHyD", "outputId": "3d8ad621-cb40-4081-be1d-3b16647dbfd0" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Successfully imported Mamba!\n" ] } ] }, { "cell_type": "markdown", "source": [ "## Verify S4D is imported correctly" ], "metadata": { "id": "v1acCbfA5AAT" } }, { "cell_type": "code", "source": [ "# check s4d\n", "try:\n", " from s4d_bits import S4DTokenClassifier as S4D\n", " print('Successfully imported S4D!')\n", "except ImportError as e:\n", " print('Could not import S4D! ', e)" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "trQ4ExZVFHGH", "outputId": "27c776dd-9165-4dc2-c134-5fa9ec4fc8df" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Successfully imported S4D!\n" ] } ] }, { "cell_type": "markdown", "source": [ "# Modular Counting" ], "metadata": { "id": "6EzDXVb8Mszk" } }, { "cell_type": "markdown", "source": [ "## Train models" ], "metadata": { "id": "pVPXhoErVUFZ" } }, { "cell_type": "code", "source": [ "%cd /content/mamba-wheels/" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "Za4QyyOPWQBw", "outputId": "f248b86b-e9a9-42c2-d3e2-2ccfb2db071b" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "/content/mamba-wheels\n" ] } ] }, { "cell_type": "code", "source": [ "from helpers.auxs import generate_inter_trial_interval, split_train_val_og, AlternatingSignalDataset, split_train_val, DataLoader\n", "import numpy as np\n", "hld_ot_int = np.random.randint(0, 99, size=200).tolist()\n", "train_inputs, train_targets, test_inputs, test_targets = generate_inter_trial_interval(num_seq=1000, n=200, iti=(20, 40), isi=10, holdout_intervals=hld_ot_int)\n", "train_inputs_split, train_targets_split, val_inputs_split, val_targets_split = split_train_val_og(train_inputs, train_targets, val_split=0.1)\n", "\n", "train_set = AlternatingSignalDataset(train_inputs_split, train_targets_split)\n", "val_set = AlternatingSignalDataset(val_inputs_split, val_targets_split)\n", "test_set = AlternatingSignalDataset(test_inputs, test_targets)\n", "# load datasets\n", "trainset, _ = split_train_val(train_set, val_split=0.2)\n", "_, valset = split_train_val(val_set, val_split=0.1)\n", "testset, _ = split_train_val(test_set, val_split=0)\n", "\n", "test_loader = DataLoader(testset, batch_size=64, shuffle=False)\n", "print(test_loader)" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "JJvybXtOVoYt", "outputId": "76372a6c-2da7-41d4-8df4-4c7ad50c340a" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "\n" ] } ] }, { "cell_type": "markdown", "source": [ "### Train Mamba model" ], "metadata": { "id": "O-awAJ1J2eG7" } }, { "cell_type": "code", "source": [ "from trainer import trainer" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "L9_Wyt246AsG", "outputId": "3a7d8dba-47a8-4f98-af72-b31d295d7b0d" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Successfully imported Mamba and S4D models\n", "Mamba: \n", "S4D: \n", "/content/mamba-wheels\n" ] } ] }, { "cell_type": "code", "source": [ "%matplotlib inline\n", "trainer(\n", " model_name='MAMBA',\n", " hidden_s=8,\n", " state_size=64,\n", " memory_size=64,\n", " epochs_length=1000,\n", " seed=123,\n", " ckpt_dir='./mamba_ckpts/',\n", " ts_length=200,\n", " lr=0.01\n", ")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 663, "referenced_widgets": [ "53741c147b034fc7b5956e6c7ece37bf", "4838166db3da46d4a85569c7d1ccdfcb", "ba92622c9bc44b11b7a997ad3a8764bd", "b509fc35f04841719bf76a721e0714b7", "84a8e6cfed5945f8bbb6d9b6560a8ddb", "bee6dd0794514e03b389d53efbb65a2b", "55169138a94c45f187c281f9c848ba4f", "23bc1161526b4150bbbdf1a88e9054d2", "7e5554cbf74a43b7a4a1465a6676c3e3", "69cb1ea890e1434687d4fb6a01c36168", "ae11348757874e39a58530601b3f12b0", "cd355626cc6f46e18c1a7522dd357242", "d18bf8c6210541d8b3f26140c4da844e", "22c17aca1ac0411d94de5b264a02da4e", "a47d690037bb4c5cb4144d4597dfe48f", "5a55b1c282bc4fe3bd0828ed7195d3a1", "5909524c735f41c2ae112d2f978df8a0", "3d95234372c949cc94c95d58db94cab7", "48be40e9979c41aca1922c2bf98a2076", "e0515b940ce640a58acae622ee201496", "4296cdfac6f84f6d873c68ab77dc6db3", "3d7dd31be9ab40e0b6163494a2b38a59" ] }, "id": "dDsYnx7EXNx7", "outputId": "bf47eaa2-3c25-4e9c-fe3d-c4077b7b04e2" }, "execution_count": null, "outputs": [ { "output_type": "display_data", "data": { "text/plain": [ "0it [00:00, ?it/s]" ], "application/vnd.jupyter.widget-view+json": { "version_major": 2, "version_minor": 0, "model_id": "53741c147b034fc7b5956e6c7ece37bf" } }, "metadata": {} }, { "output_type": "stream", "name": "stdout", "text": [ "\n", "Making a new set of Train, Test and Validation data.\n", "\n", "Datasets created!\n", "\n", "Making a new set of Train, Test and Validation data.\n", "\n", "Datasets created!\n", "\n", "Making a new set of Train, Test and Validation data.\n", "\n", "Datasets created!\n", "\n", "training MAMBA-123\n" ] }, { "output_type": "display_data", "data": { "text/plain": [ " 0%| | 0/1000 [00:00\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0mget_ipython\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_line_magic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'matplotlib'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'inline'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m trainer(\n\u001b[0m\u001b[1;32m 3\u001b[0m \u001b[0mmodel_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'MAMBA'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 4\u001b[0m \u001b[0mhidden_s\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m8\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 5\u001b[0m \u001b[0mstate_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m64\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m/content/mamba-wheels/trainer.py\u001b[0m in \u001b[0;36mtrainer\u001b[0;34m(model_name, hidden_s, state_size, memory_size, epochs_length, seed, ckpt_dir, ts_length, lr, clip_grad)\u001b[0m\n\u001b[1;32m 135\u001b[0m \u001b[0mpbar\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_description\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Epoch: %d'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mepoch\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 136\u001b[0m \u001b[0;31m# train(model, trainloader, device, optimizer, criterion, epoch, model_name, clip_grad)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 137\u001b[0;31m train(model, trainloader, device, optimizer, criterion, epoch, model_name, clip_grad,\n\u001b[0m\u001b[1;32m 138\u001b[0m all_mse_list, all_loss_list)\n\u001b[1;32m 139\u001b[0m \u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalloader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclip_grad\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m/content/mamba-wheels/helpers/auxs.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(model, trainloader, device, optimizer, criterion, epoch, model_name, clip_grad, all_mse_list, all_loss_list)\u001b[0m\n\u001b[1;32m 445\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 446\u001b[0m \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\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--> 447\u001b[0;31m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\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 448\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mclip_grad\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 449\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclip_grad_norm_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclip_grad\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.11/dist-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m 624\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 625\u001b[0m )\n\u001b[0;32m--> 626\u001b[0;31m torch.autograd.backward(\n\u001b[0m\u001b[1;32m 627\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 628\u001b[0m )\n", "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m 345\u001b[0m \u001b[0;31m# some Python versions print out the first line of a multi-line function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 346\u001b[0m \u001b[0;31m# calls in the traceback and some print out the last line\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 347\u001b[0;31m _engine_run_backward(\n\u001b[0m\u001b[1;32m 348\u001b[0m \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 349\u001b[0m \u001b[0mgrad_tensors_\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.11/dist-packages/torch/autograd/graph.py\u001b[0m in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m 821\u001b[0m \u001b[0munregister_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_register_logging_hooks_on_whole_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt_outputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 822\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 823\u001b[0;31m return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[1;32m 824\u001b[0m \u001b[0mt_outputs\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[0m\n\u001b[1;32m 825\u001b[0m ) # Calls into the C++ engine to run the backward pass\n", "\u001b[0;31mKeyboardInterrupt\u001b[0m: " ] } ] }, { "cell_type": "code", "source": [ "%cd /content/mamba-wheels/" ], "metadata": { "id": "KwDxoREyGS7Z" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "from pathlib import Path\n", "from model_eval import evaluate_custom_models\n", "from helpers.plotters import plot_results\n", "\n", "# Define your checkpoint base directory\n", "ckpt_base = Path(\"./mamba_ckpts\")\n", "\n", "# Collect all .pth checkpoints recursively\n", "checkpoint_paths = sorted([str(p) for p in ckpt_base.glob(\"**/*.pth\")])\n", "print(f\"Found {len(checkpoint_paths)} checkpoint files\")\n", "\n", "# Define your experiments\n", "phases = {\n", " \"offset\": '(False, False)',\n", " \"consec\": '(True, True)',\n", " \"double\": '(True, False)',\n", "}\n", "\n", "# Run evaluation for each phase\n", "results = []\n", "for name, phase in phases.items():\n", " print(f\"\\nRunning evaluation for: {name} phase\")\n", " result = evaluate_custom_models(checkpoint_paths, seq_len=200, phase=phase)\n", " results.append(result)\n", "\n", "# Plot results\n", "for result in results:\n", " if result:\n", " print(f\"Processing experiment: {result['type_phase']}\")\n", " print(\"Generating plots...\")\n", " figures = plot_results(result, save_plots=True)\n", " print(\"All done!\")\n", " else:\n", " print(\"No results to plot - check your model paths and dependencies!\")\n" ], "metadata": { "collapsed": true, "id": "SgUXNHPYswBE" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### Train S4D model" ], "metadata": { "id": "Uk6a8trH2k9t" } }, { "cell_type": "code", "source": [ "%cd /content/mamba-wheels/" ], "metadata": { "id": "sOTDS9kHCcfM" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "from trainer import trainer" ], "metadata": { "id": "MNbV40msVSLW" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "%matplotlib inline\n", "trainer(\n", " model_name='S4D',\n", " hidden_s=8,\n", " state_size=64,\n", " memory_size=64,\n", " epochs_length=1000,\n", " seed=879965,\n", " ckpt_dir='./s4d_ckpts/',\n", " ts_length=200,\n", " lr=0.005\n", ")" ], "metadata": { "id": "tUlWcaoXsyt4" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "from pathlib import Path\n", "from model_eval import evaluate_custom_models\n", "from helpers.plotters import plot_results\n", "\n", "# Define your checkpoint base directory\n", "ckpt_base = Path(\"./s4d_ckpts\")\n", "\n", "# Collect all .pth checkpoints recursively\n", "checkpoint_paths = sorted([str(p) for p in ckpt_base.glob(\"**/*.pth\")])\n", "print(f\"Found {len(checkpoint_paths)} checkpoint files\")\n", "\n", "# Define your experiments\n", "phases = {\n", " \"offset\": '(False, False)',\n", " \"consec\": '(True, True)',\n", " \"double\": '(True, False)',\n", "}\n", "\n", "# Run evaluation for each phase\n", "results = []\n", "for name, phase in phases.items():\n", " print(f\"\\nRunning evaluation for: {name} phase\")\n", " result = evaluate_custom_models(checkpoint_paths, seq_len=200, phase=phase)\n", " results.append(result)\n", "\n", "# Plot results\n", "for result in results:\n", " if result:\n", " print(f\"Processing experiment: {result['type_phase']}\")\n", " print(\"Generating plots...\")\n", " figures = plot_results(result, save_plots=True)\n", " print(\"All done!\")\n", " else:\n", " print(\"No results to plot - check your model paths and dependencies!\")\n", "2" ], "metadata": { "id": "qkdTeVA82sck" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "## Plot results for both models\n" ], "metadata": { "id": "88A-MIg09co8" } }, { "cell_type": "code", "source": [ "%cd /content/mamba-wheels/" ], "metadata": { "id": "YENsoqrECklW" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "from pathlib import Path\n", "from model_eval import evaluate_custom_models\n", "from helpers.plotters import plot_joint_outputs_from_two_runs\n", "\n", "# Define checkpoint directories\n", "mamba_ckpt_base = Path(\"./mamba_ckpts\")\n", "s4d_ckpt_base = Path(\"./s4d_ckpts\")\n", "\n", "# Collect all .pth checkpoints recursively\n", "mamba_ckpts = sorted([str(p) for p in mamba_ckpt_base.glob(\"**/*.pth\")])\n", "s4d_ckpts = sorted([str(p) for p in s4d_ckpt_base.glob(\"**/*.pth\")])\n", "print(f\"Found {len(mamba_ckpts)} MAMBA and {len(s4d_ckpts)} S4D checkpoint files\")\n", "\n", "# Define your experiments\n", "phases = {\n", " \"offset\": '(False, False)',\n", " \"consec\": '(True, True)',\n", " \"double\": '(True, False)',\n", "}\n", "\n", "# Run evaluation for each phase\n", "mamba_results = []\n", "s4d_results = []\n", "\n", "for name, phase in phases.items():\n", " print(f\"\\nEvaluating phase: {name}\")\n", "\n", " print(\"Running MAMBA...\")\n", " mamba_result = evaluate_custom_models(mamba_ckpts, seq_len=200, phase=phase)\n", " mamba_results.append(mamba_result)\n", "\n", " print(\"Running S4D...\")\n", " s4d_result = evaluate_custom_models(s4d_ckpts, seq_len=200, phase=phase)\n", " s4d_results.append(s4d_result)\n", "\n", "# Plot MAMBA vs S4D outputs together\n", "print(\"\\nPlotting MAMBA vs S4D outputs...\")\n", "plot_joint_outputs_from_two_runs(mamba_results, s4d_results, model_names=(\"MAMBA\", \"S4D\"), sample_idx=3)\n" ], "metadata": { "id": "KmO9oIvNwY1_" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "# Parity" ], "metadata": { "id": "yJn5XFGRM8qM" } }, { "cell_type": "markdown", "source": [ "## Loading" ], "metadata": { "id": "oa_QUf8UM9_8" } }, { "cell_type": "code", "source": [ "# 1. Install PyTorch (with CUDA if needed)\n", "# !pip install torch --upgrade\n", "\n", "# 2. Install mamba-ssm\n", "# !pip install git+https://github.com/state-spaces/mamba.git" ], "metadata": { "id": "2utGlZDTPuDf" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "! pip install lrcurve" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "2F4OXnQwvimc", "outputId": "2aa59ceb-c615-4d61-c6a9-781a8cc9ed8f" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Collecting lrcurve\n", " Downloading lrcurve-2.2.1-py3-none-any.whl.metadata (4.9 kB)\n", "Requirement already satisfied: ipython in /usr/local/lib/python3.11/dist-packages (from lrcurve) (7.34.0)\n", "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (75.2.0)\n", "Collecting jedi>=0.16 (from ipython->lrcurve)\n", " Downloading jedi-0.19.2-py2.py3-none-any.whl.metadata (22 kB)\n", "Requirement already satisfied: decorator in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (4.4.2)\n", "Requirement already satisfied: pickleshare in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (0.7.5)\n", "Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (5.7.1)\n", "Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (3.0.51)\n", "Requirement already satisfied: pygments in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (2.19.1)\n", "Requirement already satisfied: backcall in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (0.2.0)\n", "Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (0.1.7)\n", "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.11/dist-packages (from ipython->lrcurve) (4.9.0)\n", "Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.11/dist-packages (from jedi>=0.16->ipython->lrcurve) (0.8.4)\n", "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.11/dist-packages (from pexpect>4.3->ipython->lrcurve) (0.7.0)\n", "Requirement already satisfied: wcwidth in /usr/local/lib/python3.11/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython->lrcurve) (0.2.13)\n", "Downloading lrcurve-2.2.1-py3-none-any.whl (71 kB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m71.0/71.0 kB\u001b[0m \u001b[31m3.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hDownloading jedi-0.19.2-py2.py3-none-any.whl (1.6 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m27.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hInstalling collected packages: jedi, lrcurve\n", "Successfully installed jedi-0.19.2 lrcurve-2.2.1\n" ] } ] }, { "cell_type": "code", "source": [ "import random\n", "import itertools\n", "import torch\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "from torch.utils.data import DataLoader, Dataset\n", "from sklearn.decomposition import PCA\n", "import numpy as np\n", "from lrcurve import PlotLearningCurve" ], "metadata": { "id": "XJncw0fFM9mg" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "device" ], "metadata": { "id": "CMKge41mNB-I", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "ad6d13df-e1b8-41c6-846b-829fc675c613" }, "execution_count": null, "outputs": [ { "output_type": "execute_result", "data": { "text/plain": [ "device(type='cuda')" ] }, "metadata": {}, "execution_count": 5 } ] }, { "cell_type": "markdown", "source": [ "## Dataset Creation" ], "metadata": { "id": "1KkCs5uoNBTJ" } }, { "cell_type": "code", "source": [ "# data generator for mod-n parity\n", "class Mod_n_Dataset(Dataset):\n", " def __init__(self, num_samples, list_length, n):\n", " self.num_samples = num_samples\n", " self.list_length = list_length\n", " self.n = n\n", "\n", " def __len__(self):\n", " return self.num_samples\n", "\n", " def __getitem__(self, idx):\n", " random_list = [random.choice([0, 1]) for _ in range(self.list_length)]\n", " cumulative_sum = list(itertools.accumulate(random_list))\n", " mod_n_list = [x % n for x in cumulative_sum]\n", "\n", " # Convert to torch tensors\n", " random_list_tensor = torch.tensor(random_list, dtype=torch.int64)\n", " mod_n_list_tensor = torch.tensor(mod_n_list, dtype=torch.int64)\n", "\n", " return random_list_tensor, mod_n_list_tensor\n", "\n", "\n", "def generate_data_n_parity(batch_size, list_length, num_samples, n):\n", " dataset = Mod_n_Dataset(num_samples, list_length, n)\n", " dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n", " return dataloader" ], "metadata": { "id": "j-_2RJvqNE17" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "batch_size = 512 # Number of sequences per batch\n", "list_length = 8 # Length of each sequence\n", "num_samples = 10000000 # Total number of samples\n", "n = 2\n", "\n", "# Generate data\n", "dataloader_mod_n = generate_data_n_parity(batch_size, 8, 10_000_000, n)\n", "test_dataloader_n = generate_data_n_parity(batch_size, 16384, 1_000, n)\n", "\n", "# Iterate through the data\n", "input_batch, output_batch = next(iter(dataloader_mod_n))\n", "print(\"Train Input batch:\", input_batch[0])\n", "print(\"Train Output batch:\", output_batch[0])\n", "\n", "input_batch, output_batch = next(iter(test_dataloader_n))\n", "print(\"Test Input batch:\", input_batch[0])\n", "print(\"Test Output batch:\", output_batch[0])" ], "metadata": { "id": "C5vNexXUNGK3", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "e9df50ec-aad1-4218-eed8-12a31eda4799" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Train Input batch: tensor([0, 0, 0, 0, 1, 1, 1, 1])\n", "Train Output batch: tensor([0, 0, 0, 0, 1, 0, 1, 0])\n", "Test Input batch: tensor([1, 1, 0, ..., 0, 1, 0])\n", "Test Output batch: tensor([1, 0, 0, ..., 0, 1, 1])\n" ] } ] }, { "cell_type": "markdown", "source": [ "## Models" ], "metadata": { "id": "ctyFkobcNSJk" } }, { "cell_type": "markdown", "source": [ "### RNN" ], "metadata": { "id": "NhgSedexNiGI" } }, { "cell_type": "code", "source": [ "# Code copied from Jelassi's paper, adapted to RNN :\n", "# https://github.com/sjelassi/transformers_ssm_copy/blob/main/synthetic_exps/models/lstm.py\n", "\n", "import torch\n", "import math\n", "import torch.nn as nn\n", "\n", "class RNN(nn.Module):\n", "\n", " def __init__(self, embedding_dim=2, hidden_size =2, vocab_size=2, output_size=2, num_layers=1, dropout_rate=0.0, non_linearity=\"tanh\", embedding=True):\n", " super(RNN, self).__init__()\n", "\n", " self.num_layers = num_layers\n", " self.n_embd = embedding_dim\n", " self.embedding = embedding\n", " self.hidden_size = hidden_size\n", "\n", " if embedding:\n", " self.word_embeddings = nn.Embedding(vocab_size, self.n_embd)\n", "\n", " # The RNN takes word embeddings as input, and outputs hidden states\n", " # with dimensionality hidden_dim.\n", " self.rnn = nn.RNN(input_size=self.n_embd,\n", " hidden_size=self.hidden_size,\n", " num_layers=self.num_layers,\n", " batch_first=True,\n", " dropout=dropout_rate,\n", " nonlinearity=non_linearity)\n", "\n", " # The linear layer that maps from hidden state space to tag space\n", " self.head = nn.Linear(self.hidden_size, output_size)\n", "\n", " self.init_weights()\n", "\n", " def forward(self, sentence, prev_state=None):\n", " if self.embedding:\n", " embeds = self.word_embeddings(sentence)\n", " else:\n", " embeds = torch.nn.functional.one_hot(sentence, num_classes=-1).float()\n", " # print(f\"embeddings have shape {embeds.shape}\")\n", " state = self.init_hidden(sentence.shape[0], sentence.device)\n", "\n", " rnn_out, state = self.rnn(embeds)#, prev_state)\n", " # print(f\"lstm_out has shape {lstm_out.shape}\")\n", " # print(f\"state has shape {state[0].shape}\")\n", "\n", " logits = self.head(rnn_out)\n", " # print(f\"logits has shape {logits.shape}\")\n", " return logits, rnn_out # state\n", "\n", " def init_weights(self):\n", " init_range_emb = 0.1\n", " init_range_other = 1/math.sqrt(self.n_embd)\n", " if self.embedding:\n", " self.word_embeddings.weight.data.uniform_(-init_range_emb, init_range_emb)\n", " self.head.weight.data.uniform_(-init_range_other, init_range_other)\n", " self.head.bias.data.zero_()\n", " for i in range(self.num_layers):\n", " self.rnn.all_weights[i][0] = torch.FloatTensor(self.n_embd,\n", " self.n_embd).uniform_(-init_range_other, init_range_other)\n", " self.rnn.all_weights[i][1] = torch.FloatTensor(self.n_embd,\n", " self.n_embd).uniform_(-init_range_other, init_range_other)\n", "\n", " def init_hidden(self, batch_size, device):\n", " hidden = torch.zeros(self.num_layers, batch_size, self.n_embd).to(device)\n", " cell = torch.zeros(self.num_layers, batch_size, self.n_embd).to(device)\n", " return hidden, cell\n", "\n" ], "metadata": { "id": "igozXuqGNjAC" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### S4D" ], "metadata": { "id": "gTl0ogMONhDU" } }, { "cell_type": "code", "source": [ "# from this repo: https://github.com/state-spaces/s4/blob/main/models/s4/s4d.py\n", "\n", "\"\"\"Utility nn components, in particular handling activations, initializations, and normalization layers.\"\"\"\n", "\n", "from functools import partial\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import math\n", "from einops import rearrange, repeat\n", "\n", "\n", "class DropoutNd(nn.Module):\n", " def __init__(self, p: float = 0.5, tie=True, transposed=True):\n", " \"\"\"\n", " tie: tie dropout mask across sequence lengths (Dropout1d/2d/3d)\n", " \"\"\"\n", " super().__init__()\n", " if p < 0 or p >= 1:\n", " raise ValueError(\"dropout probability has to be in [0, 1), \" \"but got {}\".format(p))\n", " self.p = p\n", " self.tie = tie\n", " self.transposed = transposed\n", " self.binomial = torch.distributions.binomial.Binomial(probs=1-self.p)\n", "\n", " def forward(self, X):\n", " \"\"\"X: (batch, dim, lengths...).\"\"\"\n", " if self.training:\n", " if not self.transposed: X = rearrange(X, 'b ... d -> b d ...')\n", " # binomial = torch.distributions.binomial.Binomial(probs=1-self.p) # This is incredibly slow because of CPU -> GPU copying\n", " mask_shape = X.shape[:2] + (1,)*(X.ndim-2) if self.tie else X.shape\n", " # mask = self.binomial.sample(mask_shape)\n", " mask = torch.rand(*mask_shape, device=X.device) < 1.-self.p\n", " X = X * mask * (1.0/(1-self.p))\n", " if not self.transposed: X = rearrange(X, 'b d ... -> b ... d')\n", " return X\n", " return X\n", "\n", "\n", "\n", "\"\"\"Minimal version of S4D with extra options and features stripped out, for pedagogical purposes.\"\"\"\n", "\n", "class S4DKernel(nn.Module):\n", " \"\"\"Generate convolution kernel from diagonal SSM parameters.\"\"\"\n", "\n", " def __init__(self, d_model, N=64, dt_min=0.001, dt_max=0.1, lr=None):\n", " super().__init__()\n", " # Generate dt\n", " H = d_model\n", " log_dt = torch.rand(H) * (\n", " math.log(dt_max) - math.log(dt_min)\n", " ) + math.log(dt_min)\n", "\n", " C = torch.randn(H, N // 2, dtype=torch.cfloat)\n", " self.C = nn.Parameter(torch.view_as_real(C))\n", " self.register(\"log_dt\", log_dt, lr)\n", "\n", " log_A_real = torch.log(0.5 * torch.ones(H, N//2))\n", " A_imag = math.pi * repeat(torch.arange(N//2), 'n -> h n', h=H)\n", " self.register(\"log_A_real\", log_A_real, lr)\n", " self.register(\"A_imag\", A_imag, lr)\n", " # self.register_buffer(\"A_imag\", nn.Parameter(A_imag))\n", " # self.register_parameter(\"log_A_real\", nn.Parameter(log_A_real))\n", "\n", "\n", " def forward(self, L):\n", " \"\"\"\n", " returns: (..., c, L) where c is number of channels (default 1)\n", " \"\"\"\n", "\n", " # Materialize parameters\n", " dt = torch.exp(self.log_dt) # (H)\n", " C = torch.view_as_complex(self.C) # (H N)\n", " A = -torch.exp(self.log_A_real) + 1j * self.A_imag # (H N)\n", "\n", " # Vandermonde multiplication\n", " dtA = A * dt.unsqueeze(-1) # (H N)\n", " K = dtA.unsqueeze(-1) * torch.arange(L, device=A.device) # (H N L)\n", " C = C * (torch.exp(dtA)-1.) / A\n", " K = 2 * torch.einsum('hn, hnl -> hl', C, torch.exp(K)).real\n", "\n", " return K\n", "\n", " def register(self, name, tensor, lr=None):\n", " \"\"\"Register a tensor with a configurable learning rate and 0 weight decay\"\"\"\n", "\n", " if lr == 0.0:\n", " self.register_buffer(name, tensor)\n", " else:\n", " self.register_parameter(name, nn.Parameter(tensor))\n", "\n", " optim = {\"weight_decay\": 0.0}\n", " if lr is not None: optim[\"lr\"] = lr\n", " setattr(getattr(self, name), \"_optim\", optim)\n", "\n", "\n", "class S4D(nn.Module):\n", " def __init__(self, d_model, d_state=64, dropout=0.0, transposed=True, **kernel_args):\n", " super().__init__()\n", "\n", " self.h = d_model\n", " self.n = d_state\n", " self.d_output = self.h\n", " self.transposed = transposed\n", "\n", " self.D = nn.Parameter(torch.randn(self.h))\n", "\n", " # SSM Kernel\n", " self.kernel = S4DKernel(self.h, N=self.n, **kernel_args)\n", "\n", " # Pointwise\n", " self.activation = nn.GELU()\n", " # dropout_fn = nn.Dropout2d # NOTE: bugged in PyTorch 1.11\n", " dropout_fn = DropoutNd\n", " self.dropout = dropout_fn(dropout) if dropout > 0.0 else nn.Identity()\n", "\n", " # position-wise output transform to mix features\n", " self.output_linear = nn.Sequential(\n", " nn.Conv1d(self.h, 2*self.h, kernel_size=1),\n", " nn.GLU(dim=-2),\n", " )\n", "\n", " def forward(self, u, **kwargs): # absorbs return_output and transformer src mask\n", " \"\"\" Input and output shape (B, H, L) \"\"\"\n", " if not self.transposed: u = u.transpose(-1, -2)\n", " L = u.size(-1)\n", "\n", " # Compute SSM Kernel\n", " k = self.kernel(L=L) # (H L)\n", "\n", " # Convolution\n", " k_f = torch.fft.rfft(k, n=2*L) # (H L)\n", " u_f = torch.fft.rfft(u, n=2*L) # (B H L)\n", " # print(u_f.shape, k_f.shape)\n", " y = torch.fft.irfft(u_f*k_f, n=2*L)[..., :L] # (B H L)\n", "\n", " # Compute D term in state space equation - essentially a skip connection\n", " y = y + u * self.D.unsqueeze(-1)\n", "\n", " y = self.dropout(self.activation(y))\n", " y = self.output_linear(y)\n", " if not self.transposed: y = y.transpose(-1, -2)\n", " return y, None # Return a dummy state to satisfy this repo's interface, but this can be modified\n", "\n", "\n", "class S4DTokenClassifier(nn.Module):\n", " \"\"\"S4D Token Classifier.\"\"\"\n", "\n", " @property\n", " def bias(self) -> bool: # noqa: D102\n", " return self._bias\n", "\n", " @property\n", " def d_model(self) -> int: # noqa: D102\n", " return self._d_model\n", "\n", " @property\n", " def d_state(self) -> int: # noqa: D102\n", " return self._d_state\n", "\n", " @property\n", " def n_vocab(self) -> int: # noqa: D102\n", " return self._n_vocab\n", "\n", " @property\n", " def n_output(self) -> int: # noqa: D102\n", " return self._n_output\n", "\n", " @property\n", " def dropout(self) -> float: # noqa: D102\n", " return self._dropout\n", "\n", " @property\n", " def transposed(self) -> bool: # noqa: D102\n", " return self._transposed\n", "\n", " @property\n", " def n_layers(self) -> int: # noqa: D102\n", " return self._n_layers\n", "\n", " @property\n", " def prenorm(self) -> bool: # noqa: D102\n", " return self._prenorm\n", "\n", " @property\n", " def num_parameters(self) -> int: # noqa: D102\n", " return sum(p.numel() for p in self.parameters() if p.requires_grad)\n", "\n", " def __init__( # noqa: D107\n", " self,\n", " d_model: int | None,\n", " d_state: int,\n", " dropout: float,\n", " n_layers: int,\n", " n_vocab: int,\n", " n_output: int,\n", " pad_token_index: int | None = None,\n", " lr: float = None,\n", " transposed: bool = True,\n", " prenorm: bool = False,\n", " bias: bool = True,\n", " embed: bool = False\n", " ):\n", " super().__init__()\n", " self._d_model = d_model\n", " self._d_state = d_state\n", " self._dropout = dropout\n", " self._n_vocab = n_vocab\n", " self._n_output = n_output\n", " self._n_layers = n_layers\n", " self._prenorm = prenorm\n", " self._transposed = transposed\n", " self._bias = bias\n", " self.embed = embed\n", "\n", " if pad_token_index == None:\n", " pad_token_index = n_vocab\n", "\n", " self._pad_token_index = pad_token_index\n", " if embed:\n", " if pad_token_index is not None:\n", " assert pad_token_index == n_vocab, \"We want pad token to be equal to the vocab size\"\n", " n_vocab += 1\n", " self.embedding = nn.Embedding(num_embeddings=n_vocab, embedding_dim=d_model, padding_idx=pad_token_index)\n", " print(\"real embedding is being applied\")\n", " else:\n", " if pad_token_index != n_vocab:\n", " raise ValueError(\n", " \"When using one-hot encoding, the padding token must always be a value equivalent to the vocabulary size.\"\n", " )\n", "\n", " self.embedding = partial(F.one_hot, num_classes=(n_vocab))\n", " print(\"one-hot is being applied\")\n", " self._d_model = pad_token_index\n", "\n", " self.s4d_layers = nn.ModuleList()\n", " self.norms = nn.ModuleList()\n", " self.dropout_layers = nn.ModuleList()\n", "\n", " for _ in range(n_layers):\n", " self.s4d_layers.append(\n", " S4D(d_model=d_model, d_state=d_state, dropout=dropout, transposed=transposed, lr=lr)\n", " )\n", " self.norms.append(nn.LayerNorm(d_model))\n", " self.dropout_layers.append(nn.Dropout(dropout))\n", "\n", " self.cl_head = nn.Linear(self.d_model, self.n_output, self.bias)\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor: # noqa: D102\n", " # if not self.embed:\n", " x = x.long().squeeze()\n", " x = self.embedding(x) # (B, L, d_input) -> (B, L, d_model)\n", " # print(f\"embedded x has shape {x.shape}\")\n", " x = x.transpose(-1, -2) # (B, L, d_model) -> (B, d_model, L)\n", "\n", " for layer, norm, dropout in zip(\n", " self.s4d_layers, self.norms, self.dropout_layers\n", " ):\n", " z = x\n", " if self.prenorm:\n", " z = norm(z.transpose(-1, -2)).transpose(-1, -2)\n", "\n", " z, _ = layer(z)\n", " z = dropout(z)\n", " x = x + z\n", "\n", " if not self.prenorm:\n", " x = norm(x.transpose(-1, -2)).transpose(-1, -2)\n", "\n", " x = x.transpose(-1, -2)\n", " x = self.cl_head(x)\n", " x = torch.tanh(x)\n", "\n", " return x, None" ], "metadata": { "id": "U8gWZPgfNRqW" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### Mixtures" ], "metadata": { "id": "qmzo_PoxNp2G" } }, { "cell_type": "code", "source": [ "# ----------------------------------------------------------------------\n", "# Imports\n", "# ----------------------------------------------------------------------\n", "print(\"Importing libraries...\")\n", "import itertools\n", "import math\n", "import random\n", "from typing import Tuple\n", "\n", "import torch\n", "from torch import Tensor\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import torch.optim as optim\n", "from torch.utils.data import DataLoader, Dataset\n", "\n", "# ---------------------------------------------------------------------\n", "# Continuous-time HiPPO-LegS operator\n", "# ---------------------------------------------------------------------\n", "def hippo_leg(N: int,\n", " *,\n", " device: torch.device | str = \"cpu\",\n", " dtype: torch.dtype = torch.float32,\n", ") -> Tuple[Tensor, Tensor]:\n", " \"\"\"\n", " Returns continuous-time (A, b) with shape\n", " A : (N, N) lower-triangular\n", " b : (N,) *single* input column\n", " \"\"\"\n", " i = torch.arange(N, device=device, dtype=dtype)\n", " j = torch.arange(N, device=device, dtype=dtype)\n", " ii, jj = torch.meshgrid(i, j, indexing=\"ij\")\n", "\n", " A = torch.where(\n", " ii > jj,\n", " (-1.0) ** (ii - jj + 1) * torch.sqrt(2 * ii + 1) * torch.sqrt(2 * jj + 1),\n", " torch.zeros_like(ii),\n", " )\n", " A -= torch.diag(i + 1)\n", " b = (-1.0) ** i * torch.sqrt(2 * i + 1)\n", " return A, b\n", "\n", "\n", "# ---------------------------------------------------------------------\n", "# Bilinear discretisation\n", "# ---------------------------------------------------------------------\n", "def bilinear_discretize(A: Tensor, B: Tensor, dt: float) -> Tuple[Tensor, Tensor]:\n", " \"\"\"\n", " Bilinear (Tustin) discretisation.\n", " A : (N, N)\n", " B : (N, K) (K = 1 for scalar input)\n", " Returns (Ad, Bd) with the same shapes.\n", " \"\"\"\n", " I = torch.eye(A.size(0), device=A.device, dtype=A.dtype)\n", " inv = torch.linalg.inv(I - 0.5 * dt * A)\n", " Ad = inv @ (I + 0.5 * dt * A)\n", " Bd = inv @ (dt * B)\n", " return Ad, Bd\n", "\n", "# ---------------------------------------------------------------------\n", "# Helper that expands B to (N, embed_dim)\n", "# ---------------------------------------------------------------------\n", "def make_hippo_discrete(\n", " N: int,\n", " embed_dim: int,\n", " dt: float = 1/256,\n", " *,\n", " device: torch.device | str = \"cpu\",\n", " dtype: torch.dtype = torch.float32,\n", ") -> Tuple[Tensor, Tensor]:\n", " \"\"\"\n", " Convenience wrapper:\n", " * builds continuous HiPPO (A, b)\n", " * broadcasts b --> (N, embed_dim)\n", " * discretises\n", " Returns:\n", " Ad : (N, N)\n", " Bd : (N, embed_dim)\n", " \"\"\"\n", " A, b = hippo_leg(N, device=device, dtype=dtype)\n", "\n", " # replicate the single canonical column across embed_dim inputs\n", " B = b.unsqueeze(1).repeat(1, embed_dim) # (N, embed_dim)\n", "\n", " Ad, Bd = bilinear_discretize(A, B, dt)\n", " return Ad, Bd\n", "\n", "\"\"\"\n", "HiPPO-S4D parity task\n", "\"\"\"\n", "class S4DCell(nn.Module):\n", " \"\"\"\n", " A single recurrent step using a discretised HiPPO-LegS operator.\n", "\n", " Parameters\n", " ----------\n", " input_size : int\n", " Dimension of the one-hot input vector xₜ.\n", " state_size : int\n", " Dimensionality of the hidden state hₜ.\n", " embedding_size : int\n", " Size of the learned input embedding φ(xₜ).\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " state_size: int,\n", " embedding_size: int = 128,\n", " ):\n", " super().__init__()\n", " self.state_size = state_size\n", "\n", " # Simple two-layer MLP to embed the one-hot input\n", " self.input_mlp = nn.Sequential(\n", " nn.Linear(embedding_size, embedding_size),\n", " nn.ReLU(),\n", " nn.Linear(embedding_size, embedding_size),\n", " nn.ReLU(),\n", " )\n", "\n", " # Non-trainable HiPPO backbone (discretised)\n", " A_disc, B_disc = make_hippo_discrete(state_size, embedding_size, dt=1.0 / 256)\n", " self.register_buffer(\"A_disc\", A_disc) # (state_size, state_size)\n", " self.register_buffer(\"B_disc\", B_disc) # (state_size, embedding_size)\n", "\n", " # ------------------------------------------------------------------\n", " # forward\n", " # ------------------------------------------------------------------\n", " def forward(self, x_emb: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor:\n", " \"\"\"\n", " One recurrent update.\n", "\n", " Parameters\n", " ----------\n", " x_emb : tensor, shape (batch, embedding_size)\n", " h_prev : tensor, shape (batch, state_size)\n", "\n", " Returns\n", " -------\n", " h_next : tensor, shape (batch, state_size)\n", " \"\"\"\n", " x_emb = self.input_mlp(x_emb)\n", " h_next = h_prev @ self.A_disc.T + x_emb @ self.B_disc.T\n", " return h_next\n", "\n", "\n", "class S4DModel(nn.Module):\n", " \"\"\"\n", " HiPPO-S4D network unrolled over a sequence.\n", "\n", " * The recurrent core is `S4DCell`.\n", " * At each step we concatenate the HiPPO hidden state hₜ and\n", " the input embedding φ(xₜ), then map to logits with an MLP.\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " state_size: int,\n", " *,\n", " embedding_size: int = 128,\n", " mlp_hidden_size: int = 256,\n", " ):\n", " super().__init__()\n", " self.state_size = state_size\n", "\n", " self.cell = S4DCell(\n", " state_size,\n", " embedding_size=embedding_size,\n", " )\n", "\n", " self.readout_mlp = nn.Sequential(\n", " nn.Linear(embedding_size + state_size, mlp_hidden_size),\n", " nn.ReLU(),\n", " nn.Linear(mlp_hidden_size, embedding_size),\n", " )\n", "\n", " # ------------------------------------------------------------------\n", " # forward\n", " # ------------------------------------------------------------------\n", " def forward(self, x_sequence: torch.Tensor) -> torch.Tensor:\n", " \"\"\"\n", " Parameters\n", " ----------\n", " x_sequence : tensor, shape (seq_len, batch, embedding_size)\n", "\n", " Returns\n", " -------\n", " log_probs : tensor, shape (seq_len, batch, output_size)\n", " \"\"\"\n", " seq_len, batch_size, _ = x_sequence.shape\n", " h = torch.zeros(batch_size, self.state_size, device=x_sequence.device)\n", "\n", " outputs = []\n", " for t in range(seq_len):\n", " x_t = x_sequence[t] # (batch, input_size)\n", " h = self.cell(x_t, h) # update hidden state\n", " x_emb = self.cell.input_mlp(x_t)\n", " logits = self.readout_mlp(torch.cat([x_emb, h], dim=-1))\n", " outputs.append(logits)\n", "\n", " return torch.stack(outputs) # (seq_len, batch, output_size)" ], "metadata": { "id": "SYsKFBdpNMmS", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "c1c93e74-cbb8-40a5-d431-a06284ee67b9" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Importing libraries...\n" ] } ] }, { "cell_type": "markdown", "source": [ "#### Mamba + S4D" ], "metadata": { "id": "reLT_O3ANrle" } }, { "cell_type": "code", "source": [ "# Copyright (c) 2023, Albert Gu, Tri Dao.\n", "\n", "import math\n", "from functools import partial\n", "import json\n", "import os\n", "import copy\n", "\n", "from collections import namedtuple\n", "\n", "import torch\n", "import torch.nn as nn\n", "\n", "from mamba_ssm.models.config_mamba import MambaConfig\n", "from mamba_ssm.modules.mamba_simple import Mamba\n", "from mamba_ssm.modules.mamba2 import Mamba2\n", "from mamba_ssm.modules.mha import MHA\n", "from mamba_ssm.modules.mlp import GatedMLP\n", "from mamba_ssm.modules.block import Block\n", "from mamba_ssm.utils.generation import GenerationMixin\n", "from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf\n", "from mamba_ssm.models.mixer_seq_simple import MixerModel, _init_weights, create_block\n", "\n", "try:\n", " from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn\n", "except ImportError:\n", " RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None\n", "\n", "\n", "class MambaS4DConfig(MambaConfig):\n", "\n", " def __init__(\n", " self,\n", " s4_state_size=256,\n", " s4_mlp_hidden_size=256,\n", " s4_dropout=0.0,\n", " prenorm=False,\n", " **kwargs\n", " ):\n", " super().__init__(**kwargs)\n", "\n", " self.s4_state_size = s4_state_size\n", " self.s4_dropout = s4_dropout\n", " self.s4_mlp_hidden_size = s4_mlp_hidden_size\n", " self.prenorm = prenorm\n", "\n", "class MambaS4LMHeadModel(nn.Module, GenerationMixin):\n", "\n", " def __init__(\n", " self,\n", " config: MambaS4DConfig,\n", " initializer_cfg=None,\n", " device=None,\n", " dtype=None,\n", " ) -> None:\n", " self.config = config\n", " d_model = config.d_model\n", " n_layer = 1 # config.n_layer\n", " d_intermediate = config.d_intermediate\n", " vocab_size = config.vocab_size\n", " ssm_cfg = config.ssm_cfg\n", " attn_layer_idx = config.attn_layer_idx\n", " attn_cfg = config.attn_cfg\n", " rms_norm = config.rms_norm\n", " residual_in_fp32 = config.residual_in_fp32\n", " fused_add_norm = config.fused_add_norm\n", " pad_vocab_size_multiple = config.pad_vocab_size_multiple\n", " factory_kwargs = {\"device\": device, \"dtype\": dtype}\n", "\n", " self.prenorm = config.prenorm\n", "\n", " super().__init__()\n", " if vocab_size % pad_vocab_size_multiple != 0:\n", " vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)\n", " self.backbone = MixerModel(\n", " d_model=d_model,\n", " n_layer=n_layer,\n", " d_intermediate=d_intermediate,\n", " vocab_size=vocab_size,\n", " ssm_cfg=ssm_cfg,\n", " attn_layer_idx=attn_layer_idx,\n", " attn_cfg=attn_cfg,\n", " rms_norm=rms_norm,\n", " initializer_cfg=initializer_cfg,\n", " fused_add_norm=fused_add_norm,\n", " residual_in_fp32=residual_in_fp32,\n", " **factory_kwargs,\n", " )\n", " self.s4 = S4DModel(\n", " embedding_size=d_model,\n", " state_size=config.s4_state_size,\n", " mlp_hidden_size=config.s4_mlp_hidden_size\n", " )\n", " self.s4_norm = nn.LayerNorm(d_model)\n", " self.s4_dropout = nn.Dropout(config.s4_dropout)\n", "\n", " self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)\n", "\n", " # Initialize weights and apply final processing\n", " self.apply(\n", " partial(\n", " _init_weights,\n", " n_layer=n_layer,\n", " **(initializer_cfg if initializer_cfg is not None else {}),\n", " )\n", " )\n", " self.tie_weights()\n", "\n", " def tie_weights(self):\n", " if self.config.tie_embeddings:\n", " self.lm_head.weight = self.backbone.embedding.weight\n", "\n", " def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n", " return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n", "\n", " def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, **mixer_kwargs):\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, **mixer_kwargs)\n", "\n", " z = hidden_states\n", "\n", " if self.prenorm:\n", " z = self.s4_norm(z)\n", "\n", " z = self.s4(z)\n", " z = self.s4_dropout(z)\n", " hidden_states = hidden_states + z\n", "\n", " if not self.prenorm:\n", " hidden_states = self.s4_norm(hidden_states)\n", "\n", " if num_last_tokens > 0:\n", " hidden_states = hidden_states[:, -num_last_tokens:]\n", "\n", " lm_logits = self.lm_head(hidden_states)\n", " CausalLMOutput = namedtuple(\"CausalLMOutput\", [\"logits\"])\n", " return CausalLMOutput(logits=lm_logits)\n", "\n", " @classmethod\n", " def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):\n", " config_data = load_config_hf(pretrained_model_name)\n", " config = MambaConfig(**config_data)\n", " model = cls(config, device=device, dtype=dtype, **kwargs)\n", " model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype))\n", " return model\n", "\n", " def save_pretrained(self, save_directory):\n", " \"\"\"\n", " Minimal implementation of save_pretrained for MambaLMHeadModel.\n", " Save the model and its configuration file to a directory.\n", " \"\"\"\n", " # Ensure save_directory exists\n", " os.makedirs(save_directory, exist_ok=True)\n", "\n", " # Save the model's state_dict\n", " model_path = os.path.join(save_directory, 'pytorch_model.bin')\n", " torch.save(self.state_dict(), model_path)\n", "\n", " # Save the configuration of the model\n", " config_path = os.path.join(save_directory, 'config.json')\n", " with open(config_path, 'w') as f:\n", " json.dump(self.config.__dict__, f, indent=4)" ], "metadata": { "id": "9GWFvjYFNlde" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "#### S4D + Mamba" ], "metadata": { "id": "gct_3RteNwdQ" } }, { "cell_type": "code", "source": [ "class S4MambaLMHeadModel(nn.Module, GenerationMixin):\n", "\n", " def __init__(\n", " self,\n", " config: MambaS4DConfig,\n", " initializer_cfg=None,\n", " device=None,\n", " dtype=None,\n", " ) -> None:\n", " self.config = config\n", " d_model = config.d_model\n", " n_layer = 1 # config.n_layer\n", " d_intermediate = config.d_intermediate\n", " vocab_size = config.vocab_size\n", " ssm_cfg = config.ssm_cfg\n", " attn_layer_idx = config.attn_layer_idx\n", " attn_cfg = config.attn_cfg\n", " rms_norm = config.rms_norm\n", " residual_in_fp32 = config.residual_in_fp32\n", " fused_add_norm = config.fused_add_norm\n", " pad_vocab_size_multiple = config.pad_vocab_size_multiple\n", " factory_kwargs = {\"device\": device, \"dtype\": dtype}\n", "\n", " self.prenorm = config.prenorm\n", "\n", " super().__init__()\n", " if vocab_size % pad_vocab_size_multiple != 0:\n", " vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)\n", "\n", " self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)\n", "\n", " self.backbone = S4DMambaMixerModel(\n", " d_model=d_model,\n", " n_layer=n_layer,\n", " d_intermediate=d_intermediate,\n", " vocab_size=vocab_size,\n", " ssm_cfg=ssm_cfg,\n", " attn_layer_idx=attn_layer_idx,\n", " attn_cfg=attn_cfg,\n", " rms_norm=rms_norm,\n", " initializer_cfg=initializer_cfg,\n", " fused_add_norm=fused_add_norm,\n", " residual_in_fp32=residual_in_fp32,\n", " **factory_kwargs,\n", " )\n", " self.s4 = S4DModel(\n", " embedding_size=d_model,\n", " state_size=config.s4_state_size,\n", " mlp_hidden_size=config.s4_mlp_hidden_size\n", " )\n", " self.s4_norm = nn.LayerNorm(d_model)\n", " self.s4_dropout = nn.Dropout(config.s4_dropout)\n", "\n", " self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)\n", "\n", " # Initialize weights and apply final processing\n", " self.apply(\n", " partial(\n", " _init_weights,\n", " n_layer=n_layer,\n", " **(initializer_cfg if initializer_cfg is not None else {}),\n", " )\n", " )\n", " self.tie_weights()\n", "\n", " def tie_weights(self):\n", " if self.config.tie_embeddings:\n", " self.lm_head.weight = self.embedding.weight\n", "\n", " def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n", " return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n", "\n", " def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, **mixer_kwargs):\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", "\n", " hidden_states = self.embedding(input_ids)\n", "\n", " z = hidden_states\n", "\n", " if self.prenorm:\n", " z = self.s4_norm(z)\n", "\n", " z = self.s4(z)\n", " z = self.s4_dropout(z)\n", " hidden_states = hidden_states + z\n", "\n", " if not self.prenorm:\n", " hidden_states = self.s4_norm(hidden_states)\n", "\n", " hidden_states = self.backbone(hidden_states, inference_params=inference_params, **mixer_kwargs)\n", "\n", " if num_last_tokens > 0:\n", " hidden_states = hidden_states[:, -num_last_tokens:]\n", "\n", " lm_logits = self.lm_head(hidden_states)\n", " CausalLMOutput = namedtuple(\"CausalLMOutput\", [\"logits\"])\n", " return CausalLMOutput(logits=lm_logits)\n", "\n", " @classmethod\n", " def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):\n", " config_data = load_config_hf(pretrained_model_name)\n", " config = MambaConfig(**config_data)\n", " model = cls(config, device=device, dtype=dtype, **kwargs)\n", " model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype))\n", " return model\n", "\n", " def save_pretrained(self, save_directory):\n", " \"\"\"\n", " Minimal implementation of save_pretrained for MambaLMHeadModel.\n", " Save the model and its configuration file to a directory.\n", " \"\"\"\n", " # Ensure save_directory exists\n", " os.makedirs(save_directory, exist_ok=True)\n", "\n", " # Save the model's state_dict\n", " model_path = os.path.join(save_directory, 'pytorch_model.bin')\n", " torch.save(self.state_dict(), model_path)\n", "\n", " # Save the configuration of the model\n", " config_path = os.path.join(save_directory, 'config.json')\n", " with open(config_path, 'w') as f:\n", " json.dump(self.config.__dict__, f, indent=4)" ], "metadata": { "id": "DUXo4SJhNxlX" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "## Training" ], "metadata": { "id": "cvioVu9wN1jD" } }, { "cell_type": "markdown", "source": [ "### RNN" ], "metadata": { "id": "ubl9s5uFN54W" } }, { "cell_type": "code", "source": [ "# for RNN (with pytorch)\n", "hidden_size = 8\n", "vocab_size = 2\n", "d_embed = 2\n", "output_size = n\n", "non_linearity = \"tanh\"\n", "embedding = False\n", "model = RNN(vocab_size=vocab_size, embedding_dim=d_embed, hidden_size=hidden_size, output_size=output_size, non_linearity=non_linearity, embedding=embedding)\n", "learning_rate = 1e-3\n", "optimizer = optim.Adam(model.parameters(), lr=learning_rate, fused=torch.cuda.is_available())\n", "loss_fn = torch.nn.CrossEntropyLoss()" ], "metadata": { "id": "AQx5lL_eN5my" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "is_second_order = False\n", "show_hs = False\n", "dim_output = output_size\n", "\n", "mappings = {\n", " 'loss': { 'line': 'train', 'facet': 'loss' },\n", " 'acc': { 'line': 'train', 'facet': 'acc' },\n", " 'hs_norm': { 'line': 'train', 'facet': 'hs_norm' },}\n", "if show_hs:\n", " mappings.update({ f\"hs0_{i}\": { 'line': 'train', 'facet': f\"hs0_{i}\" } for i in range(1,hidden_size+1) })\n", " mappings.update({ f\"hs1_{i}\": { 'line': 'train', 'facet': f\"hs1_{i}\" } for i in range(1,hidden_size+1) })\n", " # 'hs_1': { 'line': 'validation', 'facet': 'hs_1\n", " # 'hs_1': { 'line': 'validation', 'facet': 'hs_1' },\n", " # 'hs_2': { 'line': 'validation', 'facet': 'hs_2' },\n", "if is_second_order:\n", " mappings.update({\n", " 'flip_00': {'line': 'validation', 'facet': 'flip_00'},\n", " 'flip_01': {'line': 'validation', 'facet': 'flip_01'},\n", " 'flip_10': {'line': 'validation', 'facet': 'flip_10'},\n", " 'flip_11': {'line': 'validation', 'facet': 'flip_11'},\n", " 'bias_norm': {'line': 'validation', 'facet': 'bias_norm'},\n", " 'ident_00': {'line': 'validation', 'facet': 'ident_00'},\n", " 'ident_01': {'line': 'validation', 'facet': 'ident_01'},\n", " 'ident_10': {'line': 'validation', 'facet': 'ident_10'},\n", " 'ident_11': {'line': 'validation', 'facet': 'ident_11'},\n", " })\n", "\n", "facet_config = {\n", " 'loss': { 'name': 'Cross-Entropy', 'limit': [0, None], 'scale': 'linear' },\n", " 'acc': { 'name': 'Accuracy', 'limit': [0, 1], 'scale': 'linear' },\n", " 'hs_norm': { 'name': 'Hidden State Norm', 'limit': [0, 1000], 'scale': 'linear' },}\n", "if show_hs:\n", " facet_config.update({ f\"hs0_{i}\": { 'name': f'Hidden State 0 {i}', 'limit': [-100, 100], 'scale': 'linear' } for i in range(1,hidden_size+1) })\n", " facet_config.update({ f\"hs1_{i}\": { 'name': f'Hidden State 1 {i}', 'limit': [-100, 100], 'scale': 'linear' } for i in range(1,hidden_size+1) })\n", "if is_second_order:\n", " facet_config.update({\n", " # 'hs_1': { 'name': 'Hidden State 1', 'limit': [-1, 1], 'scale': 'linear' },\n", " # 'hs_2': { 'name': 'Hidden State 2', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'flip_00': { 'name': 'Flip 00', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'flip_01': { 'name': 'Flip 01', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'flip_10': { 'name': 'Flip 10', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'flip_11': { 'name': 'Flip 11', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'ident_00': { 'name': 'Ident 00', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'ident_01': { 'name': 'Ident 01', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'ident_10': { 'name': 'Ident 10', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'ident_11': { 'name': 'Ident 11', 'limit': [-1, 1], 'scale': 'linear' },\n", " 'bias_norm': { 'name': 'Bias Norm', 'limit': [0, 1], 'scale': 'linear' },\n", "\n", "})\n", "\n", "plot = PlotLearningCurve(\n", " mappings = mappings,\n", " facet_config = facet_config,\n", " xaxis_config = { 'name': 'Epoch', 'limit': [0, None] }\n", ")\n", "\n", "\n", "model = model.to(device)\n", "\n", "with plot:\n", " for batch_num, batch in enumerate(dataloader_mod_n):\n", " input_batch, output_batch = batch\n", " optimizer.zero_grad()\n", " x = input_batch.to(device)\n", " # print(f\"x has shape {x.shape}\")\n", " y = output_batch.to(device)\n", " # print(f\"y has shape {y.shape}\")\n", " output, states = model(x)\n", " # print(type(states[0]), states[0].shape, output_batch[0].shape)\n", " states_1_sample = states[0][torch.where(output_batch[0]==1)]\n", " states_0_sample = states[0][torch.where(output_batch[0]==0)]\n", " # states_1_sample = states_1[0]\n", " # states_0_sample = states_0[0]\n", " states_1_mean = states_1_sample.mean(dim=0)\n", " states_0_mean = states_0_sample.mean(dim=0)\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " # TODO: Check/run this to see if it's functional\n", " # BF16/FP16\n", " # with torch.cuda.amp.autocast(enabled=mixed_precision, dtype=torch.bfloat16):\n", " prediction = torch.argmax(output, dim=2)\n", " # print(prediction)\n", " # print(y)\n", " accuracy = torch.mean((prediction == y).float())\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", " if is_second_order:\n", " W = model.i2h.weight.detach().cpu().numpy()\n", " W_ident = W[0]\n", " W_flip = W[1]\n", " # bias_norm = model.i2h.bias.norm().item()\n", " # b = model.i2h.bias.detach().cpu().numpy()\n", "\n", " # grad_step += 1\n", " batch_loss.backward()\n", " optimizer.step()\n", " if batch_num % 100 == 0:\n", " # print(batch_loss)\n", " # print(f\"accuracy = {accuracy}\")\n", " plot.append(batch_num, {\n", " 'loss': batch_loss,\n", " 'acc': accuracy,\n", " 'hs_norm': states.norm().item()\n", " # 'flip_00': W_flip[0, 0],\n", " # 'flip_01': W_flip[0, 1],\n", " # 'flip_10': W_flip[1, 0],\n", " # 'flip_11': W_flip[1, 1],\n", " # 'ident_00': W_ident[0, 0],\n", " # 'ident_01': W_ident[0, 1],\n", " # 'ident_10': W_ident[1, 0],\n", " # 'ident_11': W_ident[1, 1],\n", " # 'bias_norm': bias_norm,\n", " # 'hs_1': states[0],\n", " # 'hs_2': states[1]\n", " })\n", " if show_hs:\n", " for i in range(hidden_size):\n", " plot.append(batch_num, { f\"hs0_{i+1}\": states_0_mean[i] })\n", " plot.append(batch_num, { f\"hs1_{i+1}\": states_1_mean[i] })\n", "\n", " if is_second_order:\n", " plot.append(batch_num, {\n", " 'flip_00': W_flip[0, 0],\n", " 'flip_01': W_flip[0, 1],\n", " 'flip_10': W_flip[1, 0],\n", " 'flip_11': W_flip[1, 1],\n", " 'ident_00': W_ident[0, 0],\n", " 'ident_01': W_ident[0, 1],\n", " 'ident_10': W_ident[1, 0],\n", " 'ident_11': W_ident[1, 1],\n", " # 'bias_norm': bias_norm,\n", " 'hs_1': states[0],\n", " 'hs_2': states[1]\n", " })\n", " plot.draw()" ], "metadata": { "id": "lnIeiVH5N1GJ", "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "outputId": "42c9e600-929f-49cc-9665-ffe5c751aa64" }, "execution_count": null, "outputs": [ { "output_type": "display_data", "data": { "text/plain": [ "" ], "text/html": [ "" ] }, "metadata": {} }, { "output_type": "display_data", "data": { "text/plain": [ "" ], "text/html": [ "" ] }, "metadata": {} }, { "output_type": "error", "ename": "KeyboardInterrupt", "evalue": "", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mKeyboardInterrupt\u001b[0m Traceback (most recent call last)", "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 59\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 60\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mplot\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 61\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mbatch_num\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataloader_mod_n\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 62\u001b[0m \u001b[0minput_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_batch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 63\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\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.11/dist-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 706\u001b[0m \u001b[0;31m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 707\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# type: ignore[call-arg]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 708\u001b[0;31m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_data\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 709\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_yielded\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 710\u001b[0m if (\n", "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 762\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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 763\u001b[0m \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 764\u001b[0;31m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 765\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 766\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory_device\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.11/dist-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36mfetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m 50\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__getitems__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 51\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 52\u001b[0;31m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpossibly_batched_index\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 53\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 54\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpossibly_batched_index\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.11/dist-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 50\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__getitems__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 51\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 52\u001b[0;31m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpossibly_batched_index\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 53\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 54\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpossibly_batched_index\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;36m__getitem__\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 16\u001b[0m \u001b[0;31m# Convert to torch tensors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m \u001b[0mrandom_list_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrandom_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mint64\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 18\u001b[0m \u001b[0mmod_n_list_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmod_n_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mint64\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;31mKeyboardInterrupt\u001b[0m: " ] } ] }, { "cell_type": "code", "source": [ "for batch_num, batch in enumerate(test_dataloader_n):\n", " input_batch, output_batch = batch\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output, states = model(x)\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " # TODO: Check/run this to see if it's functional\n", " # BF16/FP16\n", " # with torch.cuda.amp.autocast(enabled=mixed_precision, dtype=torch.bfloat16):\n", " prediction = torch.argmax(output, dim=2)\n", " # print(prediction)\n", " # print(y)\n", " accuracy = torch.mean((prediction == y).float())\n", " print(accuracy)\n", " # break" ], "metadata": { "id": "CrjETPAlOBXj", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "98053363-8dd6-4f47-a5b5-0bb69c5ff7dc" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "tensor(0.9945, device='cuda:0')\n", "tensor(0.9921, device='cuda:0')\n" ] } ] }, { "cell_type": "markdown", "source": [ "### S4D" ], "metadata": { "id": "LeMGFDBcODzi" } }, { "cell_type": "code", "source": [ "# for S4D\n", "d_model = 8\n", "d_state = 16\n", "n_layers = 1\n", "output_size = n\n", "\n", "model = S4DTokenClassifier(\n", " d_model=d_model,\n", " d_state=d_state,\n", " dropout=0,\n", " n_layers=n_layers,\n", " n_vocab=8,\n", " n_output=output_size\n", ")\n", "learning_rate = 5e-4\n", "optimizer = optim.AdamW(model.parameters(), lr=learning_rate, fused=torch.cuda.is_available())\n", "loss_fn = torch.nn.CrossEntropyLoss()" ], "metadata": { "id": "9oCWzelXOFOL", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "078289e9-b52c-4643-fafc-97eb3654617e" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "one-hot is being applied\n" ] } ] }, { "cell_type": "code", "source": [ "dim_output = 8\n", "\n", "mappings = {\n", " 'loss': { 'line': 'train', 'facet': 'loss' },\n", " 'acc': { 'line': 'train', 'facet': 'acc' },\n", "}\n", "\n", "facet_config = {\n", " 'loss': { 'name': 'Cross-Entropy', 'limit': [0, None], 'scale': 'linear' },\n", " 'acc': { 'name': 'Accuracy', 'limit': [0, 1], 'scale': 'linear' },\n", "}\n", "\n", "plot = PlotLearningCurve(\n", " mappings = mappings,\n", " facet_config = facet_config,\n", " xaxis_config = { 'name': 'Epoch', 'limit': [0, None] }\n", ")\n", "\n", "model = model.to(device)\n", "\n", "with plot:\n", " for batch_num, batch in enumerate(dataloader_mod_n):\n", " input_batch, output_batch = batch\n", " optimizer.zero_grad()\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " batch_loss.backward()\n", " optimizer.step()\n", " if batch_num % 1 == 0:\n", " plot.append(batch_num, {\n", " 'loss': batch_loss,\n", " 'acc': accuracy,\n", " })\n", " plot.draw()" ], "metadata": { "id": "9Rc-ytIgOHre", "colab": { "base_uri": "https://localhost:8080/", "height": 511 }, "outputId": "186022f1-0d74-4ef7-a68a-98deaab66335" }, "execution_count": null, "outputs": [ { "output_type": "display_data", "data": { "text/plain": [ "" ], "text/html": [ "" ] }, "metadata": {} }, { "output_type": "display_data", "data": { "text/plain": [ "" ], "application/javascript": [ "window.setupLearningCurve(\"974202fa-076f-429e-b0e2-efb499e81174\", {\"id\": \"974202fa-076f-429e-b0e2-efb499e81174\", \"width\": 600, \"height\": 490, \"mappings\": {\"loss\": {\"line\": \"train\", \"facet\": \"loss\"}, \"acc\": {\"line\": \"train\", \"facet\": \"acc\"}}, \"lineConfig\": {\"train\": {\"name\": \"Train\", \"color\": \"#F8766D\"}, \"validation\": {\"name\": \"Validation\", \"color\": \"#00BFC4\"}}, \"facetConfig\": {\"loss\": {\"name\": \"Cross-Entropy\", \"limit\": [0, null], \"scale\": \"linear\"}, \"acc\": {\"name\": \"Accuracy\", \"limit\": [0, 1], \"scale\": \"linear\"}}, \"xAxisConfig\": {\"name\": \"Epoch\", \"limit\": [0, null]}});" ] }, "metadata": {} } ] }, { "cell_type": "markdown", "source": [ "### Mamba" ], "metadata": { "id": "6vPC9-XcOJU4" } }, { "cell_type": "code", "source": [ "from mamba_ssm.models.config_mamba import MambaConfig\n", "from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\n", "\n", "# for S4D\n", "mamba_config = {\n", " \"d_model\": 8,\n", " \"n_layer\": 1,\n", " \"vocab_size\": 2,\n", " \"ssm_cfg\": {},\n", " \"rms_norm\": True,\n", " \"residual_in_fp32\": True,\n", " \"fused_add_norm\": True,\n", " \"pad_vocab_size_multiple\": 8\n", "}\n", "mamba_config = MambaConfig(**mamba_config)\n", "\n", "mamba_model = MambaLMHeadModel(mamba_config)\n", "learning_rate = 5e-4\n", "optimizer = optim.AdamW(mamba_model.parameters(), lr=learning_rate, fused=torch.cuda.is_available())\n", "loss_fn = torch.nn.CrossEntropyLoss()" ], "metadata": { "id": "FD_0P15OOKF5" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "dim_output = 8\n", "\n", "mappings = {\n", " 'loss': { 'line': 'train', 'facet': 'loss' },\n", " 'acc': { 'line': 'train', 'facet': 'acc' },\n", "}\n", "\n", "facet_config = {\n", " 'loss': { 'name': 'Cross-Entropy', 'limit': [0, None], 'scale': 'linear' },\n", " 'acc': { 'name': 'Accuracy', 'limit': [0, 1], 'scale': 'linear' },\n", "}\n", "\n", "plot = PlotLearningCurve(\n", " mappings = mappings,\n", " facet_config = facet_config,\n", " xaxis_config = { 'name': 'Epoch', 'limit': [0, None] }\n", ")\n", "\n", "mamba_model = mamba_model.to(device)\n", "\n", "with plot:\n", " for batch_num, batch in enumerate(dataloader_mod_n):\n", " input_batch, output_batch = batch\n", " optimizer.zero_grad()\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = mamba_model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " batch_loss.backward()\n", " optimizer.step()\n", " if batch_num % 1 == 0:\n", " plot.append(batch_num, {\n", " 'loss': batch_loss,\n", " 'acc': accuracy,\n", " })\n", " plot.draw()\n" ], "metadata": { "id": "HAFsdPmTOMDy" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "mamba_model.eval()\n", "with torch.no_grad():\n", " for batch_num, batch in enumerate(test_dataloader_n):\n", " input_batch, output_batch = batch\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = mamba_model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " print(accuracy)" ], "metadata": { "id": "vS8Ox4YdOPgQ" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### Mamba + S4D" ], "metadata": { "id": "ciaK0OssOQ9x" } }, { "cell_type": "code", "source": [ "mamba_config = {\n", " \"d_model\": 8,\n", " \"n_layer\": 1,\n", " \"vocab_size\": 2,\n", " \"ssm_cfg\": {},\n", " \"rms_norm\": True,\n", " \"residual_in_fp32\": True,\n", " \"fused_add_norm\": True,\n", " \"pad_vocab_size_multiple\": 8,\n", " \"s4_state_size\": 16,\n", " \"s4_dropout\": 0.0,\n", " \"s4_mlp_hidden_size\": 256,\n", " \"prenorm\": False\n", "}\n", "mamba_config = MambaS4Config(**mamba_config)\n", "\n", "mamba_model = MambaS4LMHeadModel(mamba_config)\n", "learning_rate = 5e-4\n", "optimizer = optim.AdamW(mamba_model.parameters(), lr=learning_rate, fused=torch.cuda.is_available())\n", "loss_fn = torch.nn.CrossEntropyLoss()" ], "metadata": { "id": "LZozpBKiOSFD" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "dim_output = 8\n", "\n", "mappings = {\n", " 'loss': { 'line': 'train', 'facet': 'loss' },\n", " 'acc': { 'line': 'train', 'facet': 'acc' },\n", "}\n", "\n", "facet_config = {\n", " 'loss': { 'name': 'Cross-Entropy', 'limit': [0, None], 'scale': 'linear' },\n", " 'acc': { 'name': 'Accuracy', 'limit': [0, 1], 'scale': 'linear' },\n", "}\n", "\n", "plot = PlotLearningCurve(\n", " mappings = mappings,\n", " facet_config = facet_config,\n", " xaxis_config = { 'name': 'Epoch', 'limit': [0, None] }\n", ")\n", "\n", "mamba_model = mamba_model.to(device)\n", "\n", "with plot:\n", " for batch_num, batch in enumerate(dataloader_mod_n):\n", " input_batch, output_batch = batch\n", " optimizer.zero_grad()\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = mamba_model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " batch_loss.backward()\n", " optimizer.step()\n", " if batch_num % 1 == 0:\n", " plot.append(batch_num, {\n", " 'loss': batch_loss,\n", " 'acc': accuracy,\n", " })\n", " plot.draw()\n" ], "metadata": { "id": "orQ-sEwKOUUb" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "mamba_model.eval()\n", "with torch.no_grad():\n", " for batch_num, batch in enumerate(test_dataloader_n):\n", " input_batch, output_batch = batch\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = mamba_model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " print(accuracy)" ], "metadata": { "id": "0WXTdO0COXil" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### S4D + Mamba" ], "metadata": { "id": "y1MBg-Q3OYLu" } }, { "cell_type": "code", "source": [ "mamba_model = S4MambaLMHeadModel(mamba_config)\n", "learning_rate = 5e-4\n", "optimizer = optim.AdamW(mamba_model.parameters(), lr=learning_rate, fused=torch.cuda.is_available())\n", "loss_fn = torch.nn.CrossEntropyLoss()" ], "metadata": { "id": "QqwDNNDsOZz9" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "dim_output = 8\n", "\n", "mappings = {\n", " 'loss': { 'line': 'train', 'facet': 'loss' },\n", " 'acc': { 'line': 'train', 'facet': 'acc' },\n", "}\n", "\n", "facet_config = {\n", " 'loss': { 'name': 'Cross-Entropy', 'limit': [0, None], 'scale': 'linear' },\n", " 'acc': { 'name': 'Accuracy', 'limit': [0, 1], 'scale': 'linear' },\n", "}\n", "\n", "plot = PlotLearningCurve(\n", " mappings = mappings,\n", " facet_config = facet_config,\n", " xaxis_config = { 'name': 'Epoch', 'limit': [0, None] }\n", ")\n", "\n", "mamba_model = mamba_model.to(device)\n", "\n", "with plot:\n", " for batch_num, batch in enumerate(dataloader_mod_n):\n", " input_batch, output_batch = batch\n", " optimizer.zero_grad()\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = mamba_model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", "\n", " batch_loss.backward()\n", " optimizer.step()\n", " if batch_num % 1 == 0:\n", " plot.append(batch_num, {\n", " 'loss': batch_loss,\n", " 'acc': accuracy,\n", " })\n", " plot.draw()\n" ], "metadata": { "id": "PHIZMlwGOeWE" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "mamba_model.eval()\n", "with torch.no_grad():\n", " for batch_num, batch in enumerate(test_dataloader_n):\n", " input_batch, output_batch = batch\n", " x = input_batch.to(device)\n", " y = output_batch.to(device)\n", " output = mamba_model(x).logits\n", "\n", " batch_loss = loss_fn(output.view(-1, dim_output), y.view(-1).long())\n", " prediction = torch.argmax(output, dim=2)\n", " accuracy = torch.mean((prediction == y).float())\n", " print(accuracy)" ], "metadata": { "id": "_tL5cI_7Oh55" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [], "metadata": { "id": "KpMpo-HgwdWH" }, "execution_count": null, "outputs": [] } ] }