prepare_training_with_decay.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "\n",
    "InteractiveShell.ast_node_interactivity = \"all\"\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ran into all kinds of errors with the mask\n",
    "\n",
    "- Tried various things\n",
    "- Errors were not clear as the full state was getting dumped\n",
    "- Running like this was helpful - python test_training_with_weight_decay.py 2>&1 | head -50\n",
    "- But the breakthrough came with testing with a small model in prepare_training_with_decay_simple.ipynb\n",
    "- First success was with the more complex optax.multi_transform and flax.traverse_util.path_aware_map, but was good learning\n",
    "- Eventually got the mask to work\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ValueError: Custom node type mismatch: expected type: <class 'flax.core.frozen_dict.FrozenDict'>, value: State({\n",
    "\n",
    "ValueError: Mismatch custom node data: ('embedding_dropout', 'h', 'lm_head', 'ln_f', 'wpe', 'wte') != ('h', 'lm_head', 'ln_f', 'wpe', 'wte'); value: State({\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading weights from pretrained gpt: gpt2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of prepared JAX modules dict: 89\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'h': {\n",
       "    0: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    1: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    2: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    3: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    4: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    5: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    6: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    7: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    8: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    9: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    10: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    },\n",
       "    11: {\n",
       "      'attn': {\n",
       "        'c_attn': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      },\n",
       "      'ln_1': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'ln_2': {\n",
       "        'bias': False,\n",
       "        'scale': False\n",
       "      },\n",
       "      'mlp': {\n",
       "        'c_fc': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        },\n",
       "        'c_proj': {\n",
       "          'bias': False,\n",
       "          'kernel': True\n",
       "        }\n",
       "      }\n",
       "    }\n",
       "  },\n",
       "  'lm_head': {\n",
       "    'bias': False,\n",
       "    'kernel': True\n",
       "  },\n",
       "  'ln_f': {\n",
       "    'bias': False,\n",
       "    'scale': False\n",
       "  },\n",
       "  'wpe': {\n",
       "    'embedding': False\n",
       "  },\n",
       "  'wte': {\n",
       "    'embedding': False\n",
       "  }\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from gpt2 import GPT\n",
    "import flax.nnx as nnx\n",
    "import flax\n",
    "# from flax.core import FrozenDict, frozen_dict\n",
    "\n",
    "model = GPT.from_pretrained(\"gpt2\")\n",
    "\n",
    "\n",
    "def model_state_decay_mask(model: GPT):\n",
    "    flat_state = nnx.state(model).flat_state()\n",
    "    flat_mask = {}\n",
    "    for key in flat_state.keys():\n",
    "        # The grads fed to the optimizer dont have the dropout\n",
    "        if \"dropout\" in key[0]:\n",
    "            continue\n",
    "        flat_mask[key] = key[-1] not in (\"bias\", \"embedding\", \"scale\", \"count\", \"key\")\n",
    "    return nnx.State.from_flat_path(flat_mask)\n",
    "\n",
    "\n",
    "# None of the following worked but kept getting errors like expected FrozenDict so tried this\n",
    "\n",
    "# def param_decay_mask(params: FrozenDict) -> FrozenDict:\n",
    "#     \"\"\"pytree mask for non-bias parameters\"\"\"\n",
    "#     flat_params = flax.traverse_util.flatten_dict(params)\n",
    "#     flat_param_mask = {\n",
    "#         k: k[-1] not in (\"bias\", \"embedding\", \"scale\") for k in flat_params.keys()\n",
    "#     }\n",
    "#     param_mask = flax.traverse_util.unflatten_dict(flat_param_mask)\n",
    "#     return frozen_dict.freeze(param_mask)\n",
    "\n",
    "\n",
    "model_state_decay_mask(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gpt2 import GPT\n",
    "import optax\n",
    "import flax.nnx as nnx\n",
    "\n",
    "\n",
    "@nnx.jit\n",
    "def train_step(model: GPT, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric, batch):\n",
    "    x, y = batch\n",
    "\n",
    "    def loss_fn(model: GPT):\n",
    "        logits = model(x)\n",
    "        loss = optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits.reshape([-1, logits.shape[-1]]), y.reshape([-1])\n",
    "        ).mean()\n",
    "        return loss, logits\n",
    "\n",
    "    grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)\n",
    "    (loss, logits), grads = grad_fn(model)\n",
    "    metrics.update(loss=loss, logits=logits, labels=y)\n",
    "    optimizer.update(grads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading weights from pretrained gpt: gpt2\n",
      "Length of prepared JAX modules dict: 89\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'accuracy': Array(0.33658853, dtype=float32),\n",
       " 'loss': Array(4.338888, dtype=float32)}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "{'accuracy': Array(0.32535806, dtype=float32),\n",
       " 'loss': Array(4.1847277, dtype=float32)}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from gpt2 import GPT\n",
    "import optax\n",
    "import flax.nnx as nnx\n",
    "import numpy as np\n",
    "import tiktoken\n",
    "import datasets\n",
    "\n",
    "model = GPT.from_pretrained(\"gpt2\")\n",
    "model.train()\n",
    "\n",
    "tx = optax.adamw(\n",
    "    learning_rate=1e-4, weight_decay=1e-4, mask=model_state_decay_mask(model)\n",
    ")\n",
    "optimizer = nnx.Optimizer(model, tx)\n",
    "metrics = nnx.MultiMetric(\n",
    "    accuracy=nnx.metrics.Accuracy(),\n",
    "    loss=nnx.metrics.Average(\"loss\"),\n",
    ")\n",
    "\n",
    "enc = tiktoken.get_encoding(\"gpt2\")\n",
    "batch_size = 6\n",
    "block_size = 1024\n",
    "\n",
    "data = datasets.load_dataset(path=\"Trelis/tiny-shakespeare\")\n",
    "train_data = \"\\n\".join([x[\"Text\"] for x in data[\"train\"]])\n",
    "train_data = enc.encode_ordinary(train_data)\n",
    "train_data = np.array(train_data, dtype=np.uint16)\n",
    "val_data = \"\\n\".join([x[\"Text\"] for x in data[\"test\"]])\n",
    "val_data = enc.encode_ordinary(val_data)\n",
    "val_data = np.array(val_data, dtype=np.uint16)\n",
    "\n",
    "\n",
    "def get_batch(split):\n",
    "    data = train_data if split == \"train\" else val_data\n",
    "    ix = np.random.randint(len(data) - block_size, size=(batch_size,))\n",
    "    x = np.stack([data[i : i + block_size].astype(np.int32) for i in ix])\n",
    "    y = np.stack([data[i + 1 : i + 1 + block_size].astype(np.int32) for i in ix])\n",
    "    return x, y\n",
    "\n",
    "\n",
    "train_step(model, optimizer, metrics, get_batch(\"train\"))\n",
    "metrics.compute()\n",
    "\n",
    "train_step(model, optimizer, metrics, get_batch(\"train\"))\n",
    "metrics.compute()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jax",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}