{
 "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 jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx\n",
    "import optax"
   ],
   "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(nnx.Module):\n",
    "    def __init__(self, d, rngs):\n",
    "        self.body = nnx.Sequential(\n",
    "            nnx.Linear(d, d, rngs=rngs), nnx.relu,\n",
    "            nnx.Linear(d, d, rngs=rngs))\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return X + self.body(X)\n",
    "\n",
    "rngs = nnx.Rngs(42)\n",
    "net = nnx.Sequential(nnx.Linear(20, 64, rngs=rngs),\n",
    "                     ResidualBlock(64, rngs), ResidualBlock(64, rngs),\n",
    "                     nnx.Linear(64, 10, rngs=rngs))\n",
    "X = jax.random.normal(jax.random.key(0), (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=jax -->\n",
    "\n",
    "A model built from NNX modules is an object graph, and its parameters are\n",
    "`nnx.Param` variables owned by the layers. To reach one, follow ordinary\n",
    "attributes and container indices. `net.layers[3]` is the output layer and\n",
    "its `bias` variable is marked trainable by its type."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-2"
   },
   "source": [
    "bias = net.layers[3].bias\n",
    "type(bias), bias.shape"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(flax.nnx.variablelib.Param, (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.layers[1].body.layers[0].kernel.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=jax -->\n",
    "\n",
    "Gradients are not stored on parameters. `nnx.grad` returns a second state\n",
    "tree with the same parameter paths, one gradient leaf per parameter:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-4"
   },
   "source": [
    "params = nnx.state(net, nnx.Param)\n",
    "grads = nnx.grad(lambda model: model(X).sum())(net)\n",
    "jax.tree_util.tree_structure(grads) == jax.tree_util.tree_structure(params)"
   ],
   "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=jax -->\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",
    "`nnx.state(net, nnx.Param)` selects them. Its flat view yields each array\n",
    "together with the object-graph path that reaches it."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-5"
   },
   "source": [
    "[(path, leaf.shape) for path, leaf in params.flat_state()]"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('layers', 0, 'bias'), (64,)),\n (('layers', 0, 'kernel'), (20, 64)),\n (('layers', 1, 'body', 'layers', 0, 'bias'), (64,)),\n (('layers', 1, 'body', 'layers', 0, 'kernel'), (64, 64)),\n (('layers', 1, 'body', 'layers', 2, 'bias'), (64,)),\n (('layers', 1, 'body', 'layers', 2, 'kernel'), (64, 64)),\n (('layers', 2, 'body', 'layers', 0, 'bias'), (64,)),\n (('layers', 2, 'body', 'layers', 0, 'kernel'), (64, 64)),\n (('layers', 2, 'body', 'layers', 2, 'bias'), (64,)),\n (('layers', 2, 'body', 'layers', 2, 'kernel'), (64, 64)),\n (('layers', 3, 'bias'), (10,)),\n (('layers', 3, 'kernel'), (64, 10))]"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-7 fw=jax -->\n",
    "\n",
    "Read one of the paths closely.\n",
    "`('layers', 1, 'body', 'layers', 0, 'kernel')` means: child 1 of `net`, the\n",
    "first residual block, its `body`, child 0 of that sequence, and finally its\n",
    "kernel. Paths survive arbitrary nesting. Filters select the kinds of state a\n",
    "consumer needs, while checkpoints can save the full state\n",
    "(that section):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-accessing-parameters-6"
   },
   "source": [
    "list(params)"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "['layers']"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-8 fw=jax -->\n",
    "\n",
    "One tree, one naming scheme, and every consumer, whether optimizer,\n",
    "checkpoint, or debugger, walks it. This model has only `nnx.Param` state.\n",
    "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=jax -->\n",
    "\n",
    "NNX uses `Variable` subclasses to distinguish kinds of state. `nnx.Param`\n",
    "marks trainable state, `nnx.BatchStat` marks running statistics, and a custom\n",
    "subclass can mark another persistent buffer. Filters select which kinds an\n",
    "optimizer or checkpoint sees, so an ephemeral cache can remain outside the\n",
    "checkpoint filter or be passed explicitly. Here is a module that standardizes its inputs\n",
    "with statistics computed once from a reference sample:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-1"
   },
   "source": [
    "class Buffer(nnx.Variable):\n",
    "    pass\n",
    "\n",
    "class Whitener(nnx.Module):\n",
    "    def __init__(self, mean, std, rngs):\n",
    "        self.mean = Buffer(mean)\n",
    "        self.std = Buffer(std)\n",
    "        self.out = nnx.Linear(4, 2, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return self.out((X - self.mean) / self.std)\n",
    "\n",
    "sample = jax.random.normal(jax.random.key(1), (100, 4)) * jnp.arange(1., 5.)\n",
    "whiten = Whitener(sample.mean(0), sample.std(0), nnx.Rngs(2))\n",
    "[(path, value.shape) for path, value in nnx.state(whiten).flat_state()]"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('mean',), (4,)),\n (('out', 'bias'), (2,)),\n (('out', 'kernel'), (4, 2)),\n (('std',), (4,))]"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-11 fw=jax -->\n",
    "\n",
    "Both kinds belong to the object graph and can be checkpointed together. The\n",
    "optimizer selects `nnx.Param`, so it never sees the statistics."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-2"
   },
   "source": [
    "{'params': list(nnx.state(whiten, nnx.Param)),\n",
    " 'buffers': list(nnx.state(whiten, Buffer))}"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'params': ['out'], 'buffers': ['mean', 'std']}"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-12 fw=jax -->\n",
    "\n",
    "Wrapping an array in a `Variable` makes it state. We can see the resulting\n",
    "tree operation on the CPU by splitting the graph, casting every state leaf,\n",
    "and merging a low-precision clone. The original model remains unchanged:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-3"
   },
   "source": [
    "graph, state = nnx.split(whiten)\n",
    "low = nnx.merge(graph, jax.tree.map(lambda x: x.astype(jnp.float16), state))\n",
    "low.mean.dtype, whiten.mean.dtype"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(dtype('float16'), dtype('float32'))"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-13 fw=jax -->\n",
    "\n",
    "Device placement follows the same rule: `jax.device_put` moves a whole\n",
    "state tree, parameters and buffers alike. Updating the model with that tree\n",
    "moves all registered state. The same lines work on any host; on a machine\n",
    "with an accelerator, `jax.devices()[0]` is a GPU or TPU:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-parameters-and-buffers-4"
   },
   "source": [
    "moved = jax.device_put(nnx.state(whiten), jax.devices()[0])\n",
    "nnx.update(whiten, moved)\n",
    "whiten.mean.device"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "CudaDevice(id=0)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "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 = nnx.state(net, nnx.Param)\n",
    "leaves = jax.tree_util.tree_leaves(params)\n",
    "n = sum(x.size for x in leaves)\n",
    "weights = sum(x.size * x.dtype.itemsize for x in leaves)   # 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": [
    "optimizer = nnx.Optimizer(net, optax.adam(1e-3), wrt=nnx.Param)\n",
    "grads = nnx.grad(lambda model: model(X).sum())(net)\n",
    "optimizer.update(net, grads)\n",
    "moments = sum(x.size for x in jax.tree.leaves(optimizer.opt_state)\n",
    "              if hasattr(x, 'ndim') and x.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=jax -->\n",
    "\n",
    "A miniature version shows the mechanics. `nnx.Embed` provides `attend`, which multiplies a hidden state by the\n",
    "transpose of the embedding table, exactly the output projection. Calling it\n",
    "reuses the module's one kernel, so the tying is visible in the state itself:\n",
    "no head matrix exists."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-tied-parameters-1"
   },
   "source": [
    "class TinyLM(nnx.Module):\n",
    "    def __init__(self, vocab_size, d, tied=True, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.tied = tied\n",
    "        self.emb = nnx.Embed(vocab_size, d, rngs=rngs)\n",
    "        self.body = nnx.Linear(d, d, rngs=rngs)\n",
    "        self.head = None if tied else nnx.Linear(\n",
    "            d, vocab_size, use_bias=False, rngs=rngs)\n",
    "\n",
    "    def __call__(self, tokens):\n",
    "        h = nnx.relu(self.body(self.emb(tokens)))\n",
    "        if self.tied:\n",
    "            return self.emb.attend(h)   # the embedding table, used as the head\n",
    "        return self.head(h)\n",
    "\n",
    "lm = TinyLM(vocab_size=100, d=16)\n",
    "tokens = jax.random.randint(jax.random.key(1), (2, 8), 0, 100)\n",
    "[(path, value.shape)\n",
    " for path, value in nnx.state(lm, nnx.Param).flat_state()]"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('body', 'bias'), (16,)),\n (('body', 'kernel'), (16, 16)),\n (('emb', 'embedding'), (100, 16))]"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-18 fw=jax -->\n",
    "\n",
    "There is no aliasing bookkeeping to get right, because the state contains\n",
    "the table once. The parameter count reflects the $|V| \\times d$ saving, and\n",
    "an optimizer updates the table 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(x.size for x in jax.tree.leaves(nnx.state(lm, nnx.Param))),\n",
    " sum(x.size for x in jax.tree.leaves(nnx.state(untied, nnx.Param))))"
   ],
   "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=jax -->\n",
    "\n",
    "A checkpoint of the tied model stores the table once, because the checkpoint\n",
    "selects the model state. The flip side is that tied and untied checkpoints differ in\n",
    "structure: the tied one has no head kernel, so loading it into an untied\n",
    "module is a structural mismatch to resolve explicitly, not a silent aliasing\n",
    "decision."
   ]
  },
  {
   "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.emb.embedding[...] = lm.emb.embedding\n",
    "untied.body.kernel[...] = lm.body.kernel\n",
    "untied.body.bias[...] = lm.body.bias\n",
    "untied.head.kernel[...] = lm.emb.embedding.T\n",
    "g_tied = nnx.grad(lambda model: model(tokens).sum())(lm)\n",
    "g_untied = nnx.grad(lambda model: model(tokens).sum())(untied)\n",
    "jnp.allclose(g_tied['emb']['embedding'],\n",
    "             g_untied['emb']['embedding'] + g_untied['head']['kernel'].T)"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array(True, dtype=bool)"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "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=jax -->\n",
    "\n",
    "Every fine-tuning recipe (that section) rests on one primitive,\n",
    "and in NNX it lives in the optimizer: a leaf is frozen exactly when no update reaches it.\n",
    "`optax.multi_transform` routes each leaf to a transformation chosen by a\n",
    "label, and `optax.set_to_zero()` is the transformation that freezes, turning\n",
    "every gradient it handles into a zero update. The typical pattern freezes a\n",
    "pretrained backbone and trains only a new head. On a fresh copy of our\n",
    "residual MLP, labeling everything but the output layer `'freeze'` leaves 650\n",
    "of 18,634 parameters trainable:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-1"
   },
   "source": [
    "rngs = nnx.Rngs(3)\n",
    "finetune = nnx.Sequential(nnx.Linear(20, 64, rngs=rngs),\n",
    "                          ResidualBlock(64, rngs),\n",
    "                          ResidualBlock(64, rngs),\n",
    "                          nnx.Linear(64, 10, rngs=rngs))\n",
    "ft_params = nnx.state(finetune, nnx.Param)\n",
    "def labels(params):\n",
    "    return jax.tree_util.tree_map_with_path(\n",
    "        lambda path, _: ('train' if \"['layers'][3]\" in\n",
    "                         jax.tree_util.keystr(path) else 'freeze'), params)\n",
    "\n",
    "pure_params = nnx.to_pure_dict(ft_params)\n",
    "pure_labels = labels(pure_params)\n",
    "n_train = sum(x.size for x, l in zip(jax.tree_util.tree_leaves(pure_params),\n",
    "                                     jax.tree_util.tree_leaves(pure_labels))\n",
    "              if l == 'train')\n",
    "(n_train, sum(x.size for x in jax.tree_util.tree_leaves(pure_params)))"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(650, 18634)"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-23 fw=jax -->\n",
    "\n",
    "The optimizer receives the whole tree plus the labels, and the partition does\n",
    "the rest: Adam applies to the leaves labeled `'train'`, `set_to_zero` to the\n",
    "others. One gradient step then moves the head and nothing else:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-2"
   },
   "source": [
    "tx = optax.multi_transform(\n",
    "    {'train': optax.adam(0.1), 'freeze': optax.set_to_zero()}, labels)\n",
    "optimizer = nnx.Optimizer(finetune, tx, wrt=nnx.Param)\n",
    "before = nnx.to_pure_dict(nnx.state(finetune, nnx.Param))\n",
    "grads = nnx.grad(lambda model: model(X).sum())(finetune)\n",
    "optimizer.update(finetune, grads)\n",
    "stepped = nnx.to_pure_dict(nnx.state(finetune, nnx.Param))\n",
    "jax.tree_util.tree_map(lambda a, b: bool((a != b).any()), before, stepped)"
   ],
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'layers': {0: {'bias': False, 'kernel': False},\n  1: {'body': {'layers': {0: {'bias': False, 'kernel': False},\n     2: {'bias': False, 'kernel': False}}}},\n  2: {'body': {'layers': {0: {'bias': False, 'kernel': False},\n     2: {'bias': False, 'kernel': False}}}},\n  3: {'bias': True, 'kernel': True}}}"
     },
     "metadata": {},
     "execution_count": 18
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-24 fw=jax -->\n",
    "\n",
    "Only the two leaves under `layers[3]`, the head's kernel and bias, changed.\n",
    "Freezing traditionally carries two silent pitfalls, one about optimizer\n",
    "memory and one about batch normalization; with explicit state, both become\n",
    "questions you can settle by inspection."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-3"
   },
   "source": [
    "moments = sum(x.size for x in jax.tree_util.tree_leaves(optimizer.opt_state)\n",
    "              if hasattr(x, 'ndim') and x.ndim > 0)\n",
    "moments == 2 * n_train"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 19
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-25 fw=jax -->\n",
    "\n",
    "Second, freezing governs `nnx.Param` state only. Batch normalization keeps\n",
    "its running statistics as `nnx.BatchStat` variables and updates them during a\n",
    "training-mode forward pass. Freeze the layer's scale and bias and the\n",
    "statistics still drift because they never pass through the optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-4"
   },
   "source": [
    "bn = nnx.BatchNorm(4, use_running_average=False, momentum=0.9,\n",
    "                   rngs=nnx.Rngs(4))\n",
    "before = bn.mean[...]\n",
    "bn(jax.random.normal(jax.random.key(5), (8, 4)) + 3)\n",
    "bool(jnp.allclose(bn.mean, before))"
   ],
   "execution_count": 20,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "False"
     },
     "metadata": {},
     "execution_count": 20
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=parameters-state-memory-md-26 fw=jax -->\n",
    "\n",
    "The optimizer labels and `use_running_average` are orthogonal switches. The\n",
    "first stops parameter updates; the second stops the forward pass from\n",
    "updating running statistics. To pin a BatchNorm layer during fine-tuning you\n",
    "need both. `nnx.view(model, use_running_average=True)` changes the mode on a\n",
    "view of the full model."
   ]
  },
  {
   "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=jax -->\n",
    "\n",
    "In optax the average is one more piece of explicit state. `optax.ema` is a\n",
    "transformation like any other: feed it the weights after each step, in place\n",
    "of gradients, and its state carries an accumulator and a step count. With the\n",
    "default debiasing, `ema.update` returns the corrected average. Continuing the\n",
    "fine-tuning loop above, the average trails the moving head:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "parameters-state-memory-freezing-parameters-5"
   },
   "source": [
    "ema = optax.ema(decay=0.9)\n",
    "ema_state = ema.init(stepped)\n",
    "for _ in range(5):\n",
    "    grads = nnx.grad(lambda model: model(X).sum())(finetune)\n",
    "    optimizer.update(finetune, grads)\n",
    "    current = nnx.to_pure_dict(nnx.state(finetune, nnx.Param))\n",
    "    avg, ema_state = ema.update(current, ema_state)\n",
    "float(jnp.abs(avg['layers'][3]['bias']\n",
    "              - current['layers'][3]['bias']).max())"
   ],
   "execution_count": 21,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.17902737855911255"
     },
     "metadata": {},
     "execution_count": 21
    }
   ]
  },
  {
   "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=jax -->\n",
    "\n",
    "A model owns a typed state tree. `nnx.Param` marks trainable leaves;\n",
    "`nnx.BatchStat` and custom `Variable` subclasses hold persistent state that\n",
    "receives no optimizer updates. Filters select each view, and checkpoints can\n",
    "save the whole tree. Training with Adam in fp32 costs 16 bytes per parameter\n",
    "(4 weights, 4 gradients, 8 optimizer state) before activations.\n",
    "Mixed-precision totals depend on the dtypes and copies the implementation\n",
    "keeps. `nnx.Embed.attend` ties the embedding and the output head\n",
    "structurally: one leaf in the tree, gradients summed over its uses. Freezing\n",
    "is an optimizer-side partition, `optax.multi_transform` with `set_to_zero`;\n",
    "it allocates no state for frozen leaves, but it does not stop `nnx.BatchStat`\n",
    "updates in a training step."
   ]
  },
  {
   "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=jax -->\n",
    "\n",
    "3. Round-trip the tied model's `nnx.state(lm)` through serialization\n",
    "   (that section) and reload it. Where, if anywhere, is the\n",
    "   tying recorded? Explain why tying in NNX is a property of the module code\n",
    "   rather than of the checkpoint.\n",
    "4. In the tied `TinyLM`, try to freeze the embedding while training the head\n",
    "   using `optax.multi_transform` labels. What choices do you have, and what\n",
    "   does this teach about the interaction between tying and freezing?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
