Big Model Memory Leak

scratchpad.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/muellerzr/276d98b99323653069a1f189c17c9698/scratchpad.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": {
        "id": "lIYdn1woOS1n"
      },
      "outputs": [],
      "source": [
        "!pip install transformers accelerate >> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!git lfs install\n",
        "!git clone https://huggingface.co/PygmalionAI/pygmalion-6b"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZztC94-4hXMA",
        "outputId": "07338ba6-066c-4f33-dae3-65b1447362a1"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Error: Failed to call git rev-parse --git-dir: exit status 128 \n",
            "Git LFS initialized.\n",
            "Cloning into 'pygmalion-6b'...\n",
            "remote: Enumerating objects: 75, 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 75 (delta 0), reused 0 (delta 0), pack-reused 72\u001b[K\n",
            "Unpacking objects: 100% (75/75), 1.10 MiB | 5.85 MiB/s, done.\n",
            "Filtering content: 100% (14/14), 3.27 GiB | 26.54 MiB/s, done.\n",
            "Encountered 2 file(s) that may not have been copied correctly on Windows:\n",
            "\tpytorch_model-00002-of-00002.bin\n",
            "\tpytorch_model-00001-of-00002.bin\n",
            "\n",
            "See: `git lfs help smudge` for more details.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import psutil\n",
        "import gc\n",
        "from accelerate import init_empty_weights, load_checkpoint_and_dispatch\n",
        "from accelerate.hooks import remove_hook_from_submodules\n",
        "from transformers import AutoConfig, AutoModelForCausalLM\n",
        "\n",
        "available_memory = lambda: psutil.virtual_memory().available * 100 / psutil.virtual_memory().total\n",
        "\n",
        "\n",
        "available_memory()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5yjZQc27f54o",
        "outputId": "11d8aea6-d2c8-4ca7-8372-77d5846e877d"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "97.34051905889982"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint = \"PygmalionAI/pygmalion-6b\"\n",
        "config = AutoConfig.from_pretrained(checkpoint)\n",
        "\n",
        "with init_empty_weights():\n",
        "    model = AutoModelForCausalLM.from_config(config)"
      ],
      "metadata": {
        "id": "FDcZyiVChkUR"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "available_memory()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dfUOZnkOhtsC",
        "outputId": "077f6aa6-7225-48cb-ec34-62734c96fdf7"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "97.12111904380251"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "loaded_model = load_checkpoint_and_dispatch(\n",
        "    model, \"pygmalion-6b\", device_map=\"auto\", no_split_module_classes=[\"GPTJBlock\"]\n",
        ")"
      ],
      "metadata": {
        "id": "53DzFyDShuzR"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "available_memory()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6saG93edh2kZ",
        "outputId": "ff76cdcc-39ce-43fb-a2fc-43b0d8311c38"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "94.97642734983059"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "remove_hook_from_submodules(loaded_model)"
      ],
      "metadata": {
        "id": "UjNrJPW2kv_U"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "available_memory()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EIcRRYo8k0Mw",
        "outputId": "a226f600-080c-48b6-f52a-11a28df4f46d"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "94.97754228482349"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model = None\n",
        "loaded_model = None\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UovzQaNckQ2F",
        "outputId": "f7821dee-373f-42f6-8de3-55b1c0b72d36"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "available_memory()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "guspMT2MkT3v",
        "outputId": "2642f655-1209-4ecc-aa43-61cf8706e511"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "95.02998992354654"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "78Bs5LbhkbMs"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "name": "scratchpad",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU",
    "gpuClass": "premium"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}