{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-1 fw=all -->\n",
    "\n",
    "# Parameters, State, and Memory\n",
    "\n",
    "Almost everything we do to a model other than calling it operates on its\n",
    "*state*: the optimizer updates it, a checkpoint serializes it, device\n",
    "placement moves it, fine-tuning trains part of it, and the answer to \"will\n",
    "this model fit on my GPU?\" is a few lines of arithmetic over it. So far that\n",
    "state has been handled for us; the layers created the tensors and the\n",
    "training loop updated them. This section opens the box: how to reach any tensor in a model,\n",
    "which tensors are trained and which merely travel with the model, what they\n",
    "all cost in bytes, and how to share or freeze them."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-state-and-memory"
   },
   "source": [
    "import torch\n",
    "from torch import nn"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-2 fw=all -->\n",
    "\n",
    "## Accessing Parameters\n",
    "\n",
    "Our specimen is the residual MLP of that section,\n",
    "redefined here so this section stands on its own: an input layer, a stack of\n",
    "residual blocks, and an output head."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-1"
   },
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.body = nn.Sequential(nn.Linear(d, d), nn.ReLU(), nn.Linear(d, d))\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X + self.body(X)\n",
    "\n",
    "torch.manual_seed(42)\n",
    "net = nn.Sequential(nn.Linear(20, 64), ResidualBlock(64),\n",
    "                    ResidualBlock(64), nn.Linear(64, 10))\n",
    "X = torch.randn(2, 20)\n",
    "net(X).shape"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([2, 10])"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-3 fw=pytorch -->\n",
    "\n",
    "A model built from modules is a tree, and its parameters are the leaves. To\n",
    "reach one leaf, walk the tree: indexing into a `Sequential` selects a child,\n",
    "attribute access selects within it. `net[3]` is the output layer; its bias is\n",
    "an `nn.Parameter`, a tensor that announces itself to the module system as\n",
    "trainable."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-2"
   },
   "source": [
    "type(net[3].bias), net[3].bias.shape"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(torch.nn.parameter.Parameter, torch.Size([10]))"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-4 fw=all -->\n",
    "\n",
    "The same path syntax reaches arbitrarily deep. The first linear layer inside\n",
    "the first residual block sits three levels down:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-3"
   },
   "source": [
    "net[1].body[0].weight.shape"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([64, 64])"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-5 fw=pytorch -->\n",
    "\n",
    "Each parameter also carries its gradient. We have not run backpropagation on\n",
    "this model yet, so there is nothing to see:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-4"
   },
   "source": [
    "net[3].weight.grad is None"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-6 fw=pytorch -->\n",
    "\n",
    "Reaching parameters one path at a time is right for debugging a single layer.\n",
    "The optimizer, weight decay, and checkpointing instead need *every* leaf, and\n",
    "`named_parameters()` provides exactly that: a traversal of the whole tree that\n",
    "yields each parameter with its path as the name."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-5"
   },
   "source": [
    "[(name, p.shape) for name, p in net.named_parameters()]"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[('0.weight', torch.Size([64, 20])),\n ('0.bias', torch.Size([64])),\n ('1.body.0.weight', torch.Size([64, 64])),\n ('1.body.0.bias', torch.Size([64])),\n ('1.body.2.weight', torch.Size([64, 64])),\n ('1.body.2.bias', torch.Size([64])),\n ('2.body.0.weight', torch.Size([64, 64])),\n ('2.body.0.bias', torch.Size([64])),\n ('2.body.2.weight', torch.Size([64, 64])),\n ('2.body.2.bias', torch.Size([64])),\n ('3.weight', torch.Size([10, 64])),\n ('3.bias', torch.Size([10]))]"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-7 fw=pytorch -->\n",
    "\n",
    "Read one of the names closely. `1.body.0.weight` means: child\n",
    "`1` of `net` (the first residual block), its submodule `body`, that module's\n",
    "child `0`, and finally the leaf `weight`. Names are paths, so they survive any\n",
    "amount of nesting, and they are exactly the keys of the model's `state_dict`,\n",
    "the name-to-tensor mapping used for saving and loading\n",
    "(that section):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-6"
   },
   "source": [
    "list(net.state_dict()) == [name for name, _ in net.named_parameters()]"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-8 fw=pytorch -->\n",
    "\n",
    "One tree, one naming scheme, and every consumer, whether optimizer,\n",
    "checkpoint, or debugger, walks it. The equality above holds for this model\n",
    "because all of its state happens to be trainable. That is not always so."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-9 fw=all -->\n",
    "\n",
    "## Parameters and Buffers\n",
    "\n",
    "Some tensors must persist inside a model and follow it from device to device,\n",
    "yet should never receive a gradient. The canonical example is batch\n",
    "normalization [@Ioffe.Szegedy.2015]: each layer maintains a running mean\n",
    "and variance of its inputs, updated during the forward pass and used at\n",
    "prediction time. Those statistics must be saved with the model and must move\n",
    "to the GPU with it, but the optimizer has no business touching them. A\n",
    "precomputed positional table may be treated the same way. Other non-trained\n",
    "state has a different lifetime: a language model's key--value cache belongs to\n",
    "one generation request, so it follows the computation device but should not be\n",
    "stored in a model checkpoint. Causal masks are often recomputed or registered\n",
    "as non-persistent state for the same reason."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-10 fw=pytorch -->\n",
    "\n",
    "PyTorch calls such tensors *buffers*, registered with `register_buffer`. The\n",
    "rule of thumb: make it a *parameter* if the optimizer should update it, a\n",
    "persistent *buffer* if checkpoints need it, and an explicit runtime input or\n",
    "cache if its lifetime is one request. Plain Python attributes hold\n",
    "reconstructible architecture data. Here is a module that standardizes its inputs with\n",
    "statistics computed once, ahead of time, from a reference sample:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-1"
   },
   "source": [
    "class Whitener(nn.Module):\n",
    "    def __init__(self, mean, std):\n",
    "        super().__init__()\n",
    "        self.register_buffer('mean', mean)\n",
    "        self.register_buffer('std', std)\n",
    "        self.out = nn.Linear(4, 2)\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.out((X - self.mean) / self.std)\n",
    "\n",
    "sample = torch.randn(100, 4) * torch.arange(1., 5.)\n",
    "whiten = Whitener(sample.mean(0), sample.std(0))\n",
    "list(whiten.state_dict())"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "['mean', 'std', 'out.weight', 'out.bias']"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-11 fw=pytorch -->\n",
    "\n",
    "The buffers appear in the state dict, so they are checkpointed alongside the\n",
    "weights. They do not appear among the parameters, so the optimizer never sees\n",
    "them:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-2"
   },
   "source": [
    "[name for name, _ in whiten.named_parameters()]"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "['out.weight', 'out.bias']"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-12 fw=pytorch -->\n",
    "\n",
    "Registration is what makes the module system aware of a tensor. A tensor\n",
    "stored as a plain attribute is invisible to it: not saved, and not converted\n",
    "when the model moves. We can see this on the CPU by moving the module across\n",
    "*dtypes*, which uses the same machinery as moving across devices. The\n",
    "registered buffer converts; the plain attribute is left behind:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-3"
   },
   "source": [
    "whiten.note = torch.zeros(4)   # plain attribute: invisible to the module\n",
    "whiten.to(torch.float64)\n",
    "whiten.mean.dtype, whiten.note.dtype"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(torch.float64, torch.float32)"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-13 fw=pytorch -->\n",
    "\n",
    "The device version of the same fact is the classic bug: a model works on the\n",
    "CPU, then crashes after `.to('cuda')` because an unregistered tensor stayed\n",
    "behind. On a machine with a GPU the following confirms that buffers move with\n",
    "the module:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-4"
   },
   "source": [
    "if torch.cuda.is_available():\n",
    "    print('buffer lives on', whiten.to('cuda').mean.device)\n",
    "else:\n",
    "    print('no GPU here; on a CUDA machine, whiten.to(\"cuda\") '\n",
    "          'moves whiten.mean along with the parameters')"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "buffer lives on cuda:0\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-14 fw=all -->\n",
    "\n",
    "## Counting Parameters, Counting Bytes\n",
    "\n",
    "Before any training job comes the question of whether the model fits in\n",
    "memory, and the answer is arithmetic you can do on a napkin. Counting\n",
    "parameters is one line. Counting *bytes* requires remembering everything that\n",
    "training keeps per parameter: the weight itself, its gradient, and the\n",
    "optimizer's state. Adam maintains two running moments per parameter, so in\n",
    "fp32 the ledger reads:\n",
    "\n",
    "| Training state       | Precision | Bytes per parameter |\n",
    "|----------------------|-----------|---------------------|\n",
    "| Weights              | fp32      | 4                   |\n",
    "| Gradients            | fp32      | 4                   |\n",
    "| Adam first moment    | fp32      | 4                   |\n",
    "| Adam second moment   | fp32      | 4                   |\n",
    "| Total                |           | 16                  |"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-counting-parameters-counting-bytes-1"
   },
   "source": [
    "n = sum(p.numel() for p in net.parameters())\n",
    "weights, grads, adam_state = 4 * n, 4 * n, 8 * n   # bytes, all fp32\n",
    "print(f'{n} parameters: {(weights + grads + adam_state) / 2**20:.2f} MiB '\n",
    "      'for weights + gradients + Adam state')"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "18634 parameters: 0.28 MiB for weights + gradients + Adam state\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-15 fw=all -->\n",
    "\n",
    "The optimizer state is real memory, allocated tensor by tensor. After a\n",
    "single training step, Adam's two moments together hold exactly two extra\n",
    "copies of the model:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-counting-parameters-counting-bytes-2"
   },
   "source": [
    "adam = torch.optim.Adam(net.parameters())\n",
    "net(X).sum().backward()\n",
    "adam.step()\n",
    "moments = sum(t.numel() for s in adam.state.values()\n",
    "              for t in s.values() if torch.is_tensor(t) and t.ndim > 0)\n",
    "moments == 2 * n"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-16 fw=all -->\n",
    "\n",
    "For our little network the total is a third of a megabyte, which is why none\n",
    "of this mattered until now. Scale the same arithmetic to a 1-billion-parameter\n",
    "model and the standard fp32 Adam ledger gives a batch-independent floor: 4 GB\n",
    "for the weights alone and 16 GB for weights, gradients, and optimizer state,\n",
    "before storing a single activation (that section). Activations depend on\n",
    "batch size, sequence length, and architecture; they can exceed parameter state\n",
    "and are treated in that section.\n",
    "\n",
    "Large models often train in mixed precision\n",
    "[@Micikevicius.Narang.Alben.ea.2018], but storage policies differ. One\n",
    "implementation convention counts 18 bytes per parameter (fp32 master weights,\n",
    "an fp16 working copy, fp32 gradients, and two fp32 moments); the ZeRO paper\n",
    "counts 16 by keeping gradients in fp16\n",
    "[@Rajbhandari.Rasley.Ruwase.ea.2020]. Other bf16 and AMP stacks keep no\n",
    "separate working copy or choose another dtype for optimizer state. State the\n",
    "dtypes before quoting a multiplier. In the two conventions above, the fp32\n",
    "Adam moments alone cost 8 bytes per parameter.\n",
    "\n",
    "\n",
    "![Bytes per parameter under fp32 Adam training (top) versus a mixed-precision, ZeRO-style convention (bottom), drawn to the same width per byte. The two ledgers disagree on weights and gradients but agree, byte for byte, on the Adam moments m and v.](../img/bg-memory-ledger.svg)\n",
    "\n",
    "the figure lays the two conventions side by side, byte for\n",
    "byte: the moments (right half of each bar) line up exactly, and the only\n",
    "disagreement is how the weights and gradients themselves are stored.\n",
    "\n",
    "## Tied Parameters\n",
    "\n",
    "One tensor can serve several roles in a model. The standard example is the\n",
    "two ends of a language model. The input embedding is a $|V| \\times d$ table\n",
    "mapping each of $|V|$ tokens to a $d$-dimensional vector; the output\n",
    "projection maps a $d$-dimensional hidden state to $|V|$ logits, and its weight\n",
    "matrix has the same shape and an analogous meaning: one vector per token.\n",
    "*Tying* the two, using a single tensor for both roles, saves $|V| \\times d$\n",
    "parameters. The cited studies also found lower language-model perplexity in\n",
    "their experimental settings\n",
    "[@Press.Wolf.2017; @Inan.Khosravi.Socher.2017]. The savings are large: in\n",
    "GPT-2 [@Radford.Wu.Child.ea.2019] the shared $50257 \\times 768$ matrix is\n",
    "about 38.6 million of the model's 124 million parameters, roughly 31%.\n",
    "\n",
    "![One weight matrix W, two call sites: the embedding looks up a row by token id, the output head multiplies by its transpose to produce logits. Tying means both point at the same tensor, so gradients from both uses sum into one place.](../img/bg-weight-tying.svg)\n",
    "\n",
    "the figure sketches the picture behind the two call sites."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-17 fw=pytorch -->\n",
    "\n",
    "A miniature version shows the mechanics. Tying is one assignment, and it is\n",
    "aliasing, not copying:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-1"
   },
   "source": [
    "class TinyLM(nn.Module):\n",
    "    def __init__(self, vocab_size, d, tied=True):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(vocab_size, d)\n",
    "        self.body = nn.Linear(d, d)\n",
    "        self.head = nn.Linear(d, vocab_size, bias=False)\n",
    "        if tied:\n",
    "            self.head.weight = self.emb.weight   # one tensor, two roles\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        return self.head(torch.relu(self.body(self.emb(tokens))))\n",
    "\n",
    "lm = TinyLM(vocab_size=100, d=16)\n",
    "lm.head.weight is lm.emb.weight"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-18 fw=pytorch -->\n",
    "\n",
    "The module system understands the aliasing. Parameter traversal reports the\n",
    "shared tensor once, so the count below reflects the $|V| \\times d$ saving and\n",
    "the optimizer updates the tensor once per step:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-2"
   },
   "source": [
    "untied = TinyLM(vocab_size=100, d=16, tied=False)\n",
    "(sum(p.numel() for p in lm.parameters()),\n",
    " sum(p.numel() for p in untied.parameters()))"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(1872, 3472)"
     },
     "metadata": {},
     "execution_count": 15
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-19 fw=pytorch -->\n",
    "\n",
    "The state dict, by contrast, keeps *both* names, `emb.weight` and\n",
    "`head.weight`, pointing at the same storage. That is deliberate: a checkpoint\n",
    "saved from a tied model then loads into either a tied or an untied one."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-3"
   },
   "source": [
    "sd = lm.state_dict()\n",
    "(sd['emb.weight'].data_ptr() == sd['head.weight'].data_ptr(),\n",
    " [name for name, _ in lm.named_parameters()])"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(True, ['emb.weight', 'body.weight', 'body.bias'])"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-20 fw=all -->\n",
    "\n",
    "What about gradients? During backpropagation each use of the tensor\n",
    "contributes a gradient, and the contributions accumulate into the gradient of\n",
    "the single shared tensor. We can verify this against the untied twin: load\n",
    "the tied model's values into it, run the same backward pass through both, and\n",
    "the tied gradient equals the *sum* of the untied model's two gradients."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-4"
   },
   "source": [
    "untied.load_state_dict(lm.state_dict())   # same values, separate tensors\n",
    "tokens = torch.randint(0, 100, (2, 8))\n",
    "lm(tokens).sum().backward()\n",
    "untied(tokens).sum().backward()\n",
    "torch.allclose(lm.emb.weight.grad,\n",
    "               untied.emb.weight.grad + untied.head.weight.grad)"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-21 fw=all -->\n",
    "\n",
    "## Freezing Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-22 fw=pytorch -->\n",
    "\n",
    "Every fine-tuning recipe (that section) rests on one primitive:\n",
    "setting `requires_grad = False` prevents a parameter from accumulating a\n",
    "`.grad` value. Backpropagation can still pass through operations involving\n",
    "that tensor to reach trainable inputs or earlier parameters. With no gradient\n",
    "to apply, the weight stays put. The typical pattern freezes a pretrained backbone and trains only a new\n",
    "head. On a fresh copy of our residual MLP, freezing everything but the output\n",
    "layer leaves 650 of 18,634 parameters trainable:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-1"
   },
   "source": [
    "finetune = nn.Sequential(nn.Linear(20, 64), ResidualBlock(64),\n",
    "                         ResidualBlock(64), nn.Linear(64, 10))\n",
    "for p in finetune[:-1].parameters():\n",
    "    p.requires_grad = False\n",
    "\n",
    "(sum(p.numel() for p in finetune.parameters() if p.requires_grad),\n",
    " sum(p.numel() for p in finetune.parameters()))"
   ],
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(650, 18634)"
     },
     "metadata": {},
     "execution_count": 18
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-23 fw=pytorch -->\n",
    "\n",
    "Here the optimizer receives only the trainable parameters. One gradient step\n",
    "then moves the head and nothing else:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-2"
   },
   "source": [
    "head_opt = torch.optim.SGD(\n",
    "    (p for p in finetune.parameters() if p.requires_grad), lr=0.1)\n",
    "before = [p.clone() for p in finetune.parameters()]\n",
    "finetune(X).sum().backward()\n",
    "head_opt.step()\n",
    "[not torch.equal(b, p) for b, p in zip(before, finetune.parameters())]"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[False,\n False,\n False,\n False,\n False,\n False,\n False,\n False,\n False,\n False,\n True,\n True]"
     },
     "metadata": {},
     "execution_count": 19
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-24 fw=pytorch -->\n",
    "\n",
    "Only the last two entries, the head's weight and bias, changed. Two pitfalls\n",
    "deserve a warning, because both fail silently."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-3"
   },
   "source": [
    "for p in net[:-1].parameters():\n",
    "    p.requires_grad = False   # frozen, but Adam's moments remain allocated\n",
    "moments = sum(t.numel() for s in adam.state.values()\n",
    "              for t in s.values() if torch.is_tensor(t) and t.ndim > 0)\n",
    "moments == 2 * n"
   ],
   "execution_count": 20,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 20
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-25 fw=pytorch -->\n",
    "\n",
    "Second, freezing governs parameters only, and batch normalization carries\n",
    "state that is not a parameter. Its running statistics are buffers, updated by\n",
    "the *forward pass* whenever the layer is in training mode. Freeze a BatchNorm\n",
    "layer's weight and bias and its running mean keeps drifting anyway:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-4"
   },
   "source": [
    "bn = nn.BatchNorm1d(4)\n",
    "for p in bn.parameters():\n",
    "    p.requires_grad = False\n",
    "before = bn.running_mean.clone()\n",
    "bn(torch.randn(8, 4) + 3)          # train mode: stats update regardless\n",
    "torch.allclose(bn.running_mean, before)"
   ],
   "execution_count": 21,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "False"
     },
     "metadata": {},
     "execution_count": 21
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-26 fw=pytorch -->\n",
    "\n",
    "`requires_grad` and `.eval()` are orthogonal switches: the first stops\n",
    "gradients, the second stops the behaviors tied to training mode, such as\n",
    "running-statistics updates and dropout. To pin a BatchNorm layer during\n",
    "fine-tuning you need both."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-27 fw=all -->\n",
    "\n",
    "Freezing whole tensors is the bluntest form of partial training.\n",
    "Parameter-efficient methods instead add small trainable low-rank corrections\n",
    "next to frozen weights; the linear algebra behind them is developed in\n",
    "that section. A related idea maintains derived,\n",
    "non-trained state: an *exponential moving average* (EMA) of the weights kept\n",
    "alongside the trained ones and used for evaluation. Whether it improves on the\n",
    "final iterate depends on its decay and the training schedule; controlled\n",
    "examples appear in that section. Like BatchNorm statistics,\n",
    "the average is state updated outside backpropagation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-28 fw=pytorch -->\n",
    "\n",
    "PyTorch ships the mechanism in `torch.optim.swa_utils`: `AveragedModel`\n",
    "clones the model once and, given the EMA averaging function, blends the live\n",
    "weights into the copy after every optimizer step; evaluation then uses the\n",
    "averaged clone. Continuing the fine-tuning above, the average trails the\n",
    "moving head:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-5"
   },
   "source": [
    "ema_model = torch.optim.swa_utils.AveragedModel(\n",
    "    finetune, multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.9))\n",
    "for _ in range(5):\n",
    "    finetune(X).sum().backward()\n",
    "    head_opt.step()\n",
    "    head_opt.zero_grad()\n",
    "    ema_model.update_parameters(finetune)\n",
    "float((ema_model.module[3].bias - finetune[3].bias).detach().abs().max())"
   ],
   "execution_count": 22,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.6190201640129089"
     },
     "metadata": {},
     "execution_count": 22
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-29 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-30 fw=pytorch -->\n",
    "\n",
    "A model's state is a named object graph of tensors. `named_parameters()` walks the\n",
    "trainable leaves; the `state_dict` adds buffers, tensors that persist and\n",
    "move with the model but receive no gradients, such as BatchNorm running\n",
    "statistics. Training with Adam in fp32 costs 16 bytes per parameter (4 weights,\n",
    "4 gradients, 8 optimizer state) before activations. Mixed-precision totals\n",
    "depend on the dtypes and copies the implementation keeps.\n",
    "Assigning one parameter to two modules ties them: one entry in\n",
    "`named_parameters()`, gradients summed over its uses. Setting\n",
    "`requires_grad = False` freezes a parameter, but reclaims no optimizer state\n",
    "already allocated and does not stop buffer updates in training mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-31 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Write a helper that reports the byte cost of fp32 Adam training separately\n",
    "   for each top-level block of `net`. Which block\n",
    "   dominates, and would that still hold if the residual blocks were 10 times\n",
    "   wider?\n",
    "1. BatchNorm's running mean and variance are non-trained state. Suppose you\n",
    "   made them trainable parameters so that the optimizer updates them. What\n",
    "   goes wrong during training, and why does gradient descent on a running\n",
    "   average not compute the same thing as the forward-pass update rule it\n",
    "   replaces?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-32 fw=pytorch -->\n",
    "\n",
    "3. Tie two layers as in `TinyLM`, then copy the model with\n",
    "   `copy.deepcopy`. Is the copy still tied? Check with `is`, and explain how a\n",
    "   copy operation can preserve aliasing.\n",
    "4. Freeze `lm.emb.weight` in the tied `TinyLM` but leave `lm.head` alone. How\n",
    "   many trainable parameters remain? Explain what this teaches about the\n",
    "   interaction between tying and freezing."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
