Created
February 13, 2023 21:02
-
-
Save muellerzr/276d98b99323653069a1f189c17c9698 to your computer and use it in GitHub Desktop.
Big Model Memory Leak
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ | |
"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 | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment