code/freq_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "Model Analysis: This notebook loads a pretrained model checkpoint, runs inference on new text and records activations (queries and keys) for plotting."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "1. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import pickle\n",
    "from contextlib import nullcontext\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import tiktoken\n",
    "from model import GPTConfig, GPT, CausalSelfAttention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "2. Config and pytorch setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_from = 'resume' # either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl')\n",
    "base_dir = '/data/'\n",
    "out_dir = 'final-owt-ckpts/' # ignored if init_from is not 'resume'\n",
    "ckpt_fname = 'gpt2-774M-rope-ckpt.pt' # ignored if init_from is not 'resume'\n",
    "start = \"FILE:data/shakespeare_sonnet/input.txt\" # \"\\n\" or \"<|endoftext|>\" or etc. Can also specify a file, use as: \"FILE:prompt.txt\"\n",
    "# model\n",
    "n_layer = 36\n",
    "n_head = 20\n",
    "n_embd = 1280\n",
    "dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+\n",
    "norm_type = 'rmsnorm' # 'layernorm' or 'rmsnorm'\n",
    "pos_type = 'rope' # 'sinusoidal', 'rope', 'nope', 'popev1', 'popev1b', 'popev1bg', 'popev2', 'alibi'\n",
    "base_freq = 10000 # base frequency for rotary positional encoding\n",
    "rotate_fraction = 1.0\n",
    "thetab_init = 'zero'\n",
    "bias = False # do we use bias inside LayerNorm and Linear layers?\n",
    "complex_flash = True\n",
    "# sampling\n",
    "num_samples = 1 # number of samples to draw\n",
    "max_new_tokens = 500 # number of tokens generated in each sample\n",
    "temperature = 0.8 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions\n",
    "top_k = 200 # retain only the top_k most likely tokens, clamp others to have 0 probability\n",
    "seed = 1337\n",
    "device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.\n",
    "dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16'\n",
    "compile = False # use PyTorch 2.0 to compile the model to be faster\n",
    "# exec(open('configurator.py').read()) # overrides from command line or config file\n",
    "\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul\n",
    "torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn\n",
    "device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast\n",
    "ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]\n",
    "ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "3. Load model, enable qk storage and load tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model init\n",
    "if init_from == 'resume':\n",
    "    # init from a model saved in a specific directory\n",
    "    ckpt_path = os.path.join(os.path.join(base_dir, out_dir, ckpt_fname))\n",
    "    checkpoint = torch.load(ckpt_path, map_location=device)\n",
    "    gptconf = GPTConfig(**checkpoint['model_args'])\n",
    "    model = GPT(gptconf)\n",
    "    state_dict = checkpoint['model']\n",
    "    unwanted_prefix = '_orig_mod.'\n",
    "    for k,v in list(state_dict.items()):\n",
    "        if k.startswith(unwanted_prefix):\n",
    "            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)\n",
    "    model.load_state_dict(state_dict, strict=False)\n",
    "elif init_from.startswith('gpt2'):\n",
    "    # init from a given GPT-2 model\n",
    "    model = GPT.from_pretrained(init_from, dict(dropout=0.0))\n",
    "\n",
    "model.eval()\n",
    "model.to(device)\n",
    "if compile:\n",
    "    model = torch.compile(model) # requires PyTorch 2.0 (optional)\n",
    "\n",
    "# enable flags to store qk vectors in all attn layers\n",
    "for module in model.transformer.modules():\n",
    "    if isinstance(module, CausalSelfAttention):\n",
    "        module.enable_qk_storage()\n",
    "        # module.enable_attn_scores_storage()\n",
    "\n",
    "# look for the meta pickle in case it is available in the dataset folder\n",
    "load_meta = False\n",
    "if init_from == 'resume' and 'config' in checkpoint and 'dataset' in checkpoint['config']: # older checkpoints might not have these...\n",
    "    meta_path = os.path.join('data', checkpoint['config']['dataset'], 'meta.pkl')\n",
    "    load_meta = os.path.exists(meta_path)\n",
    "if load_meta:\n",
    "    print(f\"Loading meta from {meta_path}...\")\n",
    "    with open(meta_path, 'rb') as f:\n",
    "        meta = pickle.load(f)\n",
    "    # TODO want to make this more general to arbitrary encoder/decoder schemes\n",
    "    stoi, itos = meta['stoi'], meta['itos']\n",
    "    encode = lambda s: [stoi[c] for c in s]\n",
    "    decode = lambda l: ''.join([itos[i] for i in l])\n",
    "else:\n",
    "    # ok let's assume gpt-2 encodings by default\n",
    "    print(\"No meta.pkl found, assuming GPT-2 encodings...\")\n",
    "    enc = tiktoken.get_encoding(\"gpt2\")\n",
    "    encode = lambda s: enc.encode(s, allowed_special={\"<|endoftext|>\"})\n",
    "    decode = lambda l: enc.decode(l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "4. Preprocess prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# encode the beginning of the prompt\n",
    "if start.startswith('FILE:'):\n",
    "    with open(start[5:], 'r', encoding='utf-8') as f:\n",
    "        start = f.read()\n",
    "    # This pattern looks for a Roman numeral at the beginning of a line\n",
    "    # The (?m) flag enables multi-line mode, so '^' matches the start of each line\n",
    "    pattern = re.compile(r'(?m)(?=^[IVXLCDM]+\\n)')\n",
    "    # filter(None, ...) removes any empty strings from the result\n",
    "    sonnets = list(filter(None, pattern.split(start)))\n",
    "\n",
    "start_ids = [encode(sonnet) for sonnet in sonnets]\n",
    "start_ids = [torch.tensor(s, dtype=torch.long, device=device)[None, ...] for s in start_ids]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "5. Run inference and record activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "qk_store = {'queries': [], 'keys': []}\n",
    "avg_over_n = 10 # avg q/k over this many samples\n",
    "# sampled_start_ids = random.sample(start_ids, avg_over_n)\n",
    "# run generation\n",
    "with torch.no_grad():\n",
    "    with ctx:\n",
    "        for x in start_ids[10:10+avg_over_n]:\n",
    "            y = model(x)\n",
    "            # layer-wise queries and keys\n",
    "            layer_qk_store = {'queries': [], 'keys': []}\n",
    "            for module in model.transformer.modules():\n",
    "                if isinstance(module, CausalSelfAttention):\n",
    "                    # query [1, nh, T, hs]\n",
    "                    queries, keys = module.get_stored_qk()\n",
    "                    # get 2-dim chunks for RoPE\n",
    "                    if pos_type == 'rope':\n",
    "                        bs, nh, T, hs = queries.shape  # q, k shape: [bs, nh, T, hs]\n",
    "                        queries = queries.reshape(bs, nh, T, hs // 2, 2)\n",
    "                        keys = keys.reshape(bs, nh, T, hs // 2, 2)\n",
    "                        # compute L2 norm\n",
    "                        queries = queries.norm(dim=-1, p=2)\n",
    "                        keys = keys.norm(dim=-1, p=2)\n",
    "                    elif pos_type in ['popev1', 'popev1b']:\n",
    "                        assert queries.is_complex() and keys.is_complex()\n",
    "                        queries = queries.abs()\n",
    "                        keys = keys.abs()\n",
    "                    # avg over sequence dimension\n",
    "                    queries = queries.mean(dim=-2, keepdim=False)\n",
    "                    keys = keys.mean(dim=-2, keepdim=False)\n",
    "                    layer_qk_store['queries'].append(queries)\n",
    "                    layer_qk_store['keys'].append(keys)\n",
    "            # queries across all layers [1, nl, nh, hs]\n",
    "            qk_store['queries'].append(torch.stack(layer_qk_store['queries'], dim=0))\n",
    "            qk_store['keys'].append(torch.stack(layer_qk_store['keys'], dim=0))\n",
    "# average over all samples: queries [B, nl, 1, nh, hs] -> [nl, nh, hs]\n",
    "qk_store['queries'] = torch.stack(qk_store['queries'], dim=0).squeeze().mean(dim=0, keepdim=False)\n",
    "qk_store['keys'] = torch.stack(qk_store['keys'], dim=0).squeeze().mean(dim=0, keepdim=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "6. Plot frequency usage across all attn layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_freqs_of_layers(ax, data, title):\n",
    "        # plot \n",
    "        im = ax.imshow(\n",
    "                data.to(torch.float16).cpu().numpy().transpose(), \n",
    "                aspect='auto', \n",
    "                origin='lower', \n",
    "                cmap='viridis', \n",
    "                interpolation='nearest'\n",
    "        )\n",
    "        # set title\n",
    "        ax.set_title(title, fontsize=16, pad=10)\n",
    "\n",
    "        # configure x-axis (Layers)\n",
    "        ax.set_xlabel(\"Layer\", fontsize=14, labelpad=10)\n",
    "        x_ticks = np.arange(0, n_layer)\n",
    "        x_tick_labels = [str(i + 1) for i in x_ticks]\n",
    "        ax.set_xticks(x_ticks)\n",
    "        ax.set_xticklabels(x_tick_labels, fontsize=12)\n",
    "\n",
    "        # configure y-axis (Frequencies)\n",
    "        ax.set_yticks([]) # Remove default numeric y-ticks\n",
    "        # x-coordinate set to -2.8 to place the text to the left of the heatmap.\n",
    "        y_axis_size = n_embd // (n_head * 2) if pos_type == 'rope' else n_embd // n_head\n",
    "        ax.text(-1, y_axis_size * 0.25, \"Low Frequencies\",\n",
    "                ha='center', va='center', rotation=90, fontsize=14)\n",
    "        ax.text(-1, y_axis_size * 0.75, \"High Frequencies\",\n",
    "                ha='center', va='center', rotation=90, fontsize=14)\n",
    "\n",
    "        # configure color bar\n",
    "        cbar = fig.colorbar(im, ax=ax, pad=0.02)\n",
    "        cbar.set_label(\"Mean norm\", fontsize=14, labelpad=5)\n",
    "        cbar.ax.tick_params(labelsize=12)\n",
    "        return im\n",
    "\n",
    "# queries [nl, nh, hs]\n",
    "mean_q_per_layer = qk_store['queries'].mean(dim=1)\n",
    "mean_k_per_layer = qk_store['keys'].mean(dim=1)\n",
    "# init plot\n",
    "plt.style.use('seaborn-v0_8-white')\n",
    "plt.rcParams['font.family'] = 'serif'\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7), sharey=True)\n",
    "im1 = plot_freqs_of_layers(ax1, mean_q_per_layer, f\"Query frequency usage per layer\")\n",
    "im2 = plot_freqs_of_layers(ax2, mean_k_per_layer, f\"Key frequency usage per layer\")\n",
    "\n",
    "fig.tight_layout()\n",
    "# plt.show()\n",
    "plt.savefig('final-qk-viz/all_layers/774M-rope-1.png', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "7. Plot frequency usage of specific attn layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_freqs_of_heads(ax, data, title):\n",
    "        data_to_plot = data.to(torch.float16).cpu().numpy().transpose()\n",
    "        # create heatmap\n",
    "        im = ax.imshow(\n",
    "        data_to_plot,\n",
    "        aspect='auto',\n",
    "        origin='lower',\n",
    "        cmap='viridis',\n",
    "        interpolation='nearest'\n",
    "        )\n",
    "\n",
    "        # customize axes, labels and ticks\n",
    "        ax.set_title(title, fontsize=16, pad=10)\n",
    "\n",
    "        # configure x-axis (Attention Heads)\n",
    "        ax.set_xlabel(\"Attention Head\", fontsize=14, labelpad=10)\n",
    "        # Set x-ticks to correspond to head numbers (1 to 16)\n",
    "        x_ticks = np.arange(0, n_head)\n",
    "        x_tick_labels = [str(i + 1) for i in x_ticks]\n",
    "        ax.set_xticks(x_ticks)\n",
    "        ax.set_xticklabels(x_tick_labels, fontsize=12)\n",
    "\n",
    "        # configure y-axis (Frequencies)\n",
    "        ax.set_yticks([]) # Remove default numeric y-ticks\n",
    "        # Use ax.text to place custom labels for frequency ranges\n",
    "        y_axis_size = n_embd // 2 if pos_type == 'rope' else n_embd\n",
    "        ax.text(-1, y_axis_size // n_head * 0.25, \"Low Frequencies\",\n",
    "                ha='center', va='center', rotation=90, fontsize=14)\n",
    "        ax.text(-1, y_axis_size // n_head * 0.75, \"High Frequencies\",\n",
    "                ha='center', va='center', rotation=90, fontsize=14)\n",
    "\n",
    "        # configure color bar\n",
    "        cbar = fig.colorbar(im, ax=ax, pad=0.02)\n",
    "        cbar.set_label(\"Mean norm\", fontsize=14, labelpad=5)\n",
    "        cbar.ax.tick_params(labelsize=12)\n",
    "        return im\n",
    "\n",
    "# plotting the queries and keys for the selected layer\n",
    "selected_layer = 11\n",
    "# queries [L, nh, hs]\n",
    "selected_layer_q = qk_store['queries'][selected_layer]\n",
    "selected_layer_k = qk_store['keys'][selected_layer]\n",
    "\n",
    "# init plot\n",
    "plt.style.use('seaborn-v0_8-white')\n",
    "plt.rcParams['font.family'] = 'serif'\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7), sharey=True)\n",
    "im1 = plot_freqs_of_heads(ax1, selected_layer_q, f\"Query frequency usage in Layer {selected_layer+1}\")\n",
    "im2 = plot_freqs_of_heads(ax2, selected_layer_k, f\"Key frequency usage in Layer {selected_layer+1}\")\n",
    "fig.tight_layout()\n",
    "plt.savefig(f'final-qk-viz/all_heads/124M-rope-l{selected_layer+1}-1.png', dpi=300)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}