yujunjun/LLMs-from-scratch
1
0
Fork 0
mirror of https://github.com/rasbt/LLMs-from-scratch.git synced 2025-11-03 03:10:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

update mmap section

Browse Source
This commit is contained in:
rasbt 2024-10-14 14:27:19 -05:00
parent 31fb74133a
commit 3567fb656d
1 changed files with 24 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
ch05/08_memory_efficient_weight_loading/memory-efficient-state-dict.ipynb
View File

@ -752,7 +752,7 @@
"metadata": {},
"source": [
" \n",
"## 6. Using `mmap=True`"
"## 6. Using `mmap=True` (recommmended)"
]
},
{
@ -760,19 +760,20 @@
"metadata": {},
"source": [
"- As an intermediate or advanced `torch.load` user, you may wonder how these approaches compare to the `mmap=True` setting in PyTorch\n",
"- The `mmap=True` setting in PyTorch enables memory-mapped file I/O, which allows the tensor to access data directly from disk storage, thus reducing memory usage by not loading the entire file into RAM\n",
"- However, in practice, I found it to be less efficient than the sequential approaches above"
"- The `mmap=True` setting in PyTorch enables memory-mapped file I/O, which allows the tensor to access data directly from disk storage, thus reducing memory usage by not loading the entire file into RAM if RAM is limited\n",
"- Also, see the helpful comment by [mikaylagawarecki](https://github.com/rasbt/LLMs-from-scratch/issues/402)\n",
"- At first glance, it may look less efficient than the sequential approaches above:"
]
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": 37,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "7AX3vPrpv5c_",
"outputId": "e6fca10b-55c3-4e89-8674-075df5ce26e7"
"id": "GKwV0AMNemuR",
"outputId": "e207f2bf-5c87-498e-80fe-e8c4016ac711"
},
"outputs": [
{
@ -780,63 +781,32 @@
"output_type": "stream",
"text": [
"Maximum GPU memory allocated: 6.4 GB\n",
"-> Maximum CPU memory allocated: 9.9 GB\n"
"-> Maximum CPU memory allocated: 5.9 GB\n"
]
}
],
"source": [
"def baseline_mmap():\n",
" start_memory_tracking()\n",
"def best_practices():\n",
" with torch.device(\"meta\"):\n",
" model = GPTModel(BASE_CONFIG)\n",
"\n",
" model = GPTModel(BASE_CONFIG) # load model on CPU\n",
" model.load_state_dict(\n",
" torch.load(\"model.pth\", map_location=device, weights_only=True, mmap=True),\n",
" assign=True\n",
" )\n",
"\n",
" model.load_state_dict(\n",
" torch.load(\"model.pth\", map_location=\"cpu\", weights_only=True, mmap=True)\n",
" )\n",
" model.to(device) # Move model to GPU\n",
" model.eval();\n",
" print_memory_usage()\n",
"\n",
" print_memory_usage()\n",
"\n",
"peak_memory_used = memory_usage_in_gb(baseline_mmap)\n",
"peak_memory_used = memory_usage_in_gb(best_practices)\n",
"print(f\"-> Maximum CPU memory allocated: {peak_memory_used:.1f} GB\")"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "KUyK3QVRwmjR",
"outputId": "a77c191a-2f9e-4ae5-be19-8ce128e704e9"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Maximum GPU memory allocated: 12.8 GB\n",
"-> Maximum CPU memory allocated: 7.0 GB\n"
]
}
],
"cell_type": "markdown",
"metadata": {},
"source": [
"def baseline_mmap_2():\n",
" start_memory_tracking()\n",
"\n",
" model = GPTModel(BASE_CONFIG).to(device)\n",
"\n",
" model.load_state_dict(\n",
" torch.load(\"model.pth\", map_location=device, weights_only=True, mmap=True)\n",
" )\n",
" model.eval();\n",
"\n",
" print_memory_usage()\n",
"\n",
"peak_memory_used = memory_usage_in_gb(baseline_mmap_2)\n",
"print(f\"-> Maximum CPU memory allocated: {peak_memory_used:.1f} GB\")"
"- The reason why the CPU RAM usage is so high is that there's enough CPU RAM available on this machine\n",
"- However, if you were to run this on a machine with limited CPU RAM, the `mmap` approach would use less memory"
]
},
{
@ -851,8 +821,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"- This notebook is focused on simple, built-in methods for loading weights in PyTorch.\n",
"- In case none of these methods work because you (1) don't have enough CPU memory for the `load_sequentially` approach and don't have enough GPU VRAM to have 2 copies of the weights in memory (the `load_sequentially_with_meta` approach), one option is to save and load each weight tensor separately:"
"- This notebook is focused on simple, built-in methods for loading weights in PyTorch\n",
"- The recommended approach for limited CPU memory cases is the `mmap=True` approach explained enough\n",
"- Alternatively, one other option is a brute-force approach that saves and loads each weight tensor separately:"
]
},
{