assets/Accurate_Math_Train.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O3OoV2Pd-w7g"
      },
      "source": [
        "# Accurate Integer Mathematics in Transformers - Train the Model\n",
        "\n",
        "This CoLab defines and trains a Transformer model that performs integer addition, subtraction and multiplication e.g. 133357+182243=+0315600, 123450-345670=-0123230 and 000345*000823=+283935. Each digit is a separate token. For 6 digit questions, the model is given 14 \"question\" (input) tokens, and must then predict the corresponding 8 \"answer\" (output) tokens.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GtVn4muC-5p1"
      },
      "source": [
        "This CoLab trains the model, storing the results to the Colab files. Useful models are manually copied to HuggingFace."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JWXJvUUb-6in"
      },
      "source": [
        "## Tips for using the Colab\n",
        " * You can run and alter the code in this CoLab notebook yourself in Google CoLab ( https://colab.research.google.com/ ).\n",
        " * To run the notebook, in Google CoLab, **you will need to** go to Runtime > Change Runtime Type and select GPU as the hardware accelerator.\n",
        " * Some graphs are interactive!\n",
        " * Use the table of contents pane in the sidebar to navigate.\n",
        " * Collapse irrelevant sections with the dropdown arrows.\n",
        " * Search the page using the search in the sidebar, not CTRL+F."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F5K2kA0L_FHc"
      },
      "source": [
        "# Part 0: Import libraries\n",
        "Imports standard libraries. Do not read."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CBOAqd0ZuhAV"
      },
      "outputs": [],
      "source": [
        "# Janky code to do different setup when run in a Colab notebook vs VSCode\n",
        "DEVELOPMENT_MODE = True\n",
        "try:\n",
        "    import google.colab\n",
        "    IN_COLAB = True\n",
        "    print(\"Running as a Colab notebook\")\n",
        "\n",
        "    !pip install kaleido\n",
        "    !pip install transformer_lens\n",
        "    !pip install circuitsvis\n",
        "    !pip install torchtyping\n",
        "    !pip install transformers\n",
        "\n",
        "except:\n",
        "    IN_COLAB = False\n",
        "    print(\"Running as a Jupyter notebook - intended for development only!\")\n",
        "    from IPython import get_ipython\n",
        "\n",
        "    ipython = get_ipython()\n",
        "    # Code to automatically update the HookedTransformer code as its edited without restarting the kernel\n",
        "    ipython.magic(\"load_ext autoreload\")\n",
        "    ipython.magic(\"autoreload 2\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "295MWK2owFNa"
      },
      "outputs": [],
      "source": [
        "# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh\n",
        "import kaleido\n",
        "import plotly.io as pio\n",
        "\n",
        "if IN_COLAB or not DEVELOPMENT_MODE:\n",
        "    pio.renderers.default = \"colab\"\n",
        "else:\n",
        "    pio.renderers.default = \"notebook_connected\"\n",
        "print(f\"Using renderer: {pio.renderers.default}\")\n",
        "\n",
        "import plotly.express as px\n",
        "import plotly.graph_objects as go"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_h1MGZnAwHZz"
      },
      "outputs": [],
      "source": [
        "pio.templates['plotly'].layout.xaxis.title.font.size = 20\n",
        "pio.templates['plotly'].layout.yaxis.title.font.size = 20\n",
        "pio.templates['plotly'].layout.title.font.size = 30"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zJpHPxqhwMeQ"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import requests\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import numpy as np\n",
        "import tqdm.auto as tqdm\n",
        "import random\n",
        "from torch.utils.data import DataLoader\n",
        "from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer\n",
        "import matplotlib.pyplot as plt\n",
        "import circuitsvis as cv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tS-PO-BPwo8f"
      },
      "outputs": [],
      "source": [
        "import transformer_lens\n",
        "import transformer_lens.utils as utils\n",
        "from transformer_lens.hook_points import (\n",
        "    HookedRootModule,\n",
        "    HookPoint,\n",
        ")  # Hooking utilities\n",
        "from transformer_lens import HookedTransformer, HookedTransformerConfig, FactoredMatrix, ActivationCache"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Part 1A: Configuration"
      ],
      "metadata": {
        "id": "Y8SVQxjtXqa9"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kRvP20dthMqN"
      },
      "outputs": [],
      "source": [
        "# Tokens used in vocab. (Token indexes 0 to 9 represent digits 0 to 9)\n",
        "PLUS_INDEX = 10\n",
        "MINUS_INDEX = 11\n",
        "EQUALS_INDEX = 12\n",
        "MULT_INDEX = 13\n",
        "DIV_INDEX = 14\n",
        "MAX_INDEX = DIV_INDEX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Z1TVBObuOKF"
      },
      "outputs": [],
      "source": [
        "# Main configuration class for main model creation and training\n",
        "class Config():\n",
        "  #@markdown Main Model\n",
        "  n_layers: int = 3 #@param\n",
        "  n_heads: int = 4 #@param\n",
        "\n",
        "  d_vocab: int = MAX_INDEX+1\n",
        "  d_model: int = 510\n",
        "  d_mlp: int = 4 * d_model\n",
        "  d_head: int = 170\n",
        "  training_seed: int = 372001\n",
        "\n",
        "  #@markdown Data\n",
        "  n_digits: int = 6 #@param\n",
        "  n_ctx: int = 3 * n_digits + 4\n",
        "  act_fn: str = 'relu'\n",
        "  batch_size: int = 64\n",
        "\n",
        "  #@markdown Optimizer\n",
        "  n_training_steps: int = 40000 #@param\n",
        "  lr: float = 0.00008\n",
        "  weight_decay: int = 0.1\n",
        "\n",
        "  #@markdown Maths Operations\n",
        "  # Percent of questions that are multiplication, subtraction and addition.\n",
        "  perc_mult: int = 0\n",
        "  perc_sub: int = 80 #@param\n",
        "  def perc_add(self):\n",
        "    return max(0, 100 - self.perc_mult - self.perc_sub)\n",
        "\n",
        "  #@markdown Insert Model\n",
        "  insert_mode: int = 0 # 0=None 1=Init, 2=FreezeHeads 3=FreezeAll\n",
        "  insert_late: bool = False\n",
        "\n",
        "  insert_n_layers: int = 2 #@param\n",
        "  insert_n_heads: int = 3 #@param\n",
        "  insert_training_seed: int = 372001\n",
        "  insert_n_training_steps: int = 15000 #@param\n",
        "\n",
        "  # Save graphs to CoLab temp files as PDF and HTML. Can manually export files for re-use in papers.\n",
        "  save_graph_to_file: bool = True\n",
        "\n",
        "  def to_dict(self):\n",
        "    return {\n",
        "      \"n_layers\": self.n_layers,\n",
        "      \"n_heads\": self.n_heads,\n",
        "      \"d_vocab\": self.d_vocab,\n",
        "      \"d_mlp\": self.d_mlp,\n",
        "      \"d_head\": self.d_head,\n",
        "      \"training_seed\": self.training_seed,\n",
        "      \"n_digits\": self.n_digits,\n",
        "      \"n_ctx\": self.n_ctx,\n",
        "      \"act_fn\": self.act_fn,\n",
        "      \"batch_size\": self.batch_size,\n",
        "      \"n_training_steps\": self.n_training_steps,\n",
        "      \"lr\": self.lr,\n",
        "      \"weight_decay\": self.weight_decay,\n",
        "      \"perc_mult\": self.perc_mult,\n",
        "      \"perc_sub\": self.perc_sub,\n",
        "      \"insert_late\": self.insert_late,\n",
        "      \"insert_mode\": self.insert_mode,\n",
        "      \"insert_n_layers\": self.insert_n_layers,\n",
        "      \"insert_n_heads\": self.insert_n_heads,\n",
        "      \"insert_training_seed\": self.insert_training_seed,\n",
        "      \"insert_n_training_steps\": self.insert_n_training_steps,\n",
        "    }\n",
        "\n",
        "\n",
        "cfg = Config()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Rl41loM_Apt"
      },
      "source": [
        "# Part 1B: Configuration\n",
        "\n",
        "\n",
        "This CoLab can be configured to:\n",
        "- Train a model using traditional approach. For example, add_d6_l2_h3_train15K_seed372001.pth is trained from scratch using 100% addition questions to give an \"Addition\" model with a very low loss (9e-9).\n",
        "- Train a \"mixed\" model to do two tasks by inserting a \"known good\" model into the untrained composite model. For example, initialising  ins1_mix_d6_l3_h4_train20K_seed372001.pth with  add_d6_l2_h3_dm510_dh170_train15K_seed372001, and then training it on 80% subtraction questions and 20% addition questions.\n",
        "\n",
        "If we are inserting existing model weightings they are loaded from HuggingFace."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mdad0Np8kSC9"
      },
      "outputs": [],
      "source": [
        "# Which model do we want to train?\n",
        "#model_name = \"add_d5_l2_h3_train15K\"  # 5 digit addition model\n",
        "#model_name = \"add_d6_l2_h3_train15K\"  # 6 digit addition model\n",
        "#model_name = \"sub_d6_l2_h3_train30K\"  # 6 digit subtraction model\n",
        "#model_name = \"mix_d6_l3_h4_train40K\"  # 6 digit addition and subtraction model\n",
        "#model_name = \"ins1_mix_d6_l3_h4_train40K\"  # 6 digit addition / subtraction model. Initialise with addition model.\n",
        "#model_name = \"ins2_mix_d6_l4_h4_train40K\"  # 6 digit addition / subtraction model. Initialised with addition model. Reset useful heads every 100 epochs\n",
        "#model_name = \"ins3_mix_d6_l4_h3_train40K\"  # 6 digit addition / subtraction model. Initialised with addition model. Reset useful heads & MLPs every 100 epochs\n",
        "model_name = \"ins3_mix_d6_l4_h3_train60K\"  # 6 digit addition / subtraction model. Initialised with addition model. Reset useful heads & MLPs every 100 epochs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2g1bk4bekm5g"
      },
      "outputs": [],
      "source": [
        "if model_name == \"add_d5_l2_h3_train15K\" :\n",
        "  cfg.n_digits = 5\n",
        "  cfg.n_layers = 2\n",
        "  cfg.n_heads = 3\n",
        "  cfg.n_training_steps = 15000\n",
        "  cfg.perc_sub = 0\n",
        "  cfg.insert_mode = 0\n",
        "\n",
        "if model_name == \"add_d6_l2_h3_train15K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 2\n",
        "  cfg.n_heads = 3\n",
        "  cfg.n_training_steps = 15000\n",
        "  cfg.perc_sub = 0\n",
        "  cfg.insert_mode = 0\n",
        "\n",
        "if model_name == \"sub_d6_l2_h3_train30K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 2\n",
        "  cfg.n_heads = 3\n",
        "  cfg.n_training_steps = 30000\n",
        "  cfg.perc_sub = 100\n",
        "  cfg.insert_mode = 0\n",
        "\n",
        "if model_name == \"mix_d6_l3_h4_train40K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 3\n",
        "  cfg.n_heads = 4\n",
        "  cfg.n_training_steps = 40000\n",
        "  cfg.perc_sub = 66 # Train on 66% subtraction and 33% addition question batches\n",
        "  cfg.insert_mode = 0\n",
        "\n",
        "if model_name == \"ins1_mix_d6_l3_h4_train40K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 3\n",
        "  cfg.n_heads = 4\n",
        "  cfg.n_training_steps = 40000\n",
        "  cfg.perc_sub = 80 # Train on 80% subtraction and 20% addition question batches\n",
        "  cfg.insert_mode = 1 # Initialise with add_d6_l2_h3_train15K.pth.\n",
        "\n",
        "if model_name == \"ins2_mix_d6_l4_h4_train40K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 4\n",
        "  cfg.n_heads = 4\n",
        "  cfg.n_training_steps = 40000\n",
        "  cfg.perc_sub = 80 # Train on 80% subtraction and 20% addition question batches\n",
        "  cfg.insert_mode = 2 # Initialise with add_d6_l2_h3_train15K.pth. Train & reset useful heads every 100 epochs\n",
        "\n",
        "if model_name == \"ins3_mix_d6_l4_h3_train40K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 4\n",
        "  cfg.n_heads = 3\n",
        "  cfg.n_training_steps = 40000\n",
        "  cfg.perc_sub = 80 # Train on 80% subtraction and 20% addition question batches\n",
        "  cfg.insert_mode = 3 # Initialise with add_d6_l2_h3_train15K.pth. Trained & reset useful heads & MLPs every 100 epochs\n",
        "\n",
        "if model_name == \"ins3_mix_d6_l4_h3_train60K\" :\n",
        "  cfg.n_digits = 6\n",
        "  cfg.n_layers = 4\n",
        "  cfg.n_heads = 3\n",
        "  cfg.n_training_steps = 60000\n",
        "  cfg.perc_sub = 80 # Train on 80% subtraction and 20% addition question batches\n",
        "  cfg.insert_mode = 3 # Initialise with add_d6_l2_h3_train15K.pth. Trained & reset useful heads & MLPs every 100 epochs\n",
        "\n",
        "cfg.n_ctx = 3 * cfg.n_digits + 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HKgiXwaFiBj4"
      },
      "outputs": [],
      "source": [
        "def file_name_suffix(digits, layers, heads, training_steps, seed):\n",
        "  train_str = str(training_steps//1000) + \"K\"\n",
        "  return '_d{}_l{}_h{}_train{}_seed{}'.format(digits, layers, heads, train_str, seed)\n",
        "\n",
        "main_fname = '' if cfg.insert_mode == 0 else 'ins{}_'.format(cfg.insert_mode)\n",
        "main_fname += 'mul' if cfg.perc_mult == 100 else 'sub' if cfg.perc_sub == 100 else 'add' if cfg.perc_add() == 100 else 'mix'\n",
        "main_fname += file_name_suffix(cfg.n_digits, cfg.n_layers, cfg.n_heads, cfg.n_training_steps, cfg.training_seed)\n",
        "main_fname_pth = main_fname + '.pth'\n",
        "main_fname_json = main_fname + '.json'\n",
        "\n",
        "\n",
        "def print_config():\n",
        "  print(\"%Mult=\", cfg.perc_mult, \"%Sub=\", cfg.perc_sub, \"%Add=\", cfg.perc_add(), \"File=\", main_fname)\n",
        "\n",
        "print_config()\n",
        "print('Main model will save to Colab temporary file', main_fname_pth)\n",
        "print('Main model config etc will save to Colab temporary file', main_fname_json)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3tExv4rk_L0C"
      },
      "source": [
        "# Part 3: Create main_model\n",
        "This section defines the token embedding / unembedding and creates the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CqGmiutQwp22"
      },
      "outputs": [],
      "source": [
        "# Embedding / Unembedding\n",
        "\n",
        "def token_to_char(i):\n",
        "  if i < 10:\n",
        "   return str(i)\n",
        "  if i == PLUS_INDEX:\n",
        "    return \"+\"\n",
        "  if i == MINUS_INDEX:\n",
        "    return \"-\"\n",
        "  if i == EQUALS_INDEX:\n",
        "    return \"=\"\n",
        "  if i == MULT_INDEX:\n",
        "    return \"*\"\n",
        "  if i == DIV_INDEX:\n",
        "    return \"\\\\\"\n",
        "  return \"?\"\n",
        "\n",
        "def tokens_to_string(tokens):\n",
        "    tokens = utils.to_numpy(tokens)\n",
        "    return \"\".join([token_to_char(i) for i in tokens[:cfg.n_ctx]])\n",
        "\n",
        "def string_to_tokens(string, batch: bool=False):\n",
        "    lookup = {str(i):i for i in range(10)}\n",
        "    lookup['+']=PLUS_INDEX\n",
        "    lookup['-']=MINUS_INDEX\n",
        "    lookup['=']=EQUALS_INDEX\n",
        "    lookup['*']=MULT_INDEX\n",
        "    lookup['\\\\']=DIV_INDEX\n",
        "\n",
        "    tokens = [lookup[i] for i in string if i not in '\\n ']\n",
        "    if batch:\n",
        "        return torch.tensor(tokens)[None, :]\n",
        "    else:\n",
        "        return torch.tensor(tokens)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dg6uwYzew4u5"
      },
      "outputs": [],
      "source": [
        "# Transformer creation\n",
        "\n",
        "# Structure is documented at https://neelnanda-io.github.io/TransformerLens/transformer_lens.html#transformer_lens.HookedTransformerConfig.HookedTransformerConfig\n",
        "ht_cfg = HookedTransformerConfig(\n",
        "    n_layers = cfg.n_layers,\n",
        "    n_heads = cfg.n_heads,\n",
        "    d_model = cfg.d_model,\n",
        "    d_head = cfg.d_head,\n",
        "    d_mlp = cfg.d_mlp,\n",
        "    act_fn = cfg.act_fn,\n",
        "    normalization_type = 'LN',\n",
        "    d_vocab = cfg.d_vocab,\n",
        "    d_vocab_out = cfg.d_vocab,\n",
        "    n_ctx = cfg.n_ctx,\n",
        "    init_weights = True,\n",
        "    device = \"cuda\",\n",
        "    seed = cfg.training_seed,\n",
        ")\n",
        "\n",
        "main_model = HookedTransformer(ht_cfg)\n",
        "\n",
        "optimizer = optim.AdamW(main_model.parameters(),\n",
        "                        lr = cfg.lr,\n",
        "                        weight_decay = cfg.weight_decay,\n",
        "                        betas = (0.9, 0.98))\n",
        "\n",
        "max_iter = cfg.n_training_steps\n",
        "warmup_iter = max_iter // 5\n",
        "scheduler1 = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, total_iters=int(warmup_iter))\n",
        "scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=int(np.ceil((max_iter-warmup_iter))))\n",
        "scheduler  = torch.optim.lr_scheduler.SequentialLR(optimizer, schedulers=[scheduler1, scheduler2], milestones=[int(warmup_iter)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pg267eav_QSl"
      },
      "source": [
        "# Part 4: Loss Function & Data Generator\n",
        "This section defines the loss function and the training/testing data generator.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OIVxDpq0w7WY"
      },
      "outputs": [],
      "source": [
        "# Loss functions\n",
        "\n",
        "# Calculate the per-token probability by comparing a batch of prediction \"logits\" to answer \"tokens\"\n",
        "def logits_to_tokens_loss(logits, tokens):\n",
        "  # Addition answer can have one extra digit than question. Answer also has a +/- sign\n",
        "  n_answer_digits = cfg.n_digits+2\n",
        "\n",
        "  # The addition answer digit token probabilities\n",
        "  ans_logits = logits[:, -(n_answer_digits+1):-1]\n",
        "\n",
        "  # Convert raw score (logits) vector into a probability distribution.\n",
        "  # Emphasize the largest scores and suppress the smaller ones, to make them more distinguishable.\n",
        "  ans_probs = F.log_softmax(ans_logits.to(torch.float64), dim=-1)\n",
        "\n",
        "  max_prob_tokens = torch.argmax(ans_probs, dim=-1)\n",
        "\n",
        "  # The addition answer digit tokens\n",
        "  ans_tokens = tokens[:, -(n_answer_digits):]\n",
        "\n",
        "  # Extract values from the ans_probs tensor, based on indices from the ans_tokens tensor\n",
        "  ans_loss = torch.gather(ans_probs, -1, ans_tokens[:, :, None])[..., 0]\n",
        "\n",
        "  return ans_loss, max_prob_tokens\n",
        "\n",
        "\n",
        "# Calculate loss as negative of average per-token mean probability\n",
        "def loss_fn(ans_loss):\n",
        "  return -ans_loss.mean(0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "04JTfmUOliHd"
      },
      "outputs": [],
      "source": [
        "# Generate an enriched data batch for one operator type\n",
        "# \"Addition\" batch entries are formated XXXXX+YYYYY=+ZZZZZZ e.g. 550030+800020=+1350050\n",
        "# \"Subtraction\" batch entries are formated XXXXX-YYYYY=-ZZZZZZ e.g. 550030-800020=-0249990, 800020-550030=+0249990\n",
        "# \"Multiplication\" batch entries are formated 000XXX*000YYY=+ZZZZZZ e.g. 000345*000678=+233910\n",
        "def data_generator_core( batch_op ):\n",
        "\n",
        "  batch = torch.zeros((cfg.batch_size, cfg.n_ctx)).to(torch.int64)\n",
        "  x = torch.randint(0, 10, (cfg.batch_size, cfg.n_digits))\n",
        "  y = torch.randint(0, 10, (cfg.batch_size, cfg.n_digits))\n",
        "\n",
        "  if batch_op == MULT_INDEX:\n",
        "    # Convert from NNNNNN*NNNNNN= to 000NNN*000NNN= so answer (product) is NNNNNN\n",
        "    num_zeros = cfg.n_digits // 2\n",
        "    for z in range(num_zeros):\n",
        "      x[:, z] = 0\n",
        "      y[:, z] = 0\n",
        "\n",
        "  # Enrich the question data on 60% of batches to speed up training\n",
        "  random_case = random.randint(1, 5)\n",
        "  if ( batch_op == PLUS_INDEX or batch_op == MINUS_INDEX ) and (random_case <= 3):\n",
        "    # Flatten x and y to 1D tensors\n",
        "    x_flat = x.view(-1)\n",
        "    y_flat = y.view(-1)\n",
        "\n",
        "    num_elements_to_modify = int(0.40 * x.numel()) # 40%\n",
        "    indices_to_modify = torch.randperm(x_flat.numel())[:num_elements_to_modify]\n",
        "    if batch_op == PLUS_INDEX :\n",
        "      # The S2, S3, etc tasks are compound and increasing rare and so harder to learn.\n",
        "      # Increase the MakeSum9 case frequency to increase frequency of S2, S3, etc\n",
        "      if random.randint(1, 2) == 1:\n",
        "        x_flat[indices_to_modify] = 9 - y_flat[indices_to_modify]\n",
        "      else:\n",
        "        y_flat[indices_to_modify] = 9 - x_flat[indices_to_modify]\n",
        "    else:\n",
        "      # Generate more questions with negative answers (NG task) by increasing the y value\n",
        "      if random_case == 1:\n",
        "        y_flat[y_flat < 9] += 1\n",
        "\n",
        "      # The M2, M3, etc tasks are compound and increasing rare and so harder to learn.\n",
        "      # Increase the DiffZero case frequency to increase frequency of M2, M3, etc\n",
        "      if random.randint(1, 2) == 1:\n",
        "        x_flat[indices_to_modify] = y_flat[indices_to_modify]\n",
        "      else:\n",
        "        y_flat[indices_to_modify] = x_flat[indices_to_modify]\n",
        "\n",
        "\n",
        "    # Reshape x and y back to its original shape\n",
        "    x = x_flat.view(x.shape)\n",
        "    y = y_flat.view(x.shape)\n",
        "\n",
        "\n",
        "  batch[:, :cfg.n_digits] = x\n",
        "  batch[:, cfg.n_digits] = batch_op\n",
        "  batch[:, 1+cfg.n_digits:1+cfg.n_digits*2] = y\n",
        "  batch[:, 1+cfg.n_digits*2] = EQUALS_INDEX\n",
        "\n",
        "  # Convert each row into a 5-digit number\n",
        "  x_values = x[:, 0]\n",
        "  y_values = y[:, 0]\n",
        "  for dn in range(1,cfg.n_digits):\n",
        "    x_values = x_values * 10 + x[:, dn]\n",
        "    y_values = y_values * 10 + y[:, dn]\n",
        "\n",
        "  # Elementwise operations to give the 1D tensor answers\n",
        "  if batch_op == MULT_INDEX:\n",
        "    answers = x_values * y_values\n",
        "  else:\n",
        "    if batch_op == MINUS_INDEX:\n",
        "      answers = x_values - y_values\n",
        "    else:\n",
        "      answers = x_values + y_values\n",
        "\n",
        "  # Insert the answers into the batch\n",
        "  for i in range(cfg.batch_size):\n",
        "    answer = answers[i]\n",
        "\n",
        "    sign = PLUS_INDEX\n",
        "    if answer < 0:\n",
        "      sign = MINUS_INDEX\n",
        "      answer = - answer\n",
        "\n",
        "    batch[i, 2+cfg.n_digits*2] = sign\n",
        "    for j in range(cfg.n_digits+1):\n",
        "      batch[i, cfg.n_ctx-j-1] = answer % 10\n",
        "      answer = answer // 10\n",
        "      if answer == 0:\n",
        "          break\n",
        "\n",
        "  return batch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dps1g8spw-U4"
      },
      "outputs": [],
      "source": [
        "# Define \"iterator\" data generator function. Invoked using next().\n",
        "def data_generator( ):\n",
        "  torch.manual_seed(cfg.training_seed)\n",
        "  while True:\n",
        "\n",
        "    batch_rand = random.randint(1, 100)\n",
        "    batch_op = MULT_INDEX if batch_rand <= cfg.perc_mult else MINUS_INDEX if batch_rand <= cfg.perc_mult + cfg.perc_sub else PLUS_INDEX\n",
        "\n",
        "    batch = data_generator_core( batch_op )\n",
        "\n",
        "    yield batch.cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IH_rXfA2xAIG"
      },
      "outputs": [],
      "source": [
        "# Initialise the data generator\n",
        "ds = data_generator()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V6c62sGVxVhM"
      },
      "outputs": [],
      "source": [
        "# Test data generator\n",
        "tokens = next(ds)\n",
        "print(tokens[:3,:])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NvijaDhSjn0B"
      },
      "source": [
        "# Part 5: Read insert_model from HuggingFace (optional)\n",
        "\n",
        "If we are initialising the untrained model with an existing model,\n",
        "then we load the existing model from HuggingFace.\n",
        "We load both the model weights and a json file stating which nodes in the model are actually doing useful calculations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ioCE99sOl-Gt"
      },
      "outputs": [],
      "source": [
        "insert_base_name = \"\"\n",
        "insert_weights_fname = \"\"\n",
        "insert_nodes_fname = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LonIqn6KjpTe"
      },
      "outputs": [],
      "source": [
        "# Read insert_model weights from HuggingFace\n",
        "if cfg.insert_mode >= 1:\n",
        "  insert_base_name = \"add\" + file_name_suffix(cfg.n_digits, cfg.insert_n_layers, cfg.insert_n_heads, cfg.insert_n_training_steps, cfg.insert_training_seed)\n",
        "\n",
        "  insert_weights_fname = insert_base_name + \".pth\"\n",
        "\n",
        "  ht_cfg = HookedTransformerConfig(\n",
        "      n_layers = cfg.insert_n_layers,\n",
        "      n_heads = cfg.insert_n_heads,\n",
        "      d_model = cfg.d_model, # Assume constant\n",
        "      d_head = cfg.d_head, # Assume constant\n",
        "      d_mlp = cfg.d_mlp, # Assume constant\n",
        "      act_fn = cfg.act_fn, # Assume constant\n",
        "      normalization_type = 'LN',\n",
        "      d_vocab = cfg.d_vocab, # Assume constant\n",
        "      d_vocab_out = cfg.d_vocab, # Assume constant\n",
        "      n_ctx = cfg.n_ctx, # Assume constant\n",
        "      init_weights = True, # Assume constant\n",
        "      device = \"cuda\",\n",
        "      seed = cfg.insert_training_seed,\n",
        "  )\n",
        "\n",
        "  insert_model = HookedTransformer(ht_cfg)\n",
        "\n",
        "  print('Loading insertion model from', insert_weights_fname)\n",
        "  insert_model.load_state_dict(utils.download_file_from_hf(repo_name=\"PhilipQuirke/Accurate6DigitSubtraction\", file_name=insert_weights_fname, force_is_torch=True))\n",
        "  insert_model.eval()\n",
        "\n",
        "  print(\"Loaded insert model\", insert_weights_fname)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qDaxauYJoorT"
      },
      "outputs": [],
      "source": [
        "class UsefulCell():\n",
        "  # Position.Layer.Head of the cell\n",
        "  position: int  # token-position. Zero to cfg.n_ctx - 1\n",
        "  layer: int\n",
        "  head: int\n",
        "\n",
        "  # Tags related to the cell of form \"MajorVersion.MinorVersion\"\n",
        "  tags: list\n",
        "\n",
        "\n",
        "  # Is this cell an attention head? If not, it must be an MLP layer\n",
        "  def is_head(self):\n",
        "    return self.head != cfg.insert_n_heads\n",
        "\n",
        "\n",
        "  def __init__(self, position, layer, head, tags):\n",
        "    self.position = position\n",
        "    self.layer = layer\n",
        "    self.head = head\n",
        "    self.tags = tags\n",
        "\n",
        "\n",
        "  @classmethod\n",
        "  def from_dict(cls, data):\n",
        "      return cls(data['position'], data['layer'], data['head'], data['tags'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HXw7-a1E025H"
      },
      "outputs": [],
      "source": [
        "if cfg.insert_mode >= 1:\n",
        "  # Read insert_model useful node information from HuggingFace\n",
        "  huggingface_directory_url = 'https://huggingface.co/PhilipQuirke/Accurate6DigitSubtraction/raw/main/'\n",
        "  insert_nodes_fname = huggingface_directory_url + insert_base_name + '_tags.json'\n",
        "  print(insert_nodes_fname)\n",
        "\n",
        "  # Download the file\n",
        "  response = requests.get(insert_nodes_fname)\n",
        "\n",
        "  # Ensure the request was successful\n",
        "  if response.status_code == 200:\n",
        "      # Load the JSON data\n",
        "      json_data = json.loads(response.content.decode('utf-8'))\n",
        "\n",
        "      useful_cells = [UsefulCell.from_dict(item) for item in json_data]\n",
        "\n",
        "      print( \"Loaded:\", len(useful_cells), \"Sample:\", useful_cells[0].tags)\n",
        "  else:\n",
        "      print( \"Failed to download the file:\", response.status_code)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1fW5galHlQqA"
      },
      "source": [
        "# Part 6B: Transfer all of insert_model into main_model (optional)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iKe3LUst8veJ"
      },
      "outputs": [],
      "source": [
        "# Transfer all attention heads weights from the small to the main model, updating the right-most small.n_heads of main_model\n",
        "def transfer_all_heads(small_model, small_cfg, start_layer, end_layer, large_model):\n",
        "  for small_layer_no, large_layer_no in enumerate(range(start_layer, end_layer+1)):\n",
        "    large_model.blocks[large_layer_no].attn.W_Q.data[:small_cfg[\"n_heads\"]] = small_model.blocks[small_layer_no].attn.W_Q.clone().data\n",
        "    large_model.blocks[large_layer_no].attn.W_K.data[:small_cfg[\"n_heads\"]] = small_model.blocks[small_layer_no].attn.W_K.clone().data\n",
        "    large_model.blocks[large_layer_no].attn.W_V.data[:small_cfg[\"n_heads\"]] = small_model.blocks[small_layer_no].attn.W_V.clone().data\n",
        "\n",
        "    large_model.blocks[large_layer_no].attn.b_Q.data[:small_cfg[\"n_heads\"]] = small_model.blocks[small_layer_no].attn.b_Q.clone().data\n",
        "    large_model.blocks[large_layer_no].attn.b_K.data[:small_cfg[\"n_heads\"]] = small_model.blocks[small_layer_no].attn.b_K.clone().data\n",
        "    large_model.blocks[large_layer_no].attn.b_V.data[:small_cfg[\"n_heads\"]] = small_model.blocks[small_layer_no].attn.b_V.clone().data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DcaxqmJC8vsa"
      },
      "outputs": [],
      "source": [
        "# Transfer all MLP layer weights from the small to the main model, updating the right-most small.d_mlp of main_model\n",
        "def transfer_all_mlps(small_model, small_cfg, start_layer, end_layer, large_model):\n",
        "  for small_layer_no, large_layer_no in enumerate(range(start_layer, end_layer+1)):\n",
        "    large_model.blocks[large_layer_no].mlp.W_in.data[:, :small_cfg[\"d_mlp\"]] = small_model.blocks[small_layer_no].mlp.W_in.clone().data\n",
        "    large_model.blocks[large_layer_no].mlp.b_in.data[:small_cfg[\"d_mlp\"]] = small_model.blocks[small_layer_no].mlp.b_in.clone().data\n",
        "    large_model.blocks[large_layer_no].mlp.W_out.data[:small_cfg[\"d_mlp\"],] = small_model.blocks[small_layer_no].mlp.W_out.clone().data\n",
        "    large_model.blocks[large_layer_no].mlp.b_out.data = small_model.blocks[small_layer_no].mlp.b_out.clone().data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-UsGFC-b9oVC"
      },
      "outputs": [],
      "source": [
        "def transfer_all_ln(small_model, start_layer, end_layer, large_model):\n",
        "  for small_layer_no, large_layer_no in enumerate(range(start_layer, end_layer+1)):\n",
        "    large_model.blocks[large_layer_no].ln1.w.data = small_model.blocks[small_layer_no].ln1.w.clone().data\n",
        "    large_model.blocks[large_layer_no].ln1.b.data = small_model.blocks[small_layer_no].ln1.b.clone().data\n",
        "\n",
        "  large_model.ln_final.w.data = small_model.ln_final.w.clone().data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BYpItGeX9-jq"
      },
      "outputs": [],
      "source": [
        "def transfer_all_embeds(small_model, large_model):\n",
        "  large_model.embed.W_E.data = small_model.embed.W_E.clone().data\n",
        "  large_model.pos_embed.W_pos.data = small_model.pos_embed.W_pos.clone().data\n",
        "  large_model.unembed.W_U.data = small_model.unembed.W_U.clone().data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WnkU3fYKlQAk"
      },
      "outputs": [],
      "source": [
        "small_cfg = {}\n",
        "large_cfg = {}\n",
        "\n",
        "# Insert the small model weights into the large model\n",
        "def transfer_full_model(small_model, large_model, start_layer, end_layer, transfer_ln=True, transfer_embeds=True):\n",
        "  \"\"\"Args:\n",
        "  small_model: The model to transfer weights from\n",
        "  large_model: The model to transfer weights to\n",
        "  start_layer: The first layer to transfer weights to\n",
        "  end_layer: The last layer to transfer weights to (Note that this is end-inclusive!)\n",
        "  \"\"\"\n",
        "  global small_cfg\n",
        "  global large_cfg\n",
        "\n",
        "  small_cfg = {k: v for k,v in small_model.cfg.__dict__.items() if k in [\"d_head\", \"d_mlp\", \"d_model\", \"n_heads\", \"n_layers\"]}\n",
        "  large_cfg = {k: v for k,v in large_model.cfg.__dict__.items() if k in [\"d_head\", \"d_mlp\", \"d_model\", \"n_heads\", \"n_layers\"]}\n",
        "\n",
        "  # Sanity checks for large model > small model\n",
        "  assert small_cfg[\"d_model\"] == large_cfg[\"d_model\"]\n",
        "  assert small_cfg[\"d_head\"] == large_cfg[\"d_head\"]\n",
        "  assert small_cfg[\"n_layers\"] <= large_cfg[\"n_layers\"]\n",
        "  assert small_cfg[\"n_heads\"] <= large_cfg[\"n_heads\"]\n",
        "  assert small_cfg[\"d_mlp\"] <= large_cfg[\"d_mlp\"]\n",
        "\n",
        "  assert 0 <= start_layer < end_layer <= large_cfg[\"n_layers\"] # Make sure start_layer and end_layer are valid\n",
        "  assert end_layer - start_layer + 1 == small_cfg[\"n_layers\"] # Make sure the number of layers to transfer is correct\n",
        "\n",
        "  transfer_all_heads(small_model, small_cfg, start_layer, end_layer, large_model)\n",
        "  transfer_all_mlps(small_model, small_cfg, start_layer, end_layer, large_model)\n",
        "  if transfer_ln:\n",
        "    transfer_all_ln(small_model, start_layer, end_layer, large_model)\n",
        "  if transfer_embeds:\n",
        "    transfer_all_embeds(small_model, large_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Rd7d5Gjmrd7"
      },
      "outputs": [],
      "source": [
        "def insert_existing_model( first_time ):\n",
        "  if cfg.insert_mode >= 1 :\n",
        "    # Is the destination the first few or last few layers of the main_model?\n",
        "    start_layer = cfg.n_layers - cfg.insert_n_layers if cfg.insert_late else 0\n",
        "    end_layer = start_layer + cfg.insert_n_layers-1\n",
        "\n",
        "    if first_time:\n",
        "      print( \"Inserting trained small_model\", insert_weights_fname)\n",
        "      print( \"into larger, untrained main_model\", main_fname)\n",
        "      print( \"destination layers:\", start_layer, end_layer)\n",
        "\n",
        "    transfer_full_model(insert_model, main_model, start_layer, end_layer, first_time, first_time)\n",
        "\n",
        "\n",
        "insert_existing_model( True )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IppYbJCP_Ubj"
      },
      "source": [
        "# Part 6C: Transfer useful heads of insert_model into main_model (optional)\n",
        "\n",
        "Transfer just the useful attention heads from insert_model into main_model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5nvucMtl_8tt"
      },
      "outputs": [],
      "source": [
        "# Transfer one attention head's weights from the small to the main model.\n",
        "# The right-most small.n_heads of main_model are updated\n",
        "def transfer_one_head(small_model, small_layer_no, small_head_no, large_model, start_layer):\n",
        "  large_layer_no = start_layer + small_layer_no\n",
        "  large_head_no = large_cfg[\"n_heads\"] - small_cfg[\"n_heads\"] + small_head_no\n",
        "\n",
        "  large_model.blocks[large_layer_no].attn.W_Q.data[large_head_no] = small_model.blocks[small_layer_no].attn.W_Q.clone().data[small_head_no]\n",
        "  large_model.blocks[large_layer_no].attn.W_K.data[large_head_no] = small_model.blocks[small_layer_no].attn.W_K.clone().data[small_head_no]\n",
        "  large_model.blocks[large_layer_no].attn.W_V.data[large_head_no] = small_model.blocks[small_layer_no].attn.W_V.clone().data[small_head_no]\n",
        "\n",
        "  large_model.blocks[large_layer_no].attn.b_Q.data[large_head_no] = small_model.blocks[small_layer_no].attn.b_Q.clone().data[small_head_no]\n",
        "  large_model.blocks[large_layer_no].attn.b_K.data[large_head_no] = small_model.blocks[small_layer_no].attn.b_K.clone().data[small_head_no]\n",
        "  large_model.blocks[large_layer_no].attn.b_V.data[large_head_no] = small_model.blocks[small_layer_no].attn.b_V.clone().data[small_head_no]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iefg8KGG_flE"
      },
      "outputs": [],
      "source": [
        "def transfer_useful_heads(small_model, large_model):\n",
        "  if cfg.insert_mode >= 2 and len(useful_cells) > 0:\n",
        "    # Is the destination the first few or last few layers of the main_model?\n",
        "    start_layer = cfg.n_layers - cfg.insert_n_layers if cfg.insert_late else 0\n",
        "\n",
        "    # TODO: Somewhat inefficent loop as a given head may exist in the list multiple times (with different use_cell.position and use_cell.tags values)\n",
        "    for use_cell in useful_cells:\n",
        "      if use_cell.is_head():\n",
        "        transfer_one_head(small_model, use_cell.layer, use_cell.head, large_model, start_layer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4JDY_Fqa_Wfb"
      },
      "source": [
        "# Part 7: Train add/sub/mult main_model with Infinite Data\n",
        "Train main_model for n_training_steps, storing train_losses per epoch.\n",
        "\n",
        "Each training step (of n_training_steps) new training data (a batch of batch_size tokens) is generated and the model is trained and loss calculated on it. No separate \"testing\" data is needed, as the training data is unique each step. Memorisation of past training data by the model (if any) is minimally beneficial. For 6 digit addition or subtraction there are 1000 billion possible questions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UJ61_nfKxI9c"
      },
      "outputs": [],
      "source": [
        "# Train the model\n",
        "train_losses_list = []\n",
        "per_token_train_losses_list = []\n",
        "\n",
        "for epoch in tqdm.tqdm(range(cfg.n_training_steps)):\n",
        "\n",
        "  tokens = next(ds)\n",
        "  logits = main_model(tokens)\n",
        "\n",
        "  per_token_train_losses_raw, _ = logits_to_tokens_loss(logits, tokens)\n",
        "  per_token_train_losses = loss_fn(per_token_train_losses_raw)\n",
        "  per_token_train_losses_list.append(utils.to_numpy(per_token_train_losses))\n",
        "\n",
        "  train_loss = per_token_train_losses.mean()\n",
        "  train_loss.backward()\n",
        "  train_losses_list.append(train_loss.item())\n",
        "\n",
        "  optimizer.step()\n",
        "  scheduler.step()\n",
        "  optimizer.zero_grad()\n",
        "\n",
        "  if epoch % 100 == 0:\n",
        "    print(epoch, train_loss.item())\n",
        "    if cfg.insert_mode == 2:\n",
        "      # Freeze the useful attention heads from insert_model\n",
        "      transfer_useful_heads(insert_model, main_model)\n",
        "    if cfg.insert_mode == 3:\n",
        "      # Freeze the useful attention heads and MLP layers from insert_model\n",
        "      insert_existing_model( False )\n",
        "\n",
        "print(epoch, train_loss.item())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iwJBrnSJ23eb"
      },
      "outputs": [],
      "source": [
        "final_training_loss = round((train_losses_list[-5]+train_losses_list[-4]+train_losses_list[-3]+train_losses_list[-2]+train_losses_list[-1])/5,9)\n",
        "\n",
        "print( \"AvgFinalLoss\", final_training_loss)\n",
        "print( \"Final loss\", train_losses_list[-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qgzjamSn2y4j"
      },
      "outputs": [],
      "source": [
        "# These temporary Colab files can be manually downloaded from the Colab \"Files\" tab (at left).\n",
        "# The download can be manually loaded into HuggingFace so the \"Accurate Math - Analyse\" Colab can access it.\n",
        "\n",
        "print(\"Saving main model to temporary Colab file\", main_fname_pth)\n",
        "torch.save(main_model.state_dict(), main_fname_pth)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "extra_data = {\n",
        "    \"Config\": cfg.to_dict(),\n",
        "    \"AvgFinalLoss\": final_training_loss,\n",
        "    \"FinalLoss\": train_losses_list[-1],\n",
        "    \"TrainingLoss\": train_losses_list\n",
        "}\n",
        "\n",
        "print( \"Saving main model config etc to temporary Colab file:\", main_fname_json)\n",
        "save_cfg = cfg.to_dict()\n",
        "with open(main_fname_json, 'w') as file:\n",
        "    json.dump(extra_data, file)"
      ],
      "metadata": {
        "id": "aygOXqmalo71"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LzIMC7iEksnz"
      },
      "source": [
        "# Part 8: Training Loss - Addition and Subtraction\n",
        "\n",
        "On 26Jan24 ran several runs:\n",
        "\n",
        "Addition:\n",
        "- 26Jan24: add_d5_l2_h3_train15K_seed372001. AvgFinalLoss=1.6e-08. Handles 1m Qs\n",
        "- 26Jan24: add_d6_l2_h3_train15K_seed372001. AvgFinalLoss=1.7e-08. Handles 1m Qs\n",
        "\n",
        "Subtraction:\n",
        "- 26Jan24: sub_d6_l2_h3_train20K_seed372001. AvgFinalLoss=9.8-05. Fails 1m Qs\n",
        "- 26Jan24: sub_d6_l2_h3_train30K_seed372001. AvgFinalLoss=5.8e-06. Fails 1m Qs\n",
        "\n",
        "Mixed:\n",
        "- 26Jan24: sub_d6_l3_h4_train20K_seed372001. AvgFinalLoss=5e-09. Fails 1m Qs\n",
        "\n",
        "Mixed+Insert (add_d6_l2_h3_train15K_seed372001 inserted):\n",
        "- 26Jan24: ins1_mix_d6_l3_h4_train40K_seed372001. AvgFinalLoss=8e-09. Handles 1m Qs for Add and Sub\n",
        "- 26Jan24: ins2_mix_d6_l4_h4_train40K_seed372001. AvgFinalLoss=7e-09. Fails 1m Qs\n",
        "- 26Jan24: ins3_mix_d6_l4_h3_train40K_seed372001. AvgFinalLoss=2.6e-06. Fails 1m Qs\n",
        "- ??Jan24: ins3_mix_d6_l4_h3_train60K_seed372001. AvgFinalLoss=????"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LiOdvQFD_fxZ"
      },
      "source": [
        "# Part 9: Line Graphs\n",
        "\n",
        "This section analyses the training loss by graphing it at a high level.\n",
        "\n",
        "The loss curve for all digits show visible inflection points (bumps), but is too high level to help understand the algorithm.\n",
        "\n",
        "When this graph is decomposed into 'per digit' graphs, the interesting distinct 'per digit' curves appear, showing each digit is being refined semi-independently, with the model algorithm refining each digit separately."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tmyT2TFd29at"
      },
      "outputs": [],
      "source": [
        "def line(tensor, renderer=None, xaxis=\"\", yaxis=\"\", **kwargs):\n",
        "    px.line(utils.to_numpy(tensor), labels={\"x\":xaxis, \"y\":yaxis}, **kwargs).show(renderer)\n",
        "\n",
        "# Helper function to plot multiple lines\n",
        "def lines(raw_lines_list, x=None, mode='lines', labels=None, xaxis='Epoch', yaxis='Loss', title = '', log_y=False, hover=None, **kwargs):\n",
        "\n",
        "    lines_list = raw_lines_list\n",
        "    log_suffix = '' if log_y==False else ' (Log)'\n",
        "    full_title = title + log_suffix\n",
        "\n",
        "    if type(lines_list)==torch.Tensor:\n",
        "        lines_list = [lines_list[i] for i in range(lines_list.shape[0])]\n",
        "    if x is None:\n",
        "        x=np.arange(len(lines_list[0]))\n",
        "    if cfg.save_graph_to_file :\n",
        "      fig = go.Figure(layout={})\n",
        "      print(full_title)\n",
        "    else:\n",
        "      fig = go.Figure(layout={'title':full_title})\n",
        "\n",
        "    fig.update_xaxes(title=xaxis)\n",
        "    fig.update_yaxes(title=yaxis + log_suffix)\n",
        "    for c, line in enumerate(lines_list):\n",
        "        if type(line)==torch.Tensor:\n",
        "            line = utils.to_numpy(line)\n",
        "        if labels is not None:\n",
        "            label = labels[c]\n",
        "        else:\n",
        "            label = c\n",
        "        fig.add_trace(go.Scatter(x=x, y=line, mode=mode, name=label, hovertext=hover, **kwargs))\n",
        "    if log_y:\n",
        "        fig.update_layout(yaxis_type=\"log\")\n",
        "    if cfg.save_graph_to_file:\n",
        "        fig.update_layout(margin=dict(l=10, r=10, t=10, b=10),width=1200,height=300)\n",
        "\n",
        "    fig.show(bbox_inches=\"tight\")\n",
        "\n",
        "    if cfg.save_graph_to_file:\n",
        "        filename = full_title.replace(\" \", \"\").replace(\"(\", \"\").replace(\")\", \"\").replace(\"&\", \"\").replace(\",\", \"\").replace(\"%\", \"\")   +'.pdf'\n",
        "        pio.write_image(fig, filename)\n",
        "\n",
        "\n",
        "title_suffix = 'Digit Loss Curves ' + main_fname\n",
        "per_token_losses = np.stack(per_token_train_losses_list, axis=0)\n",
        "\n",
        "line(train_losses_list,\n",
        "    title=title_suffix)\n",
        "\n",
        "lines([per_token_losses[:, i] for i in range(1+cfg.n_digits)]+[train_losses_list],\n",
        "      labels = [f'digit {i}' for i in range(1+cfg.n_digits)]+['all_digits'],\n",
        "      title='Per digit'+title_suffix, log_y=False)\n",
        "\n",
        "lines([per_token_losses[:, i] for i in range(1+cfg.n_digits)]+[train_losses_list],\n",
        "      labels = [f'digit {i}' for i in range(1+cfg.n_digits)]+['all_digits'],\n",
        "      title='Per digit'+title_suffix, log_y=True)\n",
        "\n",
        "for i in range(1+cfg.n_digits):\n",
        "  print('Final Loss for digit ' + str(i) + ' is ', per_token_losses[-1, i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pgk8m02A_lxB"
      },
      "source": [
        "# Part 10: Questions Set Up\n",
        "\n",
        "Create sets of sample questions (by task) to ask the model to predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8BpOcko53GrI"
      },
      "outputs": [],
      "source": [
        "def make_varied_questions():\n",
        "  q0 = next(ds)\n",
        "  q1 = next(ds)\n",
        "  q2 = next(ds)\n",
        "  q3 = next(ds)\n",
        "\n",
        "  questions = torch.vstack((q0.cuda(), q1.cuda(), q2.cuda(), q3.cuda()))\n",
        "\n",
        "  return questions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J4bqJ5C23M-A"
      },
      "outputs": [],
      "source": [
        "verbose = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJayaDiR3PgJ"
      },
      "outputs": [],
      "source": [
        "# Build a test batch of random questions\n",
        "varied_questions = make_varied_questions();\n",
        "\n",
        "\n",
        "# Run the sample batch, gather the cache\n",
        "main_model.reset_hooks()\n",
        "main_model.set_use_attn_result(True)\n",
        "sample_logits, sample_cache = main_model.run_with_cache(varied_questions.cuda())\n",
        "print(sample_cache) # Gives names of datasets in the cache\n",
        "sample_losses_raw, sample_max_prob_tokens = logits_to_tokens_loss(sample_logits, varied_questions.cuda())\n",
        "sample_loss_mean = utils.to_numpy(loss_fn(sample_losses_raw).mean())\n",
        "print(\"Sample Mean Loss\", sample_loss_mean)\n",
        "\n",
        "\n",
        "# attn.hook_z is the \"attention head output\" hook point name (at a specified layer)\n",
        "l_attn_hook_z_name = [utils.get_act_name('z', 0, 'a'),utils.get_act_name('z', 1, 'a')] # 'blocks.0.attn.hook_z' etc\n",
        "sample_attn_z = sample_cache[l_attn_hook_z_name[0]]\n",
        "print(\"Sample\", l_attn_hook_z_name[0], sample_attn_z.shape) # gives [239, 18, 3, 170] = num_questions, cfg.n_ctx, n_heads, d_head\n",
        "mean_attn_z = torch.mean(sample_attn_z, dim=0, keepdim=True)\n",
        "print(\"Mean\", l_attn_hook_z_name[0], mean_attn_z.shape) # gives [1, 18, 3, 170] = 1, cfg.n_ctx, n_heads, d_head\n",
        "\n",
        "\n",
        "# hook_resid_pre is the \"pre residual memory update\" hook point name (at a specified layer)\n",
        "l_hook_resid_pre_name = ['blocks.0.hook_resid_pre','blocks.1.hook_resid_pre']\n",
        "\n",
        "\n",
        "# hook_resid_post is the \"post residual memory update\" hook point name (at a specified layer)\n",
        "l_hook_resid_post_name = ['blocks.0.hook_resid_post','blocks.1.hook_resid_post']\n",
        "sample_resid_post = sample_cache[l_hook_resid_post_name[0]]\n",
        "print(\"Sample\", l_hook_resid_post_name[0], sample_resid_post.shape) # gives [239, 18, 510] = num_questions, cfg.n_ctx, d_model\n",
        "mean_resid_post = torch.mean(sample_resid_post, dim=0, keepdim=True)\n",
        "print(\"Mean\", l_hook_resid_post_name[0], mean_resid_post.shape) # gives [1, 18, 510] = 1, cfg.n_ctx, d_model\n",
        "\n",
        "\n",
        "# mlp.hook_post is the \"MLP layer\" hook point name (at a specified layer)\n",
        "l_mlp_hook_post_name = [utils.get_act_name('post', 0),utils.get_act_name('post', 1)] # 'blocks.0.mlp.hook_post' etc\n",
        "sample_mlp_hook_post = sample_cache[l_mlp_hook_post_name[0]]\n",
        "print(\"Sample\", l_mlp_hook_post_name[0], sample_mlp_hook_post.shape) # gives [239, 18, 2040] = num_questions, cfg.n_ctx, d_model*4\n",
        "mean_mlp_hook_post = torch.mean(sample_mlp_hook_post, dim=0, keepdim=True)\n",
        "print(\"Mean\", l_mlp_hook_post_name[0], mean_mlp_hook_post.shape) # gives [1, 18, 2040] = 1, cfg.n_ctx, d_model*4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "deGhK1OC_1mD"
      },
      "source": [
        "# Part 11: Attention Patterns\n",
        "Attention patterns show which token(s) the model's attention heads are paying attention to in each token position of the prediction calculation.\n",
        "\n",
        "For the default CoLab set up, the  model has 3 attention heads, and performs 5 digit addition. The attention pattern is 18 by 18 squares (as 54321+77779=132100 is 18 tokens). Time proceeds vertically downwards, with one additional token being revealed horizontally at each token position, giving the overall triangle shape. This visualisation provided insights. After the question is fully revealed (at token position 11), each head starts attending to pairs of question digits from left to right (i.e. high-value digits before lower-value digits) giving the “double staircase\" shape. The three heads attend to a given digit pair in three different token position, giving a time ordering of heads."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PiVCSBSE3Xpk"
      },
      "outputs": [],
      "source": [
        "def show_token_attention_patterns(index, layer, token_at_index, use_case):\n",
        "\n",
        "  the_tokens = [str(token) for token in token_at_index.tolist()]\n",
        "  if layer == 0:\n",
        "    tokens_str = tokens_to_string(token_at_index)\n",
        "    print(\"Attention patterns for\", tokens_str)\n",
        "\n",
        "  attention_pattern=sample_cache[\"pattern\", layer, \"attn\"][index]\n",
        "  display(cv.attention.attention_patterns(\n",
        "      tokens=the_tokens,\n",
        "      attention=attention_pattern,\n",
        "      #attention_head_names=[f\"L{layer}H{i}\" for i in range(cfg.n_heads)],\n",
        "  ))\n",
        "\n",
        "\n",
        "sample_size = 3\n",
        "\n",
        "# Show attention patterns for some randomly chosen tokens\n",
        "for i in range(sample_size):\n",
        "  for layer in range(cfg.n_layers):\n",
        "    show_token_attention_patterns(i, layer, tokens[i], \"Misc\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hUCXw1k93a20"
      },
      "outputs": [],
      "source": [
        "if cfg.save_graph_to_file:\n",
        "\n",
        "  tokens_str = []\n",
        "  for i in range(cfg.n_heads):\n",
        "    one_token_str = []\n",
        "    for j in tokens[i]:\n",
        "      one_token_str.append(str(utils.to_numpy(j)))\n",
        "    tokens_str.append(one_token_str)\n",
        "\n",
        "  # Refer https://github.com/callummcdougall/CircuitsVis/blob/main/python/circuitsvis/circuitsvis_demo.ipynb\n",
        "\n",
        "  # html_object = cv.attention.from_cache(\n",
        "  #    cache = sample_cache,\n",
        "  #    tokens = tokens_str, # list of list of strings\n",
        "  #    return_mode = \"html\",\n",
        "  #)\n",
        "\n",
        "  # Create a CoLab file containing the attention pattern(s) in HTML\n",
        "  #filename = \"AttentionPattern\" + str(cfg.n_digits) + \"Digits\" + str(cfg.n_heads) + \"Heads.html\"\n",
        "  #with open(filename, \"w\") as f:\n",
        "  #    f.write(html_object.data)\n",
        "\n",
        "  # Manually download the CoLab \"html\" file and open in your local browser.\n",
        "  # Install and use the Edge extension \"FireShot\" to save a portion of the HTML page as a PDF"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "O3OoV2Pd-w7g",
        "F5K2kA0L_FHc",
        "Y8SVQxjtXqa9",
        "3tExv4rk_L0C",
        "Pg267eav_QSl",
        "NvijaDhSjn0B",
        "1fW5galHlQqA",
        "IppYbJCP_Ubj",
        "pgk8m02A_lxB",
        "deGhK1OC_1mD"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}