{
 "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 os\n",
    "from mxnet import autograd, gluon, np, npx\n",
    "from mxnet.gluon import nn\n",
    "npx.set_np()"
   ],
   "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.Block):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.body = nn.Sequential()\n",
    "        self.body.add(nn.Dense(d, activation='relu'), nn.Dense(d))\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X + self.body(X)\n",
    "\n",
    "np.random.seed(42)\n",
    "net = nn.Sequential()\n",
    "net.add(nn.Dense(64), ResidualBlock(64), ResidualBlock(64), nn.Dense(10))\n",
    "net.initialize()\n",
    "X = np.random.uniform(size=(2, 20))\n",
    "net(X).shape"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-3 fw=mxnet -->\n",
    "\n",
    "A model built from blocks 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",
    "a `Parameter`, an object that announces itself to the block system as state,\n",
    "bundling the tensor with its gradient buffer; `.data()` returns the tensor,\n",
    "and `.shape` reads through to it."
   ]
  },
  {
   "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": "(mxnet.gluon.parameter.Parameter, (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": "(64, 64)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-5 fw=mxnet -->\n",
    "\n",
    "The gradient is the second array a `Parameter` carries: `.data()` returns the\n",
    "value and `.grad()` its gradient buffer. gluon allocates the buffer alongside\n",
    "the data at initialization time, so before any backpropagation it simply\n",
    "holds zeros:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-4"
   },
   "source": [
    "net[3].bias.grad()"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-6 fw=mxnet -->\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",
    "`collect_params()` provides exactly that: a traversal of the whole tree that\n",
    "returns a dictionary of every parameter keyed by its path."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-5"
   },
   "source": [
    "[(name, p.shape) for name, p in net.collect_params().items()]"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[('0.weight', (64, 20)),\n ('0.bias', (64,)),\n ('1.body.0.weight', (64, 64)),\n ('1.body.0.bias', (64,)),\n ('1.body.1.weight', (64, 64)),\n ('1.body.1.bias', (64,)),\n ('2.body.0.weight', (64, 64)),\n ('2.body.0.bias', (64,)),\n ('2.body.1.weight', (64, 64)),\n ('2.body.1.bias', (64,)),\n ('3.weight', (10, 64)),\n ('3.bias', (10,))]"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-7 fw=mxnet -->\n",
    "\n",
    "Read one of the names closely. `1.body.0.weight` means: child `1` of `net`\n",
    "(the first residual block), its submodule `body`, that block's child `0`, and\n",
    "finally the leaf `weight`. Names are paths, so they survive any amount of\n",
    "nesting, and they are exactly the names `save_parameters` writes and\n",
    "`load_parameters` expects (that section). Nor is there a\n",
    "separate list of trainable parameters: each `Parameter` carries a `grad_req`\n",
    "attribute that tells autograd whether to record a gradient for it, and so far\n",
    "every entry says `'write'`:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-6"
   },
   "source": [
    "params = net.collect_params()\n",
    "[name for name, p in params.items() if p.grad_req == 'write'] == list(params)"
   ],
   "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=mxnet -->\n",
    "\n",
    "One dictionary, 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=mxnet -->\n",
    "\n",
    "gluon needs no second mechanism for such tensors: a `Constant` *is* a\n",
    "`Parameter`, one whose `grad_req` is hardwired to `'null'`. Assigning one to\n",
    "an attribute registers it like any other parameter, so it appears in\n",
    "`collect_params()`, is saved with the model, and moves with it across\n",
    "devices; but autograd never records a gradient for it and the optimizer never\n",
    "touches it. The rule of thumb: make it a `Parameter` if the optimizer should\n",
    "update it, a `Constant` if checkpoints need it, and an explicit runtime input\n",
    "or cache if its lifetime is one request. Plain Python attributes hold\n",
    "reconstructible architecture data. Here is a block that standardizes its\n",
    "inputs with 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.Block):\n",
    "    def __init__(self, mean, std):\n",
    "        super().__init__()\n",
    "        self.mean = gluon.Constant(mean)\n",
    "        self.std = gluon.Constant(std)\n",
    "        self.out = nn.Dense(2)\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.out((X - self.mean.data()) / self.std.data())\n",
    "\n",
    "sample = np.random.normal(0, 1, (100, 4)) * np.arange(1., 5.)\n",
    "whiten = Whitener(sample.mean(axis=0), sample.std(axis=0))\n",
    "whiten.initialize()\n",
    "whiten(sample[:2])   # build the deferred Dense layer\n",
    "list(whiten.collect_params())"
   ],
   "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=mxnet -->\n",
    "\n",
    "The constants sit in the same dictionary as the weights, so they are\n",
    "checkpointed alongside them. What keeps the optimizer away is the `grad_req`\n",
    "on each entry:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-2"
   },
   "source": [
    "{name: p.grad_req for name, p in whiten.collect_params().items()}"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'mean': 'null', 'std': 'null', 'out.weight': 'write', 'out.bias': 'write'}"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-12 fw=mxnet -->\n",
    "\n",
    "Registration is what makes the block aware of a tensor, and assignment only\n",
    "registers `Parameter`s (a `Constant` is one) and child blocks. A raw array\n",
    "stored as a plain attribute is invisible: not saved, and not converted when\n",
    "the model moves. We can see this on the CPU by moving the block across\n",
    "*dtypes*, which uses the same machinery as moving across devices. The\n",
    "registered constant converts; the plain attribute is left behind:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-3"
   },
   "source": [
    "whiten.note = np.zeros(4)   # plain attribute: invisible to the block\n",
    "whiten.cast('float64')\n",
    "whiten.mean.data().dtype, whiten.note.dtype"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(dtype('float64'), dtype('float32'))"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-13 fw=mxnet -->\n",
    "\n",
    "The device version of the same fact is the classic bug: a model works on the\n",
    "CPU, then crashes after `reset_device(npx.gpu(0))` because an unregistered\n",
    "tensor stayed behind. On a machine with a GPU the following confirms that\n",
    "constants move with the block:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-4"
   },
   "source": [
    "if npx.num_gpus() > 0:\n",
    "    whiten.reset_device(npx.gpu(0))\n",
    "    print('constant lives on', whiten.mean.data().device)\n",
    "else:\n",
    "    print('no GPU here; on a CUDA machine, whiten.reset_device(npx.gpu(0)) '\n",
    "          'moves whiten.mean along with the parameters')"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "constant lives on gpu(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": [
    "params = net.collect_params()\n",
    "n = sum(p.data().size for p in params.values())\n",
    "weights = sum(p.data().size * p.data().dtype.itemsize\n",
    "              for p in params.values())   # fp32: 4 bytes\n",
    "grads, adam_state = weights, 2 * weights\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": [
    "trainer = gluon.Trainer(net.collect_params(), 'adam')\n",
    "with autograd.record():\n",
    "    loss = net(X).sum()\n",
    "loss.backward()\n",
    "trainer.step(batch_size=1)   # the updater allocates the moments here\n",
    "trainer.save_states('adam.states')\n",
    "state_bytes = os.path.getsize('adam.states')\n",
    "state_bytes, 8 * n"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(151114, 149072)"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-16 fw=mxnet -->\n",
    "\n",
    "gluon's `Trainer` owns this state and offers no public traversal of it: the\n",
    "updater inside creates, at each parameter's first update, two zero arrays\n",
    "shaped like the weight. What it does offer is `save_states`, which serializes\n",
    "exactly that state, so up to a sliver of packaging the file is Adam's two\n",
    "moments: 8 bytes per parameter, two extra copies of the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-17 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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-18 fw=mxnet -->\n",
    "\n",
    "In gluon the fp32 master copy is one optimizer flag away: construct the\n",
    "`Trainer` with `optimizer_params={'multi_precision': True}` and, for every\n",
    "fp16 weight, the updater keeps an fp32 copy next to the two moments and\n",
    "applies updates to it. That is the master-weights term of the accounting\n",
    "above, 4 more bytes per parameter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-19 fw=all -->\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-20 fw=mxnet -->\n",
    "\n",
    "A miniature version shows the mechanics. gluon stores a `Dense` weight as\n",
    "(outputs, inputs), so the head's weight has shape $(|V|, d)$: the embedding\n",
    "table's shape exactly, no transpose required. Tying is `share_parameters`,\n",
    "which re-points the head's `'weight'` slot at the embedding's `Parameter`,\n",
    "matched by name; it is aliasing, not copying:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-1"
   },
   "source": [
    "class TinyLM(nn.Block):\n",
    "    def __init__(self, vocab_size, d):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(vocab_size, d)\n",
    "        self.body = nn.Dense(d, flatten=False)\n",
    "        self.head = nn.Dense(vocab_size, use_bias=False, flatten=False)\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        return self.head(npx.relu(self.body(self.emb(tokens))))\n",
    "\n",
    "lm = TinyLM(vocab_size=100, d=16)\n",
    "lm.initialize()\n",
    "lm.head.share_parameters(lm.emb.collect_params())   # one tensor, two roles\n",
    "tokens = np.random.randint(0, 100, (2, 8))\n",
    "lm(tokens)\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-21 fw=mxnet -->\n",
    "\n",
    "The two blocks now hold one `Parameter` between them. `collect_params()`\n",
    "keeps *both* names pointing at it, so a naive sum over the dictionary would\n",
    "count the table twice; deduplicating by identity shows the $|V| \\times d$\n",
    "saving. The `Trainer` deduplicates the same way internally and updates the\n",
    "table once per step:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-2"
   },
   "source": [
    "untied = TinyLM(vocab_size=100, d=16)   # no share_parameters: separate head\n",
    "untied.initialize()\n",
    "untied(tokens)\n",
    "\n",
    "def n_unique(net):\n",
    "    unique = {id(p): p for p in net.collect_params().values()}\n",
    "    return sum(p.data().size for p in unique.values())\n",
    "\n",
    "(n_unique(lm), n_unique(untied))"
   ],
   "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-22 fw=mxnet -->\n",
    "\n",
    "Keeping both names is deliberate: `save_parameters` writes the tensor under\n",
    "each name it appears under (pass `deduplicate=True` to store it once), so a\n",
    "checkpoint saved from a tied model loads into either a tied or an untied one."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-3"
   },
   "source": [
    "cp = lm.collect_params()\n",
    "(cp['emb.weight'] is cp['head.weight'], list(cp))"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(True, ['emb.weight', 'body.weight', 'body.bias', 'head.weight'])"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-23 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": [
    "up = untied.collect_params()\n",
    "for name, p in lm.collect_params().items():\n",
    "    up[name].set_data(p.data())   # same values, separate tensors\n",
    "with autograd.record():\n",
    "    lm(tokens).sum().backward()\n",
    "with autograd.record():\n",
    "    untied(tokens).sum().backward()\n",
    "bool(np.abs(lm.emb.weight.grad()\n",
    "            - (untied.emb.weight.grad()\n",
    "               + untied.head.weight.grad())).max() < 1e-6)"
   ],
   "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-24 fw=all -->\n",
    "\n",
    "## Freezing Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-25 fw=mxnet -->\n",
    "\n",
    "Every fine-tuning recipe (that section) rests on one primitive,\n",
    "and in gluon it is the switch we have already met: set a parameter's\n",
    "`grad_req` to `'null'` and backpropagation skips it, so the `Trainer` has\n",
    "nothing to apply and the weight stays put (gluon even frees the gradient\n",
    "buffer on the spot). It is the same mechanism that makes a `Constant`\n",
    "constant, applied at fine-tuning time, and `Block.setattr` flips it in bulk\n",
    "on any subtree. The typical pattern freezes a pretrained backbone and trains\n",
    "only a new head. On a fresh copy of our residual MLP, freezing everything but\n",
    "the output layer leaves 650 of 18,634 parameters trainable:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-1"
   },
   "source": [
    "finetune = nn.Sequential()\n",
    "finetune.add(nn.Dense(64), ResidualBlock(64),\n",
    "             ResidualBlock(64), nn.Dense(10))\n",
    "finetune.initialize()\n",
    "finetune(X)\n",
    "finetune[:-1].setattr('grad_req', 'null')\n",
    "\n",
    "(sum(p.data().size for p in finetune.collect_params().values()\n",
    "     if p.grad_req == 'write'),\n",
    " sum(p.data().size for p in finetune.collect_params().values()))"
   ],
   "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-26 fw=mxnet -->\n",
    "\n",
    "The `Trainer` can receive the whole dictionary: it reads each parameter's\n",
    "`grad_req` and skips the frozen ones, so the freeze lives on the model and\n",
    "every consumer respects it. One gradient step then moves the head and nothing\n",
    "else:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-2"
   },
   "source": [
    "head_trainer = gluon.Trainer(finetune.collect_params(), 'sgd',\n",
    "                             {'learning_rate': 0.1})\n",
    "before = {name: p.data().copy()\n",
    "          for name, p in finetune.collect_params().items()}\n",
    "with autograd.record():\n",
    "    loss = finetune(X).sum()\n",
    "loss.backward()\n",
    "head_trainer.step(batch_size=1)\n",
    "{name: bool((p.data() != before[name]).any())\n",
    " for name, p in finetune.collect_params().items()}"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'0.weight': False,\n '0.bias': False,\n '1.body.0.weight': False,\n '1.body.0.bias': False,\n '1.body.1.weight': False,\n '1.body.1.bias': False,\n '2.body.0.weight': False,\n '2.body.0.bias': False,\n '2.body.1.weight': False,\n '2.body.1.bias': False,\n '3.weight': True,\n '3.bias': True}"
     },
     "metadata": {},
     "execution_count": 19
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-27 fw=mxnet -->\n",
    "\n",
    "Only the two leaves under `3`, the head's weight and bias, changed. Two\n",
    "pitfalls deserve attention here: one about optimizer memory, and one about\n",
    "batch normalization."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-3"
   },
   "source": [
    "net[:-1].setattr('grad_req', 'null')   # frozen, but Adam's moments remain\n",
    "trainer.save_states('adam.states')\n",
    "os.path.getsize('adam.states') == state_bytes"
   ],
   "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-28 fw=mxnet -->\n",
    "\n",
    "Second, freezing governs gradients only, and batch normalization carries\n",
    "state that no gradient ever touches. Its running mean and variance are\n",
    "themselves parameters with `grad_req='null'` hardwired, updated by the\n",
    "*forward pass* whenever it runs in training mode, which in gluon means inside\n",
    "an `autograd.record()` scope. Freeze a `BatchNorm` layer's scale and shift\n",
    "and its running mean keeps drifting anyway:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-4"
   },
   "source": [
    "bn = nn.BatchNorm()\n",
    "bn.initialize()\n",
    "bn(np.zeros((8, 4)))            # first call builds the deferred shapes\n",
    "bn.gamma.grad_req = 'null'\n",
    "bn.beta.grad_req = 'null'\n",
    "before = bn.running_mean.data().copy()\n",
    "x = np.random.normal(0, 1, (8, 4)) + 3\n",
    "x.attach_grad()                 # a leaf for autograd: every Parameter is frozen\n",
    "with autograd.record():\n",
    "    bn(x).sum().backward()      # train mode: stats update regardless\n",
    "bool((bn.running_mean.data() == before).all())"
   ],
   "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-29 fw=mxnet -->\n",
    "\n",
    "`grad_req` and the recording scope are orthogonal switches: the first stops\n",
    "gradients, the second decides whether a forward pass runs in training mode,\n",
    "with everything that entails, batch statistics and dropout alike. To pin a\n",
    "BatchNorm layer during fine-tuning you need both: freeze its parameters *and*\n",
    "hold its statistics still, either by keeping it out of `record()` or by\n",
    "constructing it with `use_global_stats=True`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-30 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-31 fw=mxnet -->\n",
    "\n",
    "gluon has no built-in weight averaging, but with all state reachable through\n",
    "one dictionary it is a short loop: copy the trainable leaves once, then blend\n",
    "the new weights in after each step. Continuing the fine-tuning above, the\n",
    "average trails the moving head:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-5"
   },
   "source": [
    "ema = {name: p.data().copy()\n",
    "       for name, p in finetune.collect_params().items()\n",
    "       if p.grad_req == 'write'}\n",
    "for _ in range(5):\n",
    "    with autograd.record():\n",
    "        loss = finetune(X).sum()\n",
    "    loss.backward()\n",
    "    head_trainer.step(batch_size=1)\n",
    "    cp = finetune.collect_params()\n",
    "    for name in ema:   # weights in, average out\n",
    "        ema[name] = 0.9 * ema[name] + 0.1 * cp[name].data()\n",
    "float(np.abs(ema['3.bias'] - finetune[3].bias.data()).max())"
   ],
   "execution_count": 22,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.7371180653572083"
     },
     "metadata": {},
     "execution_count": 22
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-32 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-33 fw=mxnet -->\n",
    "\n",
    "A model's state is one dictionary of `Parameter`s keyed by path, returned by\n",
    "`collect_params()`, and every consumer, whether `Trainer`, checkpoint, or\n",
    "debugger, walks it. Each entry carries its own `grad_req` switch; a\n",
    "`Constant` is a `Parameter` with `grad_req='null'` hardwired, so persistent\n",
    "non-trained state and freezing are one mechanism. Training with Adam in fp32\n",
    "costs 16 bytes per parameter (4 weights, 4 gradients, 8 optimizer state)\n",
    "before activations. Mixed-precision totals depend on the dtypes and copies the\n",
    "implementation keeps (`multi_precision=True` adds an fp32 master copy).\n",
    "`share_parameters` ties two layers to one `Parameter`: two\n",
    "names in the dictionary, one tensor, gradients summed over its uses, updated\n",
    "once per step. Setting `grad_req='null'` freezes a parameter and frees its\n",
    "gradient buffer, but reclaims no optimizer state already allocated and does\n",
    "not stop running-statistics updates in training mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-34 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-35 fw=mxnet -->\n",
    "\n",
    "3. Round-trip the tied model through `save_parameters` and `load_parameters`\n",
    "   (that section), once with `deduplicate=False` and once with\n",
    "   `deduplicate=True`. Compare the file sizes, and check which of the two\n",
    "   files loads into the untied twin. Where, if anywhere, is the tying\n",
    "   recorded?\n",
    "4. In the tied `TinyLM`, set `lm.emb.weight.grad_req = 'null'`. What happened\n",
    "   to the head, and how many trainable parameters remain? Explain what this\n",
    "   teaches about the interaction between tying and freezing."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-mxnet",
   "display_name": "d2l (MXNet)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
