part_2/03_understanding-buffers/understanding-buffers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V6BXGeEJ_s-8"
   },
   "source": [
    "# Understanding PyTorch Buffers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aQt9Ob1Y_8EH"
   },
   "source": [
    "In essence, PyTorch buffers are tensor attributes associated with a PyTorch module or model similar to parameters, but unlike parameters, buffers are not updated during training.\n",
    "\n",
    "Buffers in PyTorch are particularly useful when dealing with GPU computations, as they need to be transferred between devices (like from CPU to GPU) alongside the model's parameters. Unlike parameters, buffers do not require gradient computation, but they still need to be on the correct device to ensure that all computations are performed correctly.\n",
    "\n",
    "In part 2, we use PyTorch buffers via `self.register_buffer`, which is only briefly explained in the book. Since the concept and purpose are not immediately clear, this code notebook offers a longer explanation with a hands-on example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dAwGo_gYLY45"
   },
   "source": [
    "## An example without buffers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0qBQC9IPAJVZ"
   },
   "source": [
    "Suppose we have the following code, which is based on code from part 2. This version has been modified to exclude buffers. It implements the causal self-attention mechanism used in LLMs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "7wx-_rokAN04"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class CausalAttentionWithoutBuffers(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length,\n",
    "                 dropout, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.d_out = d_out\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        attn_scores = queries @ keys.transpose(1, 2)\n",
    "        attn_scores.masked_fill_(\n",
    "            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nNrK-wLaNSi7"
   },
   "source": [
    "We can initialize and run the module as follows on some example data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "e1MZiIsPA0Py",
    "outputId": "ce1407c6-c082-4755-b8ad-d9adcc9f153a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]],\n",
      "\n",
      "        [[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]]])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89], # Your     (x^1)\n",
    "   [0.55, 0.87, 0.66], # journey  (x^2)\n",
    "   [0.57, 0.85, 0.64], # starts   (x^3)\n",
    "   [0.22, 0.58, 0.33], # with     (x^4)\n",
    "   [0.77, 0.25, 0.10], # one      (x^5)\n",
    "   [0.05, 0.80, 0.55]] # step     (x^6)\n",
    ")\n",
    "\n",
    "batch = torch.stack((inputs, inputs), dim=0)\n",
    "context_length = batch.shape[1]\n",
    "d_in = inputs.shape[1]\n",
    "d_out = 2\n",
    "\n",
    "ca_without_buffer = CausalAttentionWithoutBuffers(d_in, d_out, context_length, 0.0)\n",
    "\n",
    "with torch.no_grad():\n",
    "    context_vecs = ca_without_buffer(batch)\n",
    "\n",
    "print(context_vecs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7_hqz6AgCCc1"
   },
   "source": [
    "So far, everything has worked fine so far.\n",
    "\n",
    "However, when training LLMs, we typically use GPUs to accelerate the process. Therefore, let's transfer the `CausalAttentionWithoutBuffers` module onto a GPU device.\n",
    "\n",
    "Please note that this operation requires the code to be run in an environment equipped with GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PYwn44HWCPJS",
    "outputId": "d7236e0c-2a43-4770-ccc1-03c9d5d11421"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Machine has GPU: True\n"
     ]
    }
   ],
   "source": [
    "print(\"Machine has GPU:\", torch.mps.is_available())\n",
    "\n",
    "batch = batch.to(\"mps\")\n",
    "ca_without_buffer.to(\"mps\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4_lMki2_CoIR"
   },
   "source": [
    "Now, let's run the code again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 338
    },
    "id": "KE9iLcjGC1V1",
    "outputId": "ab6921c7-d7dd-44ea-9b92-1911037e3dcc"
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "expected self and mask to be on the same device, but got mask on cpu and self on mps:0",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[41], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m----> 2\u001b[0m     context_vecs \u001b[38;5;241m=\u001b[39m \u001b[43mca_without_buffer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(context_vecs)\n",
      "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/torch/nn/modules/module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1734\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1736\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/torch/nn/modules/module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1745\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1746\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1747\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m   1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n",
      "Cell \u001b[0;32mIn[1], line 23\u001b[0m, in \u001b[0;36mCausalAttentionWithoutBuffers.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     20\u001b[0m values \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mW_value(x)\n\u001b[1;32m     22\u001b[0m attn_scores \u001b[38;5;241m=\u001b[39m queries \u001b[38;5;241m@\u001b[39m keys\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m)\n\u001b[0;32m---> 23\u001b[0m \u001b[43mattn_scores\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmasked_fill_\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     24\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbool\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43mnum_tokens\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m:\u001b[49m\u001b[43mnum_tokens\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minf\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     25\u001b[0m attn_weights \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39msoftmax(\n\u001b[1;32m     26\u001b[0m     attn_scores \u001b[38;5;241m/\u001b[39m keys\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m0.5\u001b[39m, dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[1;32m     27\u001b[0m )\n\u001b[1;32m     28\u001b[0m attn_weights \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdropout(attn_weights)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: expected self and mask to be on the same device, but got mask on cpu and self on mps:0"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    context_vecs = ca_without_buffer(batch)\n",
    "\n",
    "print(context_vecs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I7V26PLrC2gk"
   },
   "source": [
    "Running the code resulted in an error. What happened? It seems like we attempted a matrix multiplication between a tensor on a GPU and a tensor on a CPU. But we moved the module to the GPU!?\n",
    "\n",
    "\n",
    "Let's double-check the device locations of some of the tensors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vvYDPBRIDHfU",
    "outputId": "4b9703a8-7035-4a2d-8643-c64d37b7abd2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W_query.device: mps:0\n",
      "mask.device: cpu\n"
     ]
    }
   ],
   "source": [
    "print(\"W_query.device:\", ca_without_buffer.W_query.weight.device)\n",
    "print(\"mask.device:\", ca_without_buffer.mask.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d11nX-FFOJ3C",
    "outputId": "1e92b0e8-dbc6-41f9-e88f-5d06e0726050"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(ca_without_buffer.mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ojay-KY-DL5M"
   },
   "source": [
    "As we can see, the `mask` was not moved onto the GPU. That's because it's not a PyTorch parameter like the weights (e.g., `W_query.weight`).\n",
    "\n",
    "This means we  have to manually move it to the GPU via `.to(\"mps\")`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QYirQ63zDYsW",
    "outputId": "304628ac-bc4c-49c2-a0e1-ecf9385ddcd9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mask.device: mps:0\n"
     ]
    }
   ],
   "source": [
    "ca_without_buffer.mask = ca_without_buffer.mask.to(\"mps\")\n",
    "print(\"mask.device:\", ca_without_buffer.mask.device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4OoTqzkpDfAm"
   },
   "source": [
    "Let's try our code again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WfF0yBZODdAZ",
    "outputId": "291cfb54-86e6-45f9-99d1-fa145319f379"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]],\n",
      "\n",
      "        [[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]]], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    context_vecs = ca_without_buffer(batch)\n",
    "\n",
    "print(context_vecs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oUrVgWuuD7UE"
   },
   "source": [
    "This time, it worked!\n",
    "\n",
    "However, remembering to move individual tensors to the GPU can be tedious. As we will see in the next section, it's easier to use `register_buffer` to register the `mask` as a buffer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "StS2wUrBLeuW"
   },
   "source": [
    "## An example with buffers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nEqD2NFzPO6l"
   },
   "source": [
    "Let's now modify the causal attention class to register the causal `mask` as a buffer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "id": "ndsYj3Zf6N8U"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class CausalAttentionWithBuffer(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length,\n",
    "                 dropout, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.d_out = d_out\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        # Old:\n",
    "        # self.mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
    "\n",
    "        # New:\n",
    "        self.register_buffer(\"mask\", torch.triu(torch.ones(context_length, context_length), diagonal=1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        attn_scores = queries @ keys.transpose(1, 2)\n",
    "        attn_scores.masked_fill_(\n",
    "            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_AL1X6y3Eb7S"
   },
   "source": [
    "Now, conveniently, if we move the module to the GPU, the mask will be located on the GPU as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8_VCxEa76j00",
    "outputId": "4d1af501-5a9e-46aa-b1ac-63bf0c68e02a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W_query.device: mps:0\n",
      "mask.device: mps:0\n"
     ]
    }
   ],
   "source": [
    "ca_with_buffer = CausalAttentionWithBuffer(d_in, d_out, context_length, 0.0)\n",
    "ca_with_buffer.to(\"mps\")\n",
    "\n",
    "print(\"W_query.device:\", ca_with_buffer.W_query.weight.device)\n",
    "print(\"mask.device:\", ca_with_buffer.mask.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TBWvKlMe7bbB",
    "outputId": "e43bf8ab-3fb9-417e-d087-560858332d86"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.0005,  0.4467],\n",
      "         [-0.0572,  0.4317],\n",
      "         [-0.0760,  0.4206],\n",
      "         [-0.0796,  0.3752],\n",
      "         [-0.0642,  0.2921],\n",
      "         [-0.0745,  0.3114]],\n",
      "\n",
      "        [[ 0.0005,  0.4467],\n",
      "         [-0.0572,  0.4317],\n",
      "         [-0.0760,  0.4206],\n",
      "         [-0.0796,  0.3752],\n",
      "         [-0.0642,  0.2921],\n",
      "         [-0.0745,  0.3114]]], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    context_vecs = ca_with_buffer(batch)\n",
    "\n",
    "print(context_vecs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xvOTh4NNPjef"
   },
   "source": [
    "As we can see above, registering a tensor as a buffer can make our lives a lot easier: We don't have to remember to move tensors to a target device like a GPU manually."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q-5YYKmJte3h"
   },
   "source": [
    "## Buffers and `state_dict`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YIHHawPbtjfp"
   },
   "source": [
    "- Another advantage of PyTorch buffers, over regular tensors, is that they get included in a model's `state_dict`\n",
    "- For example, consider the `state_dict` of the causal attention object without buffers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c217juzqtxsS",
    "outputId": "dbae3c3d-f4f8-4c70-a64f-90906561d8d9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('W_query.weight',\n",
       "              tensor([[-0.2354,  0.0191, -0.2867],\n",
       "                      [ 0.2177, -0.4919,  0.4232]], device='mps:0')),\n",
       "             ('W_key.weight',\n",
       "              tensor([[-0.4196, -0.4590, -0.3648],\n",
       "                      [ 0.2615, -0.2133,  0.2161]], device='mps:0')),\n",
       "             ('W_value.weight',\n",
       "              tensor([[-0.4900, -0.3503, -0.2120],\n",
       "                      [-0.1135, -0.4404,  0.3780]], device='mps:0'))])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ca_without_buffer.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NdmZuPaqt6aO"
   },
   "source": [
    "- The mask is not included in the `state_dict` above\n",
    "- However, the mask *is* included in the `state_dict` below, thanks to registering it as a buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uGIGQAwPt1Pl",
    "outputId": "00f9bc44-63f9-4ebc-87ea-d4b8cafd81c1"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('mask',\n",
       "              tensor([[0., 1., 1., 1., 1., 1.],\n",
       "                      [0., 0., 1., 1., 1., 1.],\n",
       "                      [0., 0., 0., 1., 1., 1.],\n",
       "                      [0., 0., 0., 0., 1., 1.],\n",
       "                      [0., 0., 0., 0., 0., 1.],\n",
       "                      [0., 0., 0., 0., 0., 0.]], device='mps:0')),\n",
       "             ('W_query.weight',\n",
       "              tensor([[ 0.2545,  0.2262,  0.1613],\n",
       "                      [ 0.4566, -0.2333,  0.1517]], device='mps:0')),\n",
       "             ('W_key.weight',\n",
       "              tensor([[ 0.0032, -0.4342, -0.1402],\n",
       "                      [-0.3856,  0.2553,  0.0519]], device='mps:0')),\n",
       "             ('W_value.weight',\n",
       "              tensor([[ 0.0566, -0.1752,  0.0028],\n",
       "                      [-0.1796,  0.1659,  0.5607]], device='mps:0'))])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ca_with_buffer.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ACC-a1Hnt4Zv"
   },
   "source": [
    "- A `state_dict` is useful when saving and loading trained PyTorch models, for example\n",
    "- In this particular case, saving and loading the `mask` is maybe not super useful, because it remains unchanged during training; so, for demonstration purposes, let's assume it was modified where all `1`'s were changed to `2`'s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RLm1Sw0cuhvy",
    "outputId": "4b2cc70f-1709-44e4-aa17-4e01353b86f8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 2., 2., 2., 2., 2.],\n",
       "        [0., 0., 2., 2., 2., 2.],\n",
       "        [0., 0., 0., 2., 2., 2.],\n",
       "        [0., 0., 0., 0., 2., 2.],\n",
       "        [0., 0., 0., 0., 0., 2.],\n",
       "        [0., 0., 0., 0., 0., 0.]], device='mps:0')"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ca_with_buffer.mask[ca_with_buffer.mask == 1.] = 2.\n",
    "ca_with_buffer.mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BIkGgGqqvp4S"
   },
   "source": [
    "- Then, if we save and load the model, we can see that the mask is restored with the modified value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "e8g0QHUhuVBw",
    "outputId": "cc7ee348-7f94-4117-e5cc-e0e01a94e906"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/b4/twk1kqn144z6m04fvtpv4n2h0000gn/T/ipykernel_7912/2763133880.py:4: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  new_ca_with_buffer.load_state_dict(torch.load(\"model.pth\"))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0., 2., 2., 2., 2., 2.],\n",
       "        [0., 0., 2., 2., 2., 2.],\n",
       "        [0., 0., 0., 2., 2., 2.],\n",
       "        [0., 0., 0., 0., 2., 2.],\n",
       "        [0., 0., 0., 0., 0., 2.],\n",
       "        [0., 0., 0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.save(ca_with_buffer.state_dict(), \"model.pth\")\n",
    "\n",
    "new_ca_with_buffer = CausalAttentionWithBuffer(d_in, d_out, context_length, 0.0)\n",
    "new_ca_with_buffer.load_state_dict(torch.load(\"model.pth\"))\n",
    "\n",
    "new_ca_with_buffer.mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0pPaJk7bvBD7"
   },
   "source": [
    "- This is not true if we don't use buffers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "D03w8vDyvBRS",
    "outputId": "28071601-120c-42da-b327-bb293793839f"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/b4/twk1kqn144z6m04fvtpv4n2h0000gn/T/ipykernel_7912/2023115829.py:6: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  new_ca_without_buffer.load_state_dict(torch.load(\"model.pth\"))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 1.],\n",
       "        [0., 0., 0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ca_without_buffer.mask[ca_without_buffer.mask == 1.] = 2.\n",
    "\n",
    "torch.save(ca_without_buffer.state_dict(), \"model.pth\")\n",
    "\n",
    "new_ca_without_buffer = CausalAttentionWithoutBuffers(d_in, d_out, context_length, 0.0)\n",
    "new_ca_without_buffer.load_state_dict(torch.load(\"model.pth\"))\n",
    "\n",
    "new_ca_without_buffer.mask"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "L4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.12 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}