{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "mount_file_id": "1OMmaLGK4ipNMEYHIs8aaZIXQW8SRHMaP",
      "authorship_tag": "ABX9TyNrlZW33/8/lSiD6XJKn0Cl",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "aae559be2b744c808b851a257c510d3f": {
          "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_14d86eba634548a795aaeaa16de65c2c",
              "IPY_MODEL_b6e1ff089de04b8fa88493e54619f583",
              "IPY_MODEL_251bec6e9d324f248193cb34f81294ee"
            ],
            "layout": "IPY_MODEL_9216f5d72acb413ba0c392eb0cb650d3"
          }
        },
        "14d86eba634548a795aaeaa16de65c2c": {
          "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_47763d25fdaf4f54ad1f194c81418462",
            "placeholder": "​",
            "style": "IPY_MODEL_662acc329ab046fbaf2f1b587480a964",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "b6e1ff089de04b8fa88493e54619f583": {
          "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_d754888b81ef4640955f33ceb009f4bf",
            "max": 21,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a682a90dad094bb18b78f2ceeab9a17d",
            "value": 21
          }
        },
        "251bec6e9d324f248193cb34f81294ee": {
          "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_694167de911446e2b31f8236962d72d1",
            "placeholder": "​",
            "style": "IPY_MODEL_d6ab5150f8bb4648a27ee458a50b179d",
            "value": " 21/21 [07:11&lt;00:00, 21.18s/it]"
          }
        },
        "9216f5d72acb413ba0c392eb0cb650d3": {
          "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
          }
        },
        "47763d25fdaf4f54ad1f194c81418462": {
          "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
          }
        },
        "662acc329ab046fbaf2f1b587480a964": {
          "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": ""
          }
        },
        "d754888b81ef4640955f33ceb009f4bf": {
          "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
          }
        },
        "a682a90dad094bb18b78f2ceeab9a17d": {
          "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": ""
          }
        },
        "694167de911446e2b31f8236962d72d1": {
          "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
          }
        },
        "d6ab5150f8bb4648a27ee458a50b179d": {
          "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": ""
          }
        },
        "132ebb644b3b4acaa0244d5d46675a2d": {
          "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_debb6035ad8b446793f540222d3ff2b0",
              "IPY_MODEL_7c3448293fc94c5fadf6b6b8a11af5c2",
              "IPY_MODEL_1632d7a3bb194108bd94ec20175a1bee"
            ],
            "layout": "IPY_MODEL_cee6859fe4c64275863619d909a41816"
          }
        },
        "debb6035ad8b446793f540222d3ff2b0": {
          "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_d0ac463b4c5f48649c2b89132fa8cc86",
            "placeholder": "​",
            "style": "IPY_MODEL_720af8fa27364d549dc0bc3f7e4bc4d2",
            "value": "Map: 100%"
          }
        },
        "7c3448293fc94c5fadf6b6b8a11af5c2": {
          "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_b99cff08bdfb445a97b1abba86d3075a",
            "max": 2508,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d783494565f04ddca8ad6d4889c2efaf",
            "value": 2508
          }
        },
        "1632d7a3bb194108bd94ec20175a1bee": {
          "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_16e861b1c3b54cbeb537834e5e48c9bc",
            "placeholder": "​",
            "style": "IPY_MODEL_616b86cd560544e081cbebea8bb97d99",
            "value": " 2508/2508 [00:00&lt;00:00, 9638.51 examples/s]"
          }
        },
        "cee6859fe4c64275863619d909a41816": {
          "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
          }
        },
        "d0ac463b4c5f48649c2b89132fa8cc86": {
          "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
          }
        },
        "720af8fa27364d549dc0bc3f7e4bc4d2": {
          "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": ""
          }
        },
        "b99cff08bdfb445a97b1abba86d3075a": {
          "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
          }
        },
        "d783494565f04ddca8ad6d4889c2efaf": {
          "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": ""
          }
        },
        "16e861b1c3b54cbeb537834e5e48c9bc": {
          "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
          }
        },
        "616b86cd560544e081cbebea8bb97d99": {
          "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",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mrm8488/shared_colab_notebooks/blob/master/fine_tune_Jamba.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fine-tune Jamba-v0.1 on A100 - 40GB VRAM using QLoRA"
      ],
      "metadata": {
        "id": "sOfeXQXBtLVO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! nvidia-smi"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RCJ4Y-3LpCOW",
        "outputId": "d2eaea7f-130a-4734-8bec-8bf5a29ded83"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Fri Mar 29 01:26:40 2024       \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |\n",
            "|-----------------------------------------+----------------------+----------------------+\n",
            "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                                         |                      |               MIG M. |\n",
            "|=========================================+======================+======================|\n",
            "|   0  NVIDIA A100-SXM4-40GB          Off | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   33C    P0              47W / 400W |      2MiB / 40960MiB |      0%      Default |\n",
            "|                                         |                      |             Disabled |\n",
            "+-----------------------------------------+----------------------+----------------------+\n",
            "                                                                                         \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| Processes:                                                                            |\n",
            "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
            "|        ID   ID                                                             Usage      |\n",
            "|=======================================================================================|\n",
            "|  No running processes found                                                           |\n",
            "+---------------------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install Flash Attention 2"
      ],
      "metadata": {
        "id": "Dd10VLMRtWxF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install ninja packaging\n",
        "! pip install flash-attn --no-build-isolation"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "panZSjSVq6Cc",
        "outputId": "20e445ac-77ae-4eb5-d73c-db1f2adbe280"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: ninja in /usr/local/lib/python3.10/dist-packages (1.11.1.1)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (24.0)\n",
            "Requirement already satisfied: flash-attn in /usr/local/lib/python3.10/dist-packages (2.5.6)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from flash-attn) (2.2.1+cu121)\n",
            "Requirement already satisfied: einops in /usr/local/lib/python3.10/dist-packages (from flash-attn) (0.7.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from flash-attn) (24.0)\n",
            "Requirement already satisfied: ninja in /usr/local/lib/python3.10/dist-packages (from flash-attn) (1.11.1.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (3.13.3)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (4.10.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (2023.6.0)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (8.9.2.26)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (12.1.0.106)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (2.19.3)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (12.1.105)\n",
            "Requirement already satisfied: triton==2.2.0 in /usr/local/lib/python3.10/dist-packages (from torch->flash-attn) (2.2.0)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch->flash-attn) (12.4.99)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->flash-attn) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->flash-attn) (1.3.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install Required Dependencies"
      ],
      "metadata": {
        "id": "hniJWDLCta-l"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ec-dRAOVo8TO",
        "outputId": "53bff736-ca5b-47c4-dc4d-1e4a39b23dbf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: transformers>=4.39.0 in /usr/local/lib/python3.10/dist-packages (4.39.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (3.13.3)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.19.3 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (0.20.3)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (1.25.2)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (24.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (2023.12.25)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (0.15.2)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (0.4.2)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.39.0) (4.66.2)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.19.3->transformers>=4.39.0) (2023.6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.19.3->transformers>=4.39.0) (4.10.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers>=4.39.0) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers>=4.39.0) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers>=4.39.0) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers>=4.39.0) (2024.2.2)\n",
            "Requirement already satisfied: mamba-ssm in /usr/local/lib/python3.10/dist-packages (1.2.0.post1)\n",
            "Requirement already satisfied: causal-conv1d>=1.2.0 in /usr/local/lib/python3.10/dist-packages (1.2.0.post2)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (2.2.1+cu121)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (24.0)\n",
            "Requirement already satisfied: ninja in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (1.11.1.1)\n",
            "Requirement already satisfied: einops in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (0.7.0)\n",
            "Requirement already satisfied: triton in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (2.2.0)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (from mamba-ssm) (4.39.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (3.13.3)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (4.10.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (2023.6.0)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (8.9.2.26)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (12.1.0.106)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (2.19.3)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->mamba-ssm) (12.1.105)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch->mamba-ssm) (12.4.99)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.19.3 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (0.20.3)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (1.25.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (2023.12.25)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (0.15.2)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (0.4.2)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers->mamba-ssm) (4.66.2)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->mamba-ssm) (2.1.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers->mamba-ssm) (2024.2.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->mamba-ssm) (1.3.0)\n",
            "Requirement already satisfied: peft in /usr/local/lib/python3.10/dist-packages (0.10.0)\n",
            "Requirement already satisfied: trl in /usr/local/lib/python3.10/dist-packages (0.8.1)\n",
            "Requirement already satisfied: bitsandbytes in /usr/local/lib/python3.10/dist-packages (0.43.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from peft) (1.25.2)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from peft) (24.0)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from peft) (5.9.5)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from peft) (6.0.1)\n",
            "Requirement already satisfied: torch>=1.13.0 in /usr/local/lib/python3.10/dist-packages (from peft) (2.2.1+cu121)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (from peft) (4.39.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from peft) (4.66.2)\n",
            "Requirement already satisfied: accelerate>=0.21.0 in /usr/local/lib/python3.10/dist-packages (from peft) (0.28.0)\n",
            "Requirement already satisfied: safetensors in /usr/local/lib/python3.10/dist-packages (from peft) (0.4.2)\n",
            "Requirement already satisfied: huggingface-hub>=0.17.0 in /usr/local/lib/python3.10/dist-packages (from peft) (0.20.3)\n",
            "Requirement already satisfied: datasets in /usr/local/lib/python3.10/dist-packages (from trl) (2.18.0)\n",
            "Requirement already satisfied: tyro>=0.5.11 in /usr/local/lib/python3.10/dist-packages (from trl) (0.7.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.17.0->peft) (3.13.3)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.17.0->peft) (2023.6.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.17.0->peft) (2.31.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.17.0->peft) (4.10.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (3.1.3)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (8.9.2.26)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (12.1.0.106)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (2.19.3)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (12.1.105)\n",
            "Requirement already satisfied: triton==2.2.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.0->peft) (2.2.0)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.13.0->peft) (12.4.99)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers->peft) (2023.12.25)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers->peft) (0.15.2)\n",
            "Requirement already satisfied: docstring-parser>=0.14.1 in /usr/local/lib/python3.10/dist-packages (from tyro>=0.5.11->trl) (0.16)\n",
            "Requirement already satisfied: rich>=11.1.0 in /usr/local/lib/python3.10/dist-packages (from tyro>=0.5.11->trl) (13.7.1)\n",
            "Requirement already satisfied: shtab>=1.5.6 in /usr/local/lib/python3.10/dist-packages (from tyro>=0.5.11->trl) (1.7.1)\n",
            "Requirement already satisfied: pyarrow>=12.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (14.0.2)\n",
            "Requirement already satisfied: pyarrow-hotfix in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (0.6)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (0.3.8)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (1.5.3)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (3.4.1)\n",
            "Requirement already satisfied: multiprocess in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (0.70.16)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets->trl) (3.9.3)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->trl) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->trl) (23.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->trl) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->trl) (6.0.5)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->trl) (1.9.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->trl) (4.0.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.17.0->peft) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.17.0->peft) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.17.0->peft) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.17.0->peft) (2024.2.2)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich>=11.1.0->tyro>=0.5.11->trl) (3.0.0)\n",
            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich>=11.1.0->tyro>=0.5.11->trl) (2.16.1)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.13.0->peft) (2.1.5)\n",
            "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets->trl) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets->trl) (2023.4)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.13.0->peft) (1.3.0)\n",
            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=11.1.0->tyro>=0.5.11->trl) (0.1.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.1->pandas->datasets->trl) (1.16.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install -U \"transformers>=4.39.0\"\n",
        "!pip install mamba-ssm \"causal-conv1d>=1.2.0\"\n",
        "!pip install peft trl bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_id = \"ai21labs/Jamba-v0.1\""
      ],
      "metadata": {
        "id": "8tJgLuhTqLCZ"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "from trl import SFTTrainer\n",
        "from peft import LoraConfig\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig"
      ],
      "metadata": {
        "id": "zhmfYJMopcB3"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    llm_int4_skip_modules=[\"mamba\"]\n",
        ")"
      ],
      "metadata": {
        "id": "iVgCYTKcp3Xr"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    trust_remote_code=True,\n",
        "    device_map='auto',\n",
        "    attn_implementation=\"flash_attention_2\",\n",
        "    quantization_config=quantization_config,\n",
        "    use_mamba_kernels=False #Disabling the mamba kernels since I have a recurrent error.\n",
        "    )"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "aae559be2b744c808b851a257c510d3f",
            "14d86eba634548a795aaeaa16de65c2c",
            "b6e1ff089de04b8fa88493e54619f583",
            "251bec6e9d324f248193cb34f81294ee",
            "9216f5d72acb413ba0c392eb0cb650d3",
            "47763d25fdaf4f54ad1f194c81418462",
            "662acc329ab046fbaf2f1b587480a964",
            "d754888b81ef4640955f33ceb009f4bf",
            "a682a90dad094bb18b78f2ceeab9a17d",
            "694167de911446e2b31f8236962d72d1",
            "d6ab5150f8bb4648a27ee458a50b179d"
          ]
        },
        "id": "ktiiRqlgp8bI",
        "outputId": "2d36509f-1b12-4458-f0f4-cf68c86a658b"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/21 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "aae559be2b744c808b851a257c510d3f"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#model.save_pretrained(\"/content/drive/MyDrive/jamba\")\n",
        "#tokenizer.save_pretrained(\"/content/drive/MyDrive/jamba\")"
      ],
      "metadata": {
        "id": "WNZbOoUKrCBt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = load_dataset(\"Abirate/english_quotes\", split=\"train\")"
      ],
      "metadata": {
        "id": "I2TWny5tqU8I"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./results\",\n",
        "    num_train_epochs=3,\n",
        "    per_device_train_batch_size=1,\n",
        "    logging_dir='./logs',\n",
        "    logging_steps=10,\n",
        "    learning_rate=2e-3 # 2.5e-5\n",
        ")"
      ],
      "metadata": {
        "id": "ee1IMXszqWyU"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "lora_config = LoraConfig(\n",
        "    r=8,\n",
        "    target_modules=[\"embed_tokens\", \"x_proj\", \"in_proj\", \"out_proj\"],\n",
        "    task_type=\"CAUSAL_LM\",\n",
        "    bias=\"none\"\n",
        ")"
      ],
      "metadata": {
        "id": "iyipI7FyqZHI"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    args=training_args,\n",
        "    peft_config=lora_config,\n",
        "    train_dataset=dataset,\n",
        "    dataset_text_field=\"quote\",\n",
        "    max_seq_length=128\n",
        ")"
      ],
      "metadata": {
        "id": "6gMx9J6Eqas3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 158,
          "referenced_widgets": [
            "132ebb644b3b4acaa0244d5d46675a2d",
            "debb6035ad8b446793f540222d3ff2b0",
            "7c3448293fc94c5fadf6b6b8a11af5c2",
            "1632d7a3bb194108bd94ec20175a1bee",
            "cee6859fe4c64275863619d909a41816",
            "d0ac463b4c5f48649c2b89132fa8cc86",
            "720af8fa27364d549dc0bc3f7e4bc4d2",
            "b99cff08bdfb445a97b1abba86d3075a",
            "d783494565f04ddca8ad6d4889c2efaf",
            "16e861b1c3b54cbeb537834e5e48c9bc",
            "616b86cd560544e081cbebea8bb97d99"
          ]
        },
        "outputId": "d8fa0622-0fe8-42e5-d6b7-e7487b4646ef"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/2508 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "132ebb644b3b4acaa0244d5d46675a2d"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/trl/trainer/sft_trainer.py:317: UserWarning: You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py:432: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches', 'even_batches', 'use_seedable_sampler']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
            "dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False, even_batches=True, use_seedable_sampler=True)\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir ./logs"
      ],
      "metadata": {
        "id": "YprXnC4fyXmM",
        "outputId": "f2348c1f-379a-4756-8fe5-92950a6a3815",
        "colab": {
          "resources": {
            "https://localhost:6006/?tensorboardColab=true": {
              "data": "[encoded data, 711032 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "text/html; charset=utf-8"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/font-roboto/oMMgfZMQthOryQo9n22dcuvvDin1pK8aKteLpeZ5c0A.woff2": {
              "data": "[encoded data, 14332 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "font/woff2"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/index.js?_file_hash=62538f61": {
              "data": "[encoded data, 7730416 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "text/javascript; charset=utf-8"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/font-roboto/RxZJdnzeo3R5zSexge8UUZBw1xU1rKptJj_0jans920.woff2": {
              "data": "[encoded data, 14384 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "font/woff2"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/plugins_listing": {
              "data": "eyJ0aW1lc2VyaWVzIjogeyJkaXNhYmxlX3JlbG9hZCI6IGZhbHNlLCAiZW5hYmxlZCI6IHRydWUsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJUaW1lIFNlcmllcyIsICJsb2FkaW5nX21lY2hhbmlzbSI6IHsidHlwZSI6ICJOR19DT01QT05FTlQifX0sICJzY2FsYXJzIjogeyJkaXNhYmxlX3JlbG9hZCI6IGZhbHNlLCAiZW5hYmxlZCI6IHRydWUsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJzY2FsYXJzIiwgImxvYWRpbmdfbWVjaGFuaXNtIjogeyJ0eXBlIjogIkNVU1RPTV9FTEVNRU5UIiwgImVsZW1lbnRfbmFtZSI6ICJ0Zi1zY2FsYXItZGFzaGJvYXJkIn19LCAiY3VzdG9tX3NjYWxhcnMiOiB7ImRpc2FibGVfcmVsb2FkIjogZmFsc2UsICJlbmFibGVkIjogZmFsc2UsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJDdXN0b20gU2NhbGFycyIsICJsb2FkaW5nX21lY2hhbmlzbSI6IHsidHlwZSI6ICJDVVNUT01fRUxFTUVOVCIsICJlbGVtZW50X25hbWUiOiAidGYtY3VzdG9tLXNjYWxhci1kYXNoYm9hcmQifX0sICJpbWFnZXMiOiB7ImRpc2FibGVfcmVsb2FkIjogZmFsc2UsICJlbmFibGVkIjogZmFsc2UsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJpbWFnZXMiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogInRmLWltYWdlLWRhc2hib2FyZCJ9fSwgImF1ZGlvIjogeyJkaXNhYmxlX3JlbG9hZCI6IGZhbHNlLCAiZW5hYmxlZCI6IGZhbHNlLCAicmVtb3ZlX2RvbSI6IGZhbHNlLCAidGFiX25hbWUiOiAiYXVkaW8iLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogInRmLWF1ZGlvLWRhc2hib2FyZCJ9fSwgImRlYnVnZ2VyLXYyIjogeyJkaXNhYmxlX3JlbG9hZCI6IGZhbHNlLCAiZW5hYmxlZCI6IGZhbHNlLCAicmVtb3ZlX2RvbSI6IGZhbHNlLCAidGFiX25hbWUiOiAiRGVidWdnZXIgVjIiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiTkdfQ09NUE9ORU5UIn19LCAiZ3JhcGhzIjogeyJkaXNhYmxlX3JlbG9hZCI6IHRydWUsICJlbmFibGVkIjogZmFsc2UsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJncmFwaHMiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogInRmLWdyYXBoLWRhc2hib2FyZCJ9fSwgImRpc3RyaWJ1dGlvbnMiOiB7ImRpc2FibGVfcmVsb2FkIjogZmFsc2UsICJlbmFibGVkIjogZmFsc2UsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJkaXN0cmlidXRpb25zIiwgImxvYWRpbmdfbWVjaGFuaXNtIjogeyJ0eXBlIjogIkNVU1RPTV9FTEVNRU5UIiwgImVsZW1lbnRfbmFtZSI6ICJ0Zi1kaXN0cmlidXRpb24tZGFzaGJvYXJkIn19LCAiaGlzdG9ncmFtcyI6IHsiZGlzYWJsZV9yZWxvYWQiOiBmYWxzZSwgImVuYWJsZWQiOiBmYWxzZSwgInJlbW92ZV9kb20iOiBmYWxzZSwgInRhYl9uYW1lIjogImhpc3RvZ3JhbXMiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogInRmLWhpc3RvZ3JhbS1kYXNoYm9hcmQifX0sICJ0ZXh0IjogeyJkaXNhYmxlX3JlbG9hZCI6IGZhbHNlLCAiZW5hYmxlZCI6IHRydWUsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJ0ZXh0IiwgImxvYWRpbmdfbWVjaGFuaXNtIjogeyJ0eXBlIjogIkNVU1RPTV9FTEVNRU5UIiwgImVsZW1lbnRfbmFtZSI6ICJ0Zi10ZXh0LWRhc2hib2FyZCJ9fSwgInByX2N1cnZlcyI6IHsiZGlzYWJsZV9yZWxvYWQiOiBmYWxzZSwgImVuYWJsZWQiOiBmYWxzZSwgInJlbW92ZV9kb20iOiBmYWxzZSwgInRhYl9uYW1lIjogIlBSIEN1cnZlcyIsICJsb2FkaW5nX21lY2hhbmlzbSI6IHsidHlwZSI6ICJDVVNUT01fRUxFTUVOVCIsICJlbGVtZW50X25hbWUiOiAidGYtcHItY3VydmUtZGFzaGJvYXJkIn19LCAicHJvZmlsZV9yZWRpcmVjdCI6IHsiZGlzYWJsZV9yZWxvYWQiOiBmYWxzZSwgImVuYWJsZWQiOiBmYWxzZSwgInJlbW92ZV9kb20iOiBmYWxzZSwgInRhYl9uYW1lIjogIlByb2ZpbGUiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogInRmLXByb2ZpbGUtcmVkaXJlY3QtZGFzaGJvYXJkIn19LCAiaHBhcmFtcyI6IHsiZGlzYWJsZV9yZWxvYWQiOiBmYWxzZSwgImVuYWJsZWQiOiBmYWxzZSwgInJlbW92ZV9kb20iOiBmYWxzZSwgInRhYl9uYW1lIjogImhwYXJhbXMiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogInRmLWhwYXJhbXMtZGFzaGJvYXJkIn19LCAibWVzaCI6IHsiZGlzYWJsZV9yZWxvYWQiOiBmYWxzZSwgImVuYWJsZWQiOiBmYWxzZSwgInJlbW92ZV9kb20iOiBmYWxzZSwgInRhYl9uYW1lIjogIm1lc2giLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiQ1VTVE9NX0VMRU1FTlQiLCAiZWxlbWVudF9uYW1lIjogIm1lc2gtZGFzaGJvYXJkIn19LCAid2l0X3JlZGlyZWN0IjogeyJkaXNhYmxlX3JlbG9hZCI6IGZhbHNlLCAiZW5hYmxlZCI6IGZhbHNlLCAicmVtb3ZlX2RvbSI6IGZhbHNlLCAidGFiX25hbWUiOiAiV2hhdC1JZiBUb29sIiwgImxvYWRpbmdfbWVjaGFuaXNtIjogeyJ0eXBlIjogIkNVU1RPTV9FTEVNRU5UIiwgImVsZW1lbnRfbmFtZSI6ICJ0Zi13aXQtcmVkaXJlY3QtZGFzaGJvYXJkIn19LCAicHJvamVjdG9yIjogeyJkaXNhYmxlX3JlbG9hZCI6IHRydWUsICJlbmFibGVkIjogZmFsc2UsICJyZW1vdmVfZG9tIjogZmFsc2UsICJ0YWJfbmFtZSI6ICJwcm9qZWN0b3IiLCAibG9hZGluZ19tZWNoYW5pc20iOiB7InR5cGUiOiAiSUZSQU1FIiwgIm1vZHVsZV9wYXRoIjogIi9kYXRhL3BsdWdpbi9wcm9qZWN0b3IvaW5kZXguanMifX19",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/icon_bundle.svg": {
              "data": "[encoded data, 19624 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "image/svg+xml; charset=utf-8"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/environment": {
              "data": "eyJ2ZXJzaW9uIjogIjIuMTUuMiIsICJkYXRhX2xvY2F0aW9uIjogIi4vbG9ncyIsICJ3aW5kb3dfdGl0bGUiOiAiIiwgImV4cGVyaW1lbnRfbmFtZSI6ICIiLCAiZXhwZXJpbWVudF9kZXNjcmlwdGlvbiI6ICIiLCAiY3JlYXRpb25fdGltZSI6IDAuMCwgImRlYnVnIjogeyJkYXRhX3Byb3ZpZGVyIjogIkdycGNEYXRhUHJvdmlkZXIoYWRkcj0nbG9jYWxob3N0OjM1MzM3JykiLCAiZmxhZ3MiOiB7ImxvZ2RpciI6ICIuL2xvZ3MiLCAibG9nZGlyX3NwZWMiOiAiIiwgImhvc3QiOiBudWxsLCAiYmluZF9hbGwiOiBmYWxzZSwgInBvcnQiOiBudWxsLCAicmV1c2VfcG9ydCI6IGZhbHNlLCAibG9hZF9mYXN0IjogImF1dG8iLCAiZXh0cmFfZGF0YV9zZXJ2ZXJfZmxhZ3MiOiAiIiwgImdycGNfY3JlZHNfdHlwZSI6ICJsb2NhbCIsICJncnBjX2RhdGFfcHJvdmlkZXIiOiAiIiwgInB1cmdlX29ycGhhbmVkX2RhdGEiOiB0cnVlLCAiZGIiOiAiIiwgImRiX2ltcG9ydCI6IGZhbHNlLCAiaW5zcGVjdCI6IGZhbHNlLCAidmVyc2lvbl90YiI6IGZhbHNlLCAidGFnIjogIiIsICJldmVudF9maWxlIjogIiIsICJwYXRoX3ByZWZpeCI6ICIiLCAid2luZG93X3RpdGxlIjogIiIsICJtYXhfcmVsb2FkX3RocmVhZHMiOiAxLCAicmVsb2FkX2ludGVydmFsIjogNS4wLCAicmVsb2FkX3Rhc2siOiAiYXV0byIsICJyZWxvYWRfbXVsdGlmaWxlIjogbnVsbCwgInJlbG9hZF9tdWx0aWZpbGVfaW5hY3RpdmVfc2VjcyI6IDg2NDAwLCAiZ2VuZXJpY19kYXRhIjogImF1dG8iLCAic2FtcGxlc19wZXJfcGx1Z2luIjoge30sICJkZXRlY3RfZmlsZV9yZXBsYWNlbWVudCI6IG51bGwsICJfX3RlbnNvcmJvYXJkX3N1YmNvbW1hbmQiOiAic2VydmUifX19",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/runs": {
              "data": "WyIuIl0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/hparams/experiment?request=%7B%22experimentName%22:%22defaultExperimentId%22%7D": {
              "data": "ewogICJuYW1lIjogIiIsCiAgImRlc2NyaXB0aW9uIjogIiIsCiAgInVzZXIiOiAiIiwKICAidGltZUNyZWF0ZWRTZWNzIjogMC4wLAogICJocGFyYW1JbmZvcyI6IFtdLAogICJtZXRyaWNJbmZvcyI6IFtdCn0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/runs": {
              "data": "WyIuIl0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/hparams/session_groups?request=%7B%22experimentName%22:%22defaultExperimentId%22,%22allowedStatuses%22:%5B%22STATUS_FAILURE%22,%22STATUS_RUNNING%22,%22STATUS_SUCCESS%22,%22STATUS_UNKNOWN%22%5D,%22colParams%22:%5B%5D,%22startIndex%22:0,%22sliceSize%22:1000000%7D": {
              "data": "ewogICJzZXNzaW9uR3JvdXBzIjogW10sCiAgInRvdGFsU2l6ZSI6IDAKfQ==",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/timeseries/tags": {
              "data": "eyJzY2FsYXJzIjogeyJydW5UYWdJbmZvIjogeyIuIjogWyJ0cmFpbi9lcG9jaCIsICJ0cmFpbi9ncmFkX25vcm0iLCAidHJhaW4vbGVhcm5pbmdfcmF0ZSIsICJ0cmFpbi9sb3NzIl19LCAidGFnRGVzY3JpcHRpb25zIjoge319LCAiaGlzdG9ncmFtcyI6IHsicnVuVGFnSW5mbyI6IHt9LCAidGFnRGVzY3JpcHRpb25zIjoge319LCAiaW1hZ2VzIjogeyJ0YWdEZXNjcmlwdGlvbnMiOiB7fSwgInRhZ1J1blNhbXBsZWRJbmZvIjoge319fQ==",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/font-roboto/vPcynSL0qHq_6dX7lKVByXYhjbSpvc47ee6xR_80Hnw.woff2": {
              "data": "[encoded data, 16472 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "font/woff2"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/font-roboto/d-6IYplOFocCacKzxwXSOJBw1xU1rKptJj_0jans920.woff2": {
              "data": "[encoded data, 14352 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "font/woff2"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/timeseries/timeSeries?requests=%5B%7B%22plugin%22:%22scalars%22,%22tag%22:%22train/grad_norm%22%7D%5D": {
              "data": "W3sicGx1Z2luIjogInNjYWxhcnMiLCAidGFnIjogInRyYWluL2dyYWRfbm9ybSIsICJydW5Ub1NlcmllcyI6IHsiLiI6IFt7IndhbGxUaW1lIjogMTcxMTY3NjEwNy41NzEwODU1LCAic3RlcCI6IDEwLCAidmFsdWUiOiAxNzcuMjc5ODQ2MTkxNDA2MjV9LCB7IndhbGxUaW1lIjogMTcxMTY3NjEyNS4zNDg1MjA1LCAic3RlcCI6IDIwLCAidmFsdWUiOiAxODQuNjU2NDk0MTQwNjI1fV19fV0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/timeseries/timeSeries?requests=%5B%7B%22plugin%22:%22scalars%22,%22tag%22:%22train/learning_rate%22%7D%5D": {
              "data": "W3sicGx1Z2luIjogInNjYWxhcnMiLCAidGFnIjogInRyYWluL2xlYXJuaW5nX3JhdGUiLCAicnVuVG9TZXJpZXMiOiB7Ii4iOiBbeyJ3YWxsVGltZSI6IDE3MTE2NzYxMDcuNTcxMTMyNCwgInN0ZXAiOiAxMCwgInZhbHVlIjogMC4wMDE5OTczNDE4Njc1MzYzMDY0fV19fV0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/timeseries/timeSeries?requests=%5B%7B%22plugin%22:%22scalars%22,%22tag%22:%22train/epoch%22%7D%5D": {
              "data": "W3sicGx1Z2luIjogInNjYWxhcnMiLCAidGFnIjogInRyYWluL2Vwb2NoIiwgInJ1blRvU2VyaWVzIjogeyIuIjogW3sid2FsbFRpbWUiOiAxNzExNjc2MTA3LjU3MTE2NzIsICJzdGVwIjogMTAsICJ2YWx1ZSI6IDAuMH1dfX1d",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/chart_worker.js?_file_hash=2c1e71c5": {
              "data": "[encoded data, 843288 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "text/javascript; charset=utf-8"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/experiment/defaultExperimentId/data/plugin/timeseries/timeSeries?requests=%5B%7B%22plugin%22:%22scalars%22,%22tag%22:%22train/loss%22%7D%5D": {
              "data": "W3sicGx1Z2luIjogInNjYWxhcnMiLCAidGFnIjogInRyYWluL2xvc3MiLCAicnVuVG9TZXJpZXMiOiB7Ii4iOiBbeyJ3YWxsVGltZSI6IDE3MTE2NzYxMDcuNTcwOTYxNSwgInN0ZXAiOiAxMCwgInZhbHVlIjogMS43MTI5MDAwNDI1MzM4NzQ1fV19fV0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/plugin/scalars/tags": {
              "data": "eyIuIjogeyJ0cmFpbi9lcG9jaCI6IHsiZGlzcGxheU5hbWUiOiAiIiwgImRlc2NyaXB0aW9uIjogIiJ9LCAidHJhaW4vbGVhcm5pbmdfcmF0ZSI6IHsiZGlzcGxheU5hbWUiOiAiIiwgImRlc2NyaXB0aW9uIjogIiJ9LCAidHJhaW4vbG9zcyI6IHsiZGlzcGxheU5hbWUiOiAiIiwgImRlc2NyaXB0aW9uIjogIiJ9LCAidHJhaW4vZ3JhZF9ub3JtIjogeyJkaXNwbGF5TmFtZSI6ICIiLCAiZGVzY3JpcHRpb24iOiAiIn19fQ==",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/plugin/scalars/scalars?run=.&tag=train%2Floss": {
              "data": "W1sxNzExNjc2MTA3LjU3MDk2MTUsIDEwLCAxLjcxMjkwMDA0MjUzMzg3NDVdLCBbMTcxMTY3NjEyNS4zNDg0NDgsIDIwLCAyLjIxMzk5OTk4NjY0ODU1OTZdLCBbMTcxMTY3NjE0OC43NTU1NTQyLCAzMCwgMi43MTcwOTk5MDUwMTQwMzhdLCBbMTcxMTY3NjE3MS4zODEwMiwgNDAsIDIuNzE5NzAwMDk4MDM3NzE5N10sIFsxNzExNjc2MjA1LjQxOTU2MywgNTAsIDIuMTA2MzAwMTE1NTg1MzI3XSwgWzE3MTE2NzYyMzEuNzU0NzE4NSwgNjAsIDMuMjQ2NzAwMDQ4NDQ2NjU1M10sIFsxNzExNjc2MjUzLjU2Njk0MSwgNzAsIDMuNjM5NDAwMDA1MzQwNTc2XSwgWzE3MTE2NzYyNzkuMjc3OTYyNCwgODAsIDQuMDM2Nzk5OTA3Njg0MzI2XSwgWzE3MTE2NzYyOTcuOTkxODc3NiwgOTAsIDQuNDUwNzk5OTQyMDE2NjAyXSwgWzE3MTE2NzYzMTguODYyOTgxNiwgMTAwLCA0LjQ5NTYwMDIyMzU0MTI2XSwgWzE3MTE2NzYzNDIuNjU5MTM0MSwgMTEwLCA2LjM3NzY5OTg1MTk4OTc0Nl0sIFsxNzExNjc2MzYzLjY0Mjk3NSwgMTIwLCA2LjMzNTcwMDAzNTA5NTIxNV0sIFsxNzExNjc2Mzc5LjU2MzA5MjIsIDEzMCwgNi43NjIzMDAwMTQ0OTU4NV0sIFsxNzExNjc2NDAyLjAyNTEwMjQsIDE0MCwgMTMuMDA4NTAwMDk5MTgyMTI5XSwgWzE3MTE2NzY0MjMuNDI5MDkzMSwgMTUwLCAxNC44MDAxMDAzMjY1MzgwODZdLCBbMTcxMTY3NjQzNS42NzA0OTksIDE2MCwgMTIuODQ5OTAwMjQ1NjY2NTA0XSwgWzE3MTE2NzY0NTcuODk4OTQzLCAxNzAsIDE3LjY1NjYwMDk1MjE0ODQzOF0sIFsxNzExNjc2NDgxLjY1NjIzNDUsIDE4MCwgMTUuNDcyNDk5ODQ3NDEyMTFdLCBbMTcxMTY3NjUwOS43MjgwOTEyLCAxOTAsIDE1LjExNDE5OTYzODM2NjddLCBbMTcxMTY3NjUzNC40NjU3NjAyLCAyMDAsIDEzLjQ2MDkwMDMwNjcwMTY2XSwgWzE3MTE2NzY1NTguMTE5NDMxMywgMjEwLCAxMi43MDU2OTk5MjA2NTQyOTddLCBbMTcxMTY3NjU3Mi45NTY2MDk1LCAyMjAsIDEyLjQ2MzYwMDE1ODY5MTQwNl0sIFsxNzExNjc2NTg2Ljg1OTQyMzQsIDIzMCwgMTQuOTEyNDAwMjQ1NjY2NTA0XSwgWzE3MTE2NzY1OTcuNjY2MTU0OSwgMjQwLCAxMy4wMzYyOTk3MDU1MDUzNzFdLCBbMTcxMTY3NjYwOS4zMDYxNDY0LCAyNTAsIDE0LjQ1NjAwMDMyODA2Mzk2NV0sIFsxNzExNjc2NjI0LjYxNjg0NjYsIDI2MCwgMTUuNTUwNzAwMTg3NjgzMTA1XSwgWzE3MTE2NzY2MzcuMzA5Mzk1NiwgMjcwLCAxMy4wNDU3MDAwNzMyNDIxODhdLCBbMTcxMTY3NjY1My42MjY3NzI0LCAyODAsIDEwLjc1MDEwMDEzNTgwMzIyM10sIFsxNzExNjc2NjY5LjA0NjUyMjEsIDI5MCwgOS4yMTAzMDA0NDU1NTY2NF0sIFsxNzExNjc2NjgyLjgyMTQxNjEsIDMwMCwgOC42OTc3OTk2ODI2MTcxODhdLCBbMTcxMTY3NjY5NC4zMjQ2OTY1LCAzMTAsIDguMjIyNDk5ODQ3NDEyMTFdLCBbMTcxMTY3NjcxMy4wNDc0MTY3LCAzMjAsIDYuOTYxMjk5ODk2MjQwMjM0XSwgWzE3MTE2NzY3MjguMTc5MDUzOCwgMzMwLCA3LjgxNzY5OTkwOTIxMDIwNV0sIFsxNzExNjc2NzQzLjM5ODg1OTUsIDM0MCwgNy4xODY2MDAyMDgyODI0NzFdLCBbMTcxMTY3Njc2MS4xMDIxMTUyLCAzNTAsIDguNjEyNTAwMTkwNzM0ODYzXSwgWzE3MTE2NzY3NzkuMzEyNDI3LCAzNjAsIDcuMjY5Mjk5OTgzOTc4MjcxNV0sIFsxNzExNjc2Nzk2LjQ0ODYyNjMsIDM3MCwgNy4zMzUyOTk5Njg3MTk0ODJdLCBbMTcxMTY3NjgwOC4zOTQzNDI0LCAzODAsIDYuODIyNDAwMDkzMDc4NjEzXSwgWzE3MTE2NzY4MTguMDk3NzkzMywgMzkwLCA2LjgxNjI5OTkxNTMxMzcyMV0sIFsxNzExNjc2ODI5LjE3NjMxMTMsIDQwMCwgNi44NDU2MDAxMjgxNzM4MjhdLCBbMTcxMTY3Njg0Ni42NDM0NDQzLCA0MTAsIDYuOTYxNTAwMTY3ODQ2NjhdLCBbMTcxMTY3Njg1Ni42NDUxNzgsIDQyMCwgNy4wODkyMDAwMTk4MzY0MjZdLCBbMTcxMTY3Njg3My4yODczNTc4LCA0MzAsIDcuMDczNTAwMTU2NDAyNTg4XSwgWzE3MTE2NzY4ODUuODQ4NTI4MSwgNDQwLCA3LjE2NDAwMDAzNDMzMjI3NV0sIFsxNzExNjc2ODk4LjMzODgyODMsIDQ1MCwgNi4zODk5OTk4NjY0ODU1OTZdLCBbMTcxMTY3NjkwOC4yNzg0MDE2LCA0NjAsIDYuNTk1Nzk5OTIyOTQzMTE1XSwgWzE3MTE2NzY5MjIuNzg1NjI4MywgNDcwLCA2LjM4MTEwMDE3Nzc2NDg5M10sIFsxNzExNjc2OTQxLjYxNDQ1MzYsIDQ4MCwgNi41MTY2MDAxMzE5ODg1MjVdLCBbMTcxMTY3Njk2MS4xNjE5MDc3LCA0OTAsIDYuNzMwODAwMTUxODI0OTUxXSwgWzE3MTE2NzY5NzIuNTcxNDE4OCwgNTAwLCA3LjYwOTg5OTk5NzcxMTE4Ml0sIFsxNzExNjc2OTg1LjIxMzI2MDIsIDUxMCwgNi42MzY4OTk5NDgxMjAxMTddLCBbMTcxMTY3NzAwNS41OTM1NzgsIDUyMCwgNy4yMTg5MDAyMDM3MDQ4MzRdXQ==",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/plugin/scalars/scalars?run=.&tag=train%2Fgrad_norm": {
              "data": "W1sxNzExNjc2MTA3LjU3MTA4NTUsIDEwLCAxNzcuMjc5ODQ2MTkxNDA2MjVdLCBbMTcxMTY3NjEyNS4zNDg1MjA1LCAyMCwgMTg0LjY1NjQ5NDE0MDYyNV0sIFsxNzExNjc2MTQ4Ljc1NTYyOSwgMzAsIDIwMS4zMzQyODk1NTA3ODEyNV0sIFsxNzExNjc2MTcxLjM4MTA5MTgsIDQwLCA2My41NDE2NDg4NjQ3NDYwOTRdLCBbMTcxMTY3NjIwNS40MTk2NTg3LCA1MCwgNzkuNDM0ODYwMjI5NDkyMTldLCBbMTcxMTY3NjIzMS43NTQ4Mzg1LCA2MCwgOTQuNjE5MzAwODQyMjg1MTZdLCBbMTcxMTY3NjI1My41NjcwMzkzLCA3MCwgMTUzLjQxNDEwODI3NjM2NzJdLCBbMTcxMTY3NjI3OS4yNzgwNDE2LCA4MCwgMjY5Ljg4MTI4NjYyMTA5Mzc1XSwgWzE3MTE2NzYyOTcuOTkyMDA5LCA5MCwgMTgxLjQyNTI3NzcwOTk2MDk0XSwgWzE3MTE2NzYzMTguODYzMTQ0OSwgMTAwLCAxODUuODA3MjgxNDk0MTQwNjJdLCBbMTcxMTY3NjM0Mi42NTkyMTQ3LCAxMTAsIDIwODUuNDg0MTMwODU5Mzc1XSwgWzE3MTE2NzYzNjMuNjQzMDQ5LCAxMjAsIDE5Ny44MTM5OTUzNjEzMjgxMl0sIFsxNzExNjc2Mzc5LjU2MzIyMjYsIDEzMCwgMjkzLjU5MDk0MjM4MjgxMjVdLCBbMTcxMTY3NjQwMi4wMjUyMDUxLCAxNDAsIDI4NzMxLjEzODY3MTg3NV0sIFsxNzExNjc2NDIzLjQyOTE3MjMsIDE1MCwgNjE0LjEzMjg3MzUzNTE1NjJdLCBbMTcxMTY3NjQzNS42NzA1NzUxLCAxNjAsIDM0MzguODUxMDc0MjE4NzVdLCBbMTcxMTY3NjQ1Ny44OTkwMzk3LCAxNzAsIDEyMjcwLjc5MTk5MjE4NzVdLCBbMTcxMTY3NjQ4MS42NTYzNzgzLCAxODAsIDExNDAuODQxOTE4OTQ1MzEyNV0sIFsxNzExNjc2NTA5LjcyODE5MjYsIDE5MCwgNzYyMTQuMjQyMTg3NV0sIFsxNzExNjc2NTM0LjQ2NTg2MDEsIDIwMCwgMzU3NzYuOTEwMTU2MjVdLCBbMTcxMTY3NjU1OC4xMTk1NDA1LCAyMTAsIDQ5NDg3LjIwNzAzMTI1XSwgWzE3MTE2NzY1NzIuOTU2NzY0MiwgMjIwLCAxODIuNDE3NTI2MjQ1MTE3Ml0sIFsxNzExNjc2NTg2Ljg1OTQ5NjQsIDIzMCwgMjg1MC4yOTE1MDM5MDYyNV0sIFsxNzExNjc2NTk3LjY2NjIzMjMsIDI0MCwgMTIuNTg2NTgyMTgzODM3ODldLCBbMTcxMTY3NjYwOS4zMDYyMjU4LCAyNTAsIDE1NjguMjQzMTY0MDYyNV0sIFsxNzExNjc2NjI0LjYxNjk0MzgsIDI2MCwgODE0LjY5NjQxMTEzMjgxMjVdLCBbMTcxMTY3NjYzNy4zMDk1MTgzLCAyNzAsIDEyLjYzMDg2NTA5NzA0NTg5OF0sIFsxNzExNjc2NjUzLjYyNjg4MTEsIDI4MCwgOS41Mzk2NzY2NjYyNTk3NjZdLCBbMTcxMTY3NjY2OS4wNDY1OTY4LCAyOTAsIDMxLjg3MTUwMDAxNTI1ODc5XSwgWzE3MTE2NzY2ODIuODIxNDkwOCwgMzAwLCA4LjIyODk4NzY5Mzc4NjYyMV0sIFsxNzExNjc2Njk0LjMyNDc3MzMsIDMxMCwgNS43NjE0MjgzNTYxNzA2NTRdLCBbMTcxMTY3NjcxMy4wNDc1NDY5LCAzMjAsIDQuNzY1NzY1MTkwMTI0NTEyXSwgWzE3MTE2NzY3MjguMTc5MTI3MiwgMzMwLCA4LjIyNzczMzYxMjA2MDU0N10sIFsxNzExNjc2NzQzLjM5ODk0NiwgMzQwLCA5LjA0MTc5NDc3NjkxNjUwNF0sIFsxNzExNjc2NzYxLjEwMjE5NTcsIDM1MCwgNC40NDU0NzQ2MjQ2MzM3ODldLCBbMTcxMTY3Njc3OS4zMTI1MTIyLCAzNjAsIDQuMzI2MDE2NDI2MDg2NDI2XSwgWzE3MTE2NzY3OTYuNDQ4NzAxLCAzNzAsIDYyMC45ODY0NTAxOTUzMTI1XSwgWzE3MTE2NzY4MDguMzk0NDE4NSwgMzgwLCA1LjY1Mjk1MDc2MzcwMjM5M10sIFsxNzExNjc2ODE4LjA5Nzg5NzUsIDM5MCwgNS4wODY5NDQxMDMyNDA5NjddLCBbMTcxMTY3NjgyOS4xNzY0MDUsIDQwMCwgNC4xNTg1NjQ1Njc1NjU5MThdLCBbMTcxMTY3Njg0Ni42NDM1MzIsIDQxMCwgMy41NjQzNzQ2ODUyODc0NzU2XSwgWzE3MTE2NzY4NTYuNjQ1MjU0MSwgNDIwLCA5LjY5MDkyNDY0NDQ3MDIxNV0sIFsxNzExNjc2ODczLjI4NzQ1MjIsIDQzMCwgNS4wMDIzODUxMzk0NjUzMzJdLCBbMTcxMTY3Njg4NS44NDg2NTcxLCA0NDAsIDMuMzE0MTI2MjUzMTI4MDUxOF0sIFsxNzExNjc2ODk4LjMzODkwMDMsIDQ1MCwgNi40OTYyNjg3NDkyMzcwNjA1XSwgWzE3MTE2NzY5MDguMjc4NDczNCwgNDYwLCA0LjI0NzE1OTQ4MTA0ODU4NF0sIFsxNzExNjc2OTIyLjc4NTcwMjUsIDQ3MCwgMy4zMDM4NDA2MzcyMDcwMzEyXSwgWzE3MTE2NzY5NDEuNjE0NTI4NCwgNDgwLCAyLjU0MTExNDU2ODcxMDMyN10sIFsxNzExNjc2OTYxLjE2MTk4MzcsIDQ5MCwgMi45OTYyNjA2NDMwMDUzNzFdLCBbMTcxMTY3Njk3Mi41NzE1Mzc1LCA1MDAsIDQuODk1NDc0NDMzODk4OTI2XSwgWzE3MTE2NzY5ODUuMjEzMzgzNywgNTEwLCAzLjAwODU2NzA5NDgwMjg1NjRdLCBbMTcxMTY3NzAwNS41OTM2NTg3LCA1MjAsIDQuOTQ0NDU4NDg0NjQ5NjU4XV0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/plugin/scalars/scalars?run=.&tag=train%2Flearning_rate": {
              "data": "W1sxNzExNjc2MTA3LjU3MTEzMjQsIDEwLCAwLjAwMTk5NzM0MTg2NzUzNjMwNjRdLCBbMTcxMTY3NjEyNS4zNDg1NDY1LCAyMCwgMC4wMDE5OTQ2ODM2NDAwNzc3MV0sIFsxNzExNjc2MTQ4Ljc1NTY1NTgsIDMwLCAwLjAwMTk5MjAyNTQxMjYxOTExNF0sIFsxNzExNjc2MTcxLjM4MTExODMsIDQwLCAwLjAwMTk4OTM2NzQxNzk5MTE2MTNdLCBbMTcxMTY3NjIwNS40MTk2ODUsIDUwLCAwLjAwMTk4NjcwOTE5MDUzMjU2NV0sIFsxNzExNjc2MjMxLjc1NDg2NTIsIDYwLCAwLjAwMTk4NDA1MDk2MzA3Mzk2OV0sIFsxNzExNjc2MjUzLjU2NzA2NjQsIDcwLCAwLjAwMTk4MTM5Mjk2ODQ0NjAxNjNdLCBbMTcxMTY3NjI3OS4yNzgwNzcsIDgwLCAwLjAwMTk3ODczNDc0MDk4NzQyXSwgWzE3MTE2NzYyOTcuOTkyMDM1NCwgOTAsIDAuMDAxOTc2MDc2NTEzNTI4ODI0XSwgWzE3MTE2NzYzMTguODYzMTkxMSwgMTAwLCAwLjAwMTk3MzQxODI4NjA3MDIyNzZdLCBbMTcxMTY3NjM0Mi42NTkyNDA3LCAxMTAsIDAuMDAxOTcwNzYwMjkxNDQyMjc1XSwgWzE3MTE2NzYzNjMuNjQzMDkyMiwgMTIwLCAwLjAwMTk2ODEwMjA2Mzk4MzY3OV0sIFsxNzExNjc2Mzc5LjU2MzI0ODYsIDEzMCwgMC4wMDE5NjU0NDM4MzY1MjUwODI2XSwgWzE3MTE2NzY0MDIuMDI1MjMwNCwgMTQwLCAwLjAwMTk2Mjc4NTg0MTg5NzEzXSwgWzE3MTE2NzY0MjMuNDI5MTk4NSwgMTUwLCAwLjAwMTk2MDEyNzYxNDQzODUzMzhdLCBbMTcxMTY3NjQzNS42NzA2MDExLCAxNjAsIDAuMDAxOTU3NDY5Mzg2OTc5OTM3Nl0sIFsxNzExNjc2NDU3Ljg5OTA2NSwgMTcwLCAwLjAwMTk1NDgxMTE1OTUyMTM0MTNdLCBbMTcxMTY3NjQ4MS42NTY0MDQzLCAxODAsIDAuMDAxOTUyMTUzMTY0ODkzMzg4N10sIFsxNzExNjc2NTA5LjcyODIxODgsIDE5MCwgMC4wMDE5NDk0OTQ5Mzc0MzQ3OTI1XSwgWzE3MTE2NzY1MzQuNDY1ODg2NCwgMjAwLCAwLjAwMTk0NjgzNjgyNjM5MTUxODFdLCBbMTcxMTY3NjU1OC4xMTk1NjY0LCAyMTAsIDAuMDAxOTQ0MTc4NTk4OTMyOTIxOV0sIFsxNzExNjc2NTcyLjk1Njc5MDIsIDIyMCwgMC4wMDE5NDE1MjA0ODc4ODk2NDc1XSwgWzE3MTE2NzY1ODYuODU5NTIzLCAyMzAsIDAuMDAxOTM4ODYyMjYwNDMxMDUxM10sIFsxNzExNjc2NTk3LjY2NjI1ODYsIDI0MCwgMC4wMDE5MzYyMDQxNDkzODc3NzY5XSwgWzE3MTE2NzY2MDkuMzA2MjUxLCAyNTAsIDAuMDAxOTMzNTQ2MDM4MzQ0NTAyNF0sIFsxNzExNjc2NjI0LjYxNjk2OTgsIDI2MCwgMC4wMDE5MzA4ODc4MTA4ODU5MDYyXSwgWzE3MTE2NzY2MzcuMzA5NTU2MiwgMjcwLCAwLjAwMTkyODIyOTY5OTg0MjYzMThdLCBbMTcxMTY3NjY1My42MjY5MDczLCAyODAsIDAuMDAxOTI1NTcxNDcyMzg0MDM1Nl0sIFsxNzExNjc2NjY5LjA0NjYyMjgsIDI5MCwgMC4wMDE5MjI5MTMzNjEzNDA3NjEyXSwgWzE3MTE2NzY2ODIuODIxNTMwOCwgMzAwLCAwLjAwMTkyMDI1NTEzMzg4MjE2NV0sIFsxNzExNjc2Njk0LjMyNDc5OSwgMzEwLCAwLjAwMTkxNzU5NzAyMjgzODg5MDZdLCBbMTcxMTY3NjcxMy4wNDc1NzM4LCAzMjAsIDAuMDAxOTE0OTM4OTExNzk1NjE2MV0sIFsxNzExNjc2NzI4LjE3OTE1MzQsIDMzMCwgMC4wMDE5MTIyODA2ODQzMzcwMl0sIFsxNzExNjc2NzQzLjM5ODk4MywgMzQwLCAwLjAwMTkwOTYyMjU3MzI5Mzc0NTVdLCBbMTcxMTY3Njc2MS4xMDIyMjIsIDM1MCwgMC4wMDE5MDY5NjQzNDU4MzUxNDkzXSwgWzE3MTE2NzY3NzkuMzEyNTQ5OCwgMzYwLCAwLjAwMTkwNDMwNjIzNDc5MTg3NDldLCBbMTcxMTY3Njc5Ni40NDg3MjcxLCAzNzAsIDAuMDAxOTAxNjQ4MDA3MzMzMjc4N10sIFsxNzExNjc2ODA4LjM5NDQ0MzgsIDM4MCwgMC4wMDE4OTg5ODk4OTYyOTAwMDQzXSwgWzE3MTE2NzY4MTguMDk3OTIzNSwgMzkwLCAwLjAwMTg5NjMzMTc4NTI0NjcyOTldLCBbMTcxMTY3NjgyOS4xNzY0MzI4LCA0MDAsIDAuMDAxODkzNjczNTU3Nzg4MTMzNl0sIFsxNzExNjc2ODQ2LjY0MzU1OTcsIDQxMCwgMC4wMDE4OTEwMTU0NDY3NDQ4NTkyXSwgWzE3MTE2NzY4NTYuNjQ1MjgwNiwgNDIwLCAwLjAwMTg4ODM1NzIxOTI4NjI2M10sIFsxNzExNjc2ODczLjI4NzQ3NzcsIDQzMCwgMC4wMDE4ODU2OTkxMDgyNDI5ODg2XSwgWzE3MTE2NzY4ODUuODQ4NjgzNCwgNDQwLCAwLjAwMTg4MzA0MDg4MDc4NDM5MjRdLCBbMTcxMTY3Njg5OC4zMzg5MjU2LCA0NTAsIDAuMDAxODgwMzgyNzY5NzQxMTE4XSwgWzE3MTE2NzY5MDguMjc4NDk5NCwgNDYwLCAwLjAwMTg3NzcyNDY1ODY5Nzg0MzZdLCBbMTcxMTY3NjkyMi43ODU3MjgyLCA0NzAsIDAuMDAxODc1MDY2NDMxMjM5MjQ3M10sIFsxNzExNjc2OTQxLjYxNDU1NDYsIDQ4MCwgMC4wMDE4NzI0MDgzMjAxOTU5NzNdLCBbMTcxMTY3Njk2MS4xNjIwMSwgNDkwLCAwLjAwMTg2OTc1MDA5MjczNzM3NjddLCBbMTcxMTY3Njk3Mi41NzE1NjMsIDUwMCwgMC4wMDE4NjcwOTE5ODE2OTQxMDIzXSwgWzE3MTE2NzY5ODUuMjEzNDEwOSwgNTEwLCAwLjAwMTg2NDQzMzc1NDIzNTUwNl0sIFsxNzExNjc3MDA1LjU5MzY4NDcsIDUyMCwgMC4wMDE4NjE3NzU2NDMxOTIyMzE3XV0=",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/data/plugin/scalars/scalars?run=.&tag=train%2Fepoch": {
              "data": "W1sxNzExNjc2MTA3LjU3MTE2NzIsIDEwLCAwLjBdLCBbMTcxMTY3NjEyNS4zNDg1NjgsIDIwLCAwLjAwOTk5OTk5OTc3NjQ4MjU4Ml0sIFsxNzExNjc2MTQ4Ljc1NTY3NzcsIDMwLCAwLjAwOTk5OTk5OTc3NjQ4MjU4Ml0sIFsxNzExNjc2MTcxLjM4MTEzOTgsIDQwLCAwLjAxOTk5OTk5OTU1Mjk2NTE2NF0sIFsxNzExNjc2MjA1LjQxOTcwNiwgNTAsIDAuMDE5OTk5OTk5NTUyOTY1MTY0XSwgWzE3MTE2NzYyMzEuNzU0ODg2OSwgNjAsIDAuMDE5OTk5OTk5NTUyOTY1MTY0XSwgWzE3MTE2NzYyNTMuNTY3MDg4LCA3MCwgMC4wMjk5OTk5OTkzMjk0NDc3NDZdLCBbMTcxMTY3NjI3OS4yNzgwOTgsIDgwLCAwLjAyOTk5OTk5OTMyOTQ0Nzc0Nl0sIFsxNzExNjc2Mjk3Ljk5MjA1NjgsIDkwLCAwLjAzOTk5OTk5OTEwNTkzMDMzXSwgWzE3MTE2NzYzMTguODYzMjI1LCAxMDAsIDAuMDM5OTk5OTk5MTA1OTMwMzNdLCBbMTcxMTY3NjM0Mi42NTkyNjIsIDExMCwgMC4wMzk5OTk5OTkxMDU5MzAzM10sIFsxNzExNjc2MzYzLjY0MzExNDMsIDEyMCwgMC4wNTAwMDAwMDA3NDUwNTgwNl0sIFsxNzExNjc2Mzc5LjU2MzI2OTQsIDEzMCwgMC4wNTAwMDAwMDA3NDUwNTgwNl0sIFsxNzExNjc2NDAyLjAyNTI1MSwgMTQwLCAwLjA1OTk5OTk5ODY1ODg5NTQ5XSwgWzE3MTE2NzY0MjMuNDI5MjE5NSwgMTUwLCAwLjA1OTk5OTk5ODY1ODg5NTQ5XSwgWzE3MTE2NzY0MzUuNjcwNjIyMywgMTYwLCAwLjA1OTk5OTk5ODY1ODg5NTQ5XSwgWzE3MTE2NzY0NTcuODk5MDk0MywgMTcwLCAwLjA3MDAwMDAwMDI5ODAyMzIyXSwgWzE3MTE2NzY0ODEuNjU2NDI1MiwgMTgwLCAwLjA3MDAwMDAwMDI5ODAyMzIyXSwgWzE3MTE2NzY1MDkuNzI4MjQwNSwgMTkwLCAwLjA3OTk5OTk5ODIxMTg2MDY2XSwgWzE3MTE2NzY1MzQuNDY1OTA3NiwgMjAwLCAwLjA3OTk5OTk5ODIxMTg2MDY2XSwgWzE3MTE2NzY1NTguMTE5NTg3NCwgMjEwLCAwLjA3OTk5OTk5ODIxMTg2MDY2XSwgWzE3MTE2NzY1NzIuOTU2ODEyMSwgMjIwLCAwLjA5MDAwMDAwMzU3NjI3ODY5XSwgWzE3MTE2NzY1ODYuODU5NTQ0NSwgMjMwLCAwLjA5MDAwMDAwMzU3NjI3ODY5XSwgWzE3MTE2NzY1OTcuNjY2MjgsIDI0MCwgMC4xMDAwMDAwMDE0OTAxMTYxMl0sIFsxNzExNjc2NjA5LjMwNjI3MjUsIDI1MCwgMC4xMDAwMDAwMDE0OTAxMTYxMl0sIFsxNzExNjc2NjI0LjYxNjk5MDYsIDI2MCwgMC4xMDAwMDAwMDE0OTAxMTYxMl0sIFsxNzExNjc2NjM3LjMwOTU5MjIsIDI3MCwgMC4xMDk5OTk5OTk0MDM5NTM1NV0sIFsxNzExNjc2NjUzLjYyNjkyODMsIDI4MCwgMC4xMDk5OTk5OTk0MDM5NTM1NV0sIFsxNzExNjc2NjY5LjA0NjY0NCwgMjkwLCAwLjExOTk5OTk5NzMxNzc5MDk5XSwgWzE3MTE2NzY2ODIuODIxNTU4NywgMzAwLCAwLjExOTk5OTk5NzMxNzc5MDk5XSwgWzE3MTE2NzY2OTQuMzI0ODIwMywgMzEwLCAwLjExOTk5OTk5NzMxNzc5MDk5XSwgWzE3MTE2NzY3MTMuMDQ3NTk1MywgMzIwLCAwLjEyOTk5OTk5NTIzMTYyODQyXSwgWzE3MTE2NzY3MjguMTc5MTc0NywgMzMwLCAwLjEyOTk5OTk5NTIzMTYyODQyXSwgWzE3MTE2NzY3NDMuMzk5MDA1MiwgMzQwLCAwLjE0MDAwMDAwMDU5NjA0NjQ1XSwgWzE3MTE2NzY3NjEuMTAyMjQzNywgMzUwLCAwLjE0MDAwMDAwMDU5NjA0NjQ1XSwgWzE3MTE2NzY3NzkuMzEyNTcyNywgMzYwLCAwLjE0MDAwMDAwMDU5NjA0NjQ1XSwgWzE3MTE2NzY3OTYuNDQ4NzQ4NiwgMzcwLCAwLjE1MDAwMDAwNTk2MDQ2NDQ4XSwgWzE3MTE2NzY4MDguMzk0NDY1MiwgMzgwLCAwLjE1MDAwMDAwNTk2MDQ2NDQ4XSwgWzE3MTE2NzY4MTguMDk3OTQ0MywgMzkwLCAwLjE1OTk5OTk5NjQyMzcyMTNdLCBbMTcxMTY3NjgyOS4xNzY0NTUsIDQwMCwgMC4xNTk5OTk5OTY0MjM3MjEzXSwgWzE3MTE2NzY4NDYuNjQzNTgxLCA0MTAsIDAuMTU5OTk5OTk2NDIzNzIxM10sIFsxNzExNjc2ODU2LjY0NTMyNDcsIDQyMCwgMC4xNzAwMDAwMDE3ODgxMzkzNF0sIFsxNzExNjc2ODczLjI4NzQ5OTIsIDQzMCwgMC4xNzAwMDAwMDE3ODgxMzkzNF0sIFsxNzExNjc2ODg1Ljg0ODcwNSwgNDQwLCAwLjE4MDAwMDAwNzE1MjU1NzM3XSwgWzE3MTE2NzY4OTguMzM4OTQ2OCwgNDUwLCAwLjE4MDAwMDAwNzE1MjU1NzM3XSwgWzE3MTE2NzY5MDguMjc4NTIwMywgNDYwLCAwLjE4MDAwMDAwNzE1MjU1NzM3XSwgWzE3MTE2NzY5MjIuNzg1NzQ5NCwgNDcwLCAwLjE4OTk5OTk5NzYxNTgxNDJdLCBbMTcxMTY3Njk0MS42MTQ1NzU2LCA0ODAsIDAuMTg5OTk5OTk3NjE1ODE0Ml0sIFsxNzExNjc2OTYxLjE2MjAzMSwgNDkwLCAwLjIwMDAwMDAwMjk4MDIzMjI0XSwgWzE3MTE2NzY5NzIuNTcxNTkxNiwgNTAwLCAwLjIwMDAwMDAwMjk4MDIzMjI0XSwgWzE3MTE2NzY5ODUuMjEzNDMyLCA1MTAsIDAuMjAwMDAwMDAyOTgwMjMyMjRdLCBbMTcxMTY3NzAwNS41OTM3MDYxLCA1MjAsIDAuMjA5OTk5OTkzNDQzNDg5MDddXQ==",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/json"
                ]
              ],
              "status": 200,
              "status_text": ""
            },
            "https://localhost:6006/font-roboto/Hgo13k-tfSpn0qi1SFdUfZBw1xU1rKptJj_0jans920.woff2": {
              "data": "[encoded data, 14504 characters]",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "font/woff2"
                ]
              ],
              "status": 200,
              "status_text": ""
            }
          },
          "base_uri": "https://localhost:8080/",
          "height": 821
        }
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "        (async () => {\n",
              "            const url = new URL(await google.colab.kernel.proxyPort(6006, {'cache': true}));\n",
              "            url.searchParams.set('tensorboardColab', 'true');\n",
              "            const iframe = document.createElement('iframe');\n",
              "            iframe.src = url;\n",
              "            iframe.setAttribute('width', '100%');\n",
              "            iframe.setAttribute('height', '800');\n",
              "            iframe.setAttribute('frameborder', 0);\n",
              "            document.body.appendChild(iframe);\n",
              "        })();\n",
              "    "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.train()"
      ],
      "metadata": {
        "id": "9gGKIIyQrZTi",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "4985905a-041b-4709-d31b-fa318d3eb6f7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:transformers_modules.ai21labs.Jamba-v0.1.a5b97acdf170eba3a43ed555f61dbbf36b7e4555.modeling_jamba:The input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in torch.float16.\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/nn/modules.py:391: UserWarning: Input type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference or training speed.\n",
            "  warnings.warn('Input type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference or training speed.')\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='532' max='7524' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [ 532/7524 15:27 < 3:23:50, 0.57 it/s, Epoch 0.21/3]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>10</td>\n",
              "      <td>1.712900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>20</td>\n",
              "      <td>2.214000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>30</td>\n",
              "      <td>2.717100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>40</td>\n",
              "      <td>2.719700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>50</td>\n",
              "      <td>2.106300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>60</td>\n",
              "      <td>3.246700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>70</td>\n",
              "      <td>3.639400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>80</td>\n",
              "      <td>4.036800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>90</td>\n",
              "      <td>4.450800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>100</td>\n",
              "      <td>4.495600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>110</td>\n",
              "      <td>6.377700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>120</td>\n",
              "      <td>6.335700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>130</td>\n",
              "      <td>6.762300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>140</td>\n",
              "      <td>13.008500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>150</td>\n",
              "      <td>14.800100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>160</td>\n",
              "      <td>12.849900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>170</td>\n",
              "      <td>17.656600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>180</td>\n",
              "      <td>15.472500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>190</td>\n",
              "      <td>15.114200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>200</td>\n",
              "      <td>13.460900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>210</td>\n",
              "      <td>12.705700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>220</td>\n",
              "      <td>12.463600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>230</td>\n",
              "      <td>14.912400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>240</td>\n",
              "      <td>13.036300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>250</td>\n",
              "      <td>14.456000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>260</td>\n",
              "      <td>15.550700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>270</td>\n",
              "      <td>13.045700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>280</td>\n",
              "      <td>10.750100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>290</td>\n",
              "      <td>9.210300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>300</td>\n",
              "      <td>8.697800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>310</td>\n",
              "      <td>8.222500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>320</td>\n",
              "      <td>6.961300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>330</td>\n",
              "      <td>7.817700</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>340</td>\n",
              "      <td>7.186600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>350</td>\n",
              "      <td>8.612500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>360</td>\n",
              "      <td>7.269300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>370</td>\n",
              "      <td>7.335300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>380</td>\n",
              "      <td>6.822400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>390</td>\n",
              "      <td>6.816300</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>400</td>\n",
              "      <td>6.845600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>410</td>\n",
              "      <td>6.961500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>420</td>\n",
              "      <td>7.089200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>430</td>\n",
              "      <td>7.073500</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>440</td>\n",
              "      <td>7.164000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>450</td>\n",
              "      <td>6.390000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>460</td>\n",
              "      <td>6.595800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>470</td>\n",
              "      <td>6.381100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>480</td>\n",
              "      <td>6.516600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>490</td>\n",
              "      <td>6.730800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>500</td>\n",
              "      <td>7.609900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>510</td>\n",
              "      <td>6.636900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>520</td>\n",
              "      <td>7.218900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>530</td>\n",
              "      <td>7.160500</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/peft/utils/save_and_load.py:139: UserWarning: Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`.\n",
            "  warnings.warn(\"Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`.\")\n"
          ]
        }
      ]
    }
  ]
}