WW_SVDSMoothing_TinyLLaMAipynb.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CalculatedContent/WeightWatcher-Examples/blob/main/WW_SVDSMoothing_TinyLLaMAipynb.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Lb85_Xp4gKSk",
        "outputId": "cd11f023-2468-4fba-ab95-a0d4a62554b4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: weightwatcher in /usr/local/lib/python3.10/dist-packages (0.7.4.7)\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (0.27.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (1.23.5)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (1.5.3)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (3.7.1)\n",
            "Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (0.1.6)\n",
            "Requirement already satisfied: powerlaw in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (1.5)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (1.2.2)\n",
            "Requirement already satisfied: safetensors in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (0.4.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from weightwatcher) (4.66.1)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from accelerate) (23.2)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate) (5.9.5)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from accelerate) (6.0.1)\n",
            "Requirement already satisfied: torch>=1.10.0 in /usr/local/lib/python3.10/dist-packages (from accelerate) (2.1.0+cu121)\n",
            "Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.10/dist-packages (from accelerate) (0.20.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (4.9.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (2.1.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface-hub->accelerate) (2.31.0)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (1.2.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (4.48.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (1.4.5)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (9.4.0)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (3.1.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->weightwatcher) (2.8.2)\n",
            "Requirement already satisfied: traitlets in /usr/local/lib/python3.10/dist-packages (from matplotlib-inline->weightwatcher) (5.7.1)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->weightwatcher) (2023.4)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from powerlaw->weightwatcher) (1.11.4)\n",
            "Requirement already satisfied: mpmath in /usr/local/lib/python3.10/dist-packages (from powerlaw->weightwatcher) (1.3.0)\n",
            "Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->weightwatcher) (1.3.2)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->weightwatcher) (3.2.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib->weightwatcher) (1.16.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.10.0->accelerate) (2.1.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub->accelerate) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub->accelerate) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub->accelerate) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub->accelerate) (2024.2.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install weightwatcher accelerate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "id": "OAkGnR1FDJyg",
        "outputId": "54288d5e-9650-4cb1-edb9-b5be36a14023"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:weightwatcher:PyTorch is available but CUDA is not. Defaulting to SciPy for SVD\n",
            "WARNING:weightwatcher:Import error , reetting to svd accurate methods\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.7.4.7'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 2
        }
      ],
      "source": [
        "import weightwatcher as ww\n",
        "ww.__version__"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "_aRmY5RkgE0i",
        "outputId": "e0dbd85d-0918-416b-9b32-30c300cea242"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.27.2'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "import accelerate\n",
        "accelerate.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lm9AtCtfhAaN"
      },
      "source": [
        "### Get TinyLlama directly from github"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iUS9y5yAgQhL",
        "outputId": "67988548-cfee-46f3-b418-d9e8de3c15d8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'TinyLlama-1.1B-intermediate-step-955k-token-2T'...\n",
            "remote: Enumerating objects: 25, done.\u001b[K\n",
            "remote: Counting objects: 100% (3/3), done.\u001b[K\n",
            "remote: Compressing objects: 100% (3/3), done.\u001b[K\n",
            "remote: Total 25 (delta 0), reused 0 (delta 0), pack-reused 22\u001b[K\n",
            "Unpacking objects: 100% (25/25), 481.34 KiB | 6.78 MiB/s, done.\n",
            "Filtering content: 100% (3/3), 201.26 MiB | 1022.00 KiB/s, done.\n",
            "Encountered 2 file(s) that may not have been copied correctly on Windows:\n",
            "\tmodel.safetensors\n",
            "\tpytorch_model.bin\n",
            "\n",
            "See: `git lfs help smudge` for more details.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://huggingface.co/TinyLlama/TinyLlama-1.1B-intermediate-step-955k-token-2T/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "coIUIi1Wg6_h",
        "outputId": "9136a167-719c-4296-a777-e50744446980"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "config.json\t\tpytorch_model.bin\t tokenizer_config.json\n",
            "generation_config.json\tREADME.md\t\t tokenizer.json\n",
            "model.safetensors\tspecial_tokens_map.json  tokenizer.model\n"
          ]
        }
      ],
      "source": [
        "!ls TinyLlama-1.1B-intermediate-step-955k-token-2T"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-laz1wRwneZX"
      },
      "source": [
        "### Lets make sure HF reads the pytorch_model.bin, not the safetensors file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "P57NoLcUnaNh"
      },
      "outputs": [],
      "source": [
        "!rm TinyLlama-1.1B-intermediate-step-955k-token-2T/model.safetensors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "VhVTOmcog44d"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "tinyLLaMA_folder = \"TinyLlama-1.1B-intermediate-step-955k-token-2T\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "amSpTi-ao9mA"
      },
      "source": [
        "### Create Folder for Smoothed Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "0Z8GnVVwo-Kh"
      },
      "outputs": [],
      "source": [
        "smoothed_model_folder = \"smoothed_TinyLLaMA\"\n",
        "smoothed_model_filename = os.path.join(smoothed_model_folder, \"pytorch_model.bin\")\n",
        "!cp -r $tinyLLaMA_folder $smoothed_model_folder\n",
        "!rm $smoothed_model_filename"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "t0uV_qDepJdW",
        "outputId": "9f1ef5eb-5e59-40e9-f2a6-ffb1f933512b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "config.json\t\tREADME.md\t\t tokenizer_config.json\ttokenizer.model\n",
            "generation_config.json\tspecial_tokens_map.json  tokenizer.json\n"
          ]
        }
      ],
      "source": [
        "!ls $smoothed_model_folder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9z0fCCgkgtqk"
      },
      "source": [
        "### Check that TinyLlama can be loaded properly and generate text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "i-CbsjK3iVmw"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer, pipeline\n",
        "import torch\n",
        "\n",
        "# Initialize the tokenizer from the local directory\n",
        "tokenizer = AutoTokenizer.from_pretrained(tinyLLaMA_folder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "eXpYNAABgl3d"
      },
      "outputs": [],
      "source": [
        "# Manually set the device you want to use (e.g., 'cuda' for GPU or 'cpu' for CPU)\n",
        "# If you want to automatically use GPU if available, you can use torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "# Initialize the pipeline and specify the local model directory\n",
        "text_generation_pipeline = pipeline(\n",
        "    \"text-generation\",\n",
        "    model=tinyLLaMA_folder,\n",
        "    tokenizer=tokenizer,\n",
        "    device=device,  # Use the manually specified device\n",
        "    framework=\"pt\",  # Specify the framework 'pt' for PyTorch\n",
        "#    torch_dtype=torch.float16. # use for GPU\n",
        ")\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p-8bL-fThTdk",
        "outputId": "d0336f49-9df4-41c0-97fb-7fe35ebdd530"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[{'generated_text': 'Who was the first US president ?\\nThe first US president was George Washington.\\nWhat was the'}]\n"
          ]
        }
      ],
      "source": [
        "# Use the pipeline for text generation (as an example)\n",
        "generated_text = text_generation_pipeline(\"Who was the first US president ?\", max_length=20)\n",
        "print(generated_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3a1sB93iim7T"
      },
      "source": [
        "### Load TinyLLaMA into memory\n",
        "\n",
        "SVDSMoothing does not support reading the safetensors format yet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "6Rr0ufbdhKkU"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "tinyLLaMA_filename = os.path.join(tinyLLaMA_folder, \"pytorch_model.bin\")\n",
        "tinyLLaMA = torch.load(tinyLLaMA_filename)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4RXVYdOHi8mT"
      },
      "source": [
        "### Test WeightWatcher runs and can read the model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 423
        },
        "id": "aniBTcHjgexs",
        "outputId": "4392e594-dca7-4ed7-d311-77594e72a2b5"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "     layer_id                              name     M      N       Q  \\\n",
              "0           1                           lm_head  2048  32000  15.625   \n",
              "1           2                model.embed_tokens  2048  32000  15.625   \n",
              "2           4   model.layers.0.self_attn.q_proj  2048   2048   1.000   \n",
              "3           5   model.layers.0.self_attn.k_proj   256   2048   8.000   \n",
              "4           6   model.layers.0.self_attn.v_proj   256   2048   8.000   \n",
              "..        ...                               ...   ...    ...     ...   \n",
              "151       195  model.layers.21.self_attn.v_proj   256   2048   8.000   \n",
              "152       196  model.layers.21.self_attn.o_proj  2048   2048   1.000   \n",
              "153       198     model.layers.21.mlp.gate_proj  2048   5632   2.750   \n",
              "154       199       model.layers.21.mlp.up_proj  2048   5632   2.750   \n",
              "155       200     model.layers.21.mlp.down_proj  2048   5632   2.750   \n",
              "\n",
              "    layer_type                          longname  num_evals  rf  \n",
              "0        dense                           lm_head       2048   1  \n",
              "1        dense                model.embed_tokens       2048   1  \n",
              "2        dense   model.layers.0.self_attn.q_proj       2048   1  \n",
              "3        dense   model.layers.0.self_attn.k_proj        256   1  \n",
              "4        dense   model.layers.0.self_attn.v_proj        256   1  \n",
              "..         ...                               ...        ...  ..  \n",
              "151      dense  model.layers.21.self_attn.v_proj        256   1  \n",
              "152      dense  model.layers.21.self_attn.o_proj       2048   1  \n",
              "153      dense     model.layers.21.mlp.gate_proj       2048   1  \n",
              "154      dense       model.layers.21.mlp.up_proj       2048   1  \n",
              "155      dense     model.layers.21.mlp.down_proj       2048   1  \n",
              "\n",
              "[156 rows x 9 columns]"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-7e3d8d9e-4b3e-4d65-a7c2-f72447d94f2d\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>layer_id</th>\n",
              "      <th>name</th>\n",
              "      <th>M</th>\n",
              "      <th>N</th>\n",
              "      <th>Q</th>\n",
              "      <th>layer_type</th>\n",
              "      <th>longname</th>\n",
              "      <th>num_evals</th>\n",
              "      <th>rf</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>lm_head</td>\n",
              "      <td>2048</td>\n",
              "      <td>32000</td>\n",
              "      <td>15.625</td>\n",
              "      <td>dense</td>\n",
              "      <td>lm_head</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>2</td>\n",
              "      <td>model.embed_tokens</td>\n",
              "      <td>2048</td>\n",
              "      <td>32000</td>\n",
              "      <td>15.625</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.embed_tokens</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>4</td>\n",
              "      <td>model.layers.0.self_attn.q_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>2048</td>\n",
              "      <td>1.000</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.0.self_attn.q_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>5</td>\n",
              "      <td>model.layers.0.self_attn.k_proj</td>\n",
              "      <td>256</td>\n",
              "      <td>2048</td>\n",
              "      <td>8.000</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.0.self_attn.k_proj</td>\n",
              "      <td>256</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>6</td>\n",
              "      <td>model.layers.0.self_attn.v_proj</td>\n",
              "      <td>256</td>\n",
              "      <td>2048</td>\n",
              "      <td>8.000</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.0.self_attn.v_proj</td>\n",
              "      <td>256</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>151</th>\n",
              "      <td>195</td>\n",
              "      <td>model.layers.21.self_attn.v_proj</td>\n",
              "      <td>256</td>\n",
              "      <td>2048</td>\n",
              "      <td>8.000</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.21.self_attn.v_proj</td>\n",
              "      <td>256</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>152</th>\n",
              "      <td>196</td>\n",
              "      <td>model.layers.21.self_attn.o_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>2048</td>\n",
              "      <td>1.000</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.21.self_attn.o_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>153</th>\n",
              "      <td>198</td>\n",
              "      <td>model.layers.21.mlp.gate_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>5632</td>\n",
              "      <td>2.750</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.21.mlp.gate_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>154</th>\n",
              "      <td>199</td>\n",
              "      <td>model.layers.21.mlp.up_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>5632</td>\n",
              "      <td>2.750</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.21.mlp.up_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>155</th>\n",
              "      <td>200</td>\n",
              "      <td>model.layers.21.mlp.down_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>5632</td>\n",
              "      <td>2.750</td>\n",
              "      <td>dense</td>\n",
              "      <td>model.layers.21.mlp.down_proj</td>\n",
              "      <td>2048</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>156 rows × 9 columns</p>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-7e3d8d9e-4b3e-4d65-a7c2-f72447d94f2d')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-7e3d8d9e-4b3e-4d65-a7c2-f72447d94f2d button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-7e3d8d9e-4b3e-4d65-a7c2-f72447d94f2d');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-4d31fa27-9c5e-4d22-9771-6824e043c64a\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-4d31fa27-9c5e-4d22-9771-6824e043c64a')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-4d31fa27-9c5e-4d22-9771-6824e043c64a button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "    </div>\n",
              "  </div>\n"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "source": [
        "watcher = ww.WeightWatcher(model=tinyLLaMA)\n",
        "details = watcher.describe()\n",
        "details"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r7j6x4Adjb9w"
      },
      "source": [
        "### Select the MLP only layers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "CqAF5fm9s9T6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a66e0c3a-69d8-4771-e3ac-c903f2da41a1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of layers with MLP in name 66\n"
          ]
        }
      ],
      "source": [
        "import pandas as pd\n",
        "\n",
        "# Assuming 'details' is your DataFrame\n",
        "D = details[details['name'].astype(str).str.contains('mlp')]\n",
        "\n",
        "# Now, extract 'layer_id' column as a list of ids\n",
        "mlp_layer_ids = list(D['layer_id'].to_numpy())\n",
        "\n",
        "# If you want to see the result\n",
        "print(\"Number of layers with MLP in name\" , len(mlp_layer_ids))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "11BOSOTvjWzx"
      },
      "source": [
        "### Run LASER (SVDSMoothing)\n",
        "\n",
        "This can take some time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g86_fUKE44oN"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "watcher = ww.WeightWatcher(model=tinyLLaMA, log_level=logging.WARNING)\n",
        "smoothed_model = watcher.SVDSmoothing(layers=mlp_layer_ids)\n",
        "# alternatively\n",
        "#smoothed_model = watcher.SVDSmoothing(layers=mlp_layer_ids), method='svd',  percent=0.80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hh1RuqxgkitB"
      },
      "source": [
        "### Save model to disk"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WRsDo0D0khPB"
      },
      "outputs": [],
      "source": [
        "torch.save(smoothed_model, smoothed_model_filename)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "co0u59SZop74"
      },
      "source": [
        "### Generate text with Smoothed Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iI4sE_1yk-uS"
      },
      "outputs": [],
      "source": [
        "# Manually set the device you want to use (e.g., 'cuda' for GPU or 'cpu' for CPU)\n",
        "# If you want to automatically use GPU if available, you can use torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "# Initialize the pipeline and specify the local model directory\n",
        "smoothed_generation_pipeline = pipeline(\n",
        "    \"text-generation\",\n",
        "    model=smoothed_model_folder,\n",
        "    tokenizer=tokenizer,\n",
        "    device=device,  # Use the manually specified device\n",
        "    framework=\"pt\",  # Specify the framework 'pt' for PyTorch\n",
        "#    torch_dtype=torch.float16. # use for GPU\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dAkZua4flZjM"
      },
      "outputs": [],
      "source": [
        "# Use the pipeline for text generation (as an example)\n",
        "generated_text = text_generation_pipeline(\"Who was the first US president ?\", max_length=30)\n",
        "print(generated_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r7IDLglCmq8v"
      },
      "outputs": [],
      "source": [
        "# Use the pipeline for text generation (as an example)\n",
        "generated_text = smoothed_generation_pipeline(\"Who was the first US president ?\", max_length=30)\n",
        "print(generated_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P7-HF4Pj8jzD"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "authorship_tag": "ABX9TyOWTdtnJaVx79YCB4CtMB+Q",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}