plot_curves.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "433492dd-6736-4e98-8df2-cf673bd42165",
   "metadata": {},
   "source": [
    "# Data Efficient Neural Scaling Laws via Model Reusing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2028fd5b-42f2-4633-beda-5632b43d6185",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "**Peihao Wang, Rameswar Panda, Zhangyang (Atlas) Wang**\n",
    "\n",
    "This notebook is aimed at helping reproduce Figure 4(a) of ICML 2023 paper [Data Efficient Neural Scaling Law via Model Reusing](https://openreview.net/pdf?id=iXYnIz4RRx).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a6bb96c-6139-422b-9af4-b22e3de9e046",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import shutil\n",
    "import pickle\n",
    "import random\n",
    "from typing import Dict, List, Tuple\n",
    "from datetime import datetime\n",
    "import time\n",
    "import math\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "from transformers import (\n",
    "    PreTrainedModel,\n",
    "    PreTrainedTokenizer,\n",
    "    BertConfig,\n",
    "    BertForMaskedLM,\n",
    "    BertTokenizer,\n",
    "    GPT2Config,\n",
    "    GPT2LMHeadModel,\n",
    "    GPT2Tokenizer,\n",
    "    OpenAIGPTConfig,\n",
    "    OpenAIGPTLMHeadModel,\n",
    "    OpenAIGPTTokenizer,\n",
    "    RobertaConfig,\n",
    "    RobertaForMaskedLM,\n",
    "    RobertaTokenizer,\n",
    ")\n",
    "\n",
    "from data import CoLDataset\n",
    "from model import SimpleBertForMaskedLM, SimpleRobertaForMaskedLM\n",
    "\n",
    "try:\n",
    "    from torch.utils.tensorboard import SummaryWriter\n",
    "except ImportError:\n",
    "    from tensorboardX import SummaryWriter\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8643494c-bee8-45c2-b1e2-1e404033df44",
   "metadata": {},
   "source": [
    "## Utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98fb0d8-f97b-4afa-956d-0731c8f3c0bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "MODEL_CLASSES = {\n",
    "    \"gpt2\": (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer),\n",
    "    \"openai-gpt\": (OpenAIGPTConfig, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),\n",
    "    \"bert\": (BertConfig, SimpleBertForMaskedLM, BertTokenizer),\n",
    "    \"roberta\": (RobertaConfig, SimpleRobertaForMaskedLM, RobertaTokenizer),\n",
    "}\n",
    "\n",
    "\n",
    "def set_seed(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42373cee-8773-435f-abb9-295d668373f3",
   "metadata": {},
   "source": [
    "## Mask language modeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4979a4a7-f89a-47b8-971a-9676a56d3e2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def mask_tokens(inputs: torch.Tensor, tokenizer: PreTrainedTokenizer, mlm_probability=0.15) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "    \"\"\" Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. \"\"\"\n",
    "\n",
    "    if tokenizer.mask_token is None:\n",
    "        raise ValueError(\n",
    "            \"This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer.\"\n",
    "        )\n",
    "\n",
    "    labels = inputs.clone()\n",
    "    # We sample a few tokens in each sequence for masked-LM training (with probability mlm_probability defaults to 0.15 in Bert/RoBERTa)\n",
    "    probability_matrix = torch.full(labels.shape, mlm_probability)\n",
    "    special_tokens_mask = [\n",
    "        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()\n",
    "    ]\n",
    "    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)\n",
    "    if tokenizer._pad_token is not None:\n",
    "        padding_mask = labels.eq(tokenizer.pad_token_id)\n",
    "        probability_matrix.masked_fill_(padding_mask, value=0.0)\n",
    "    masked_indices = torch.bernoulli(probability_matrix).bool()\n",
    "    labels[~masked_indices] = -100  # We only compute loss on masked tokens\n",
    "\n",
    "    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])\n",
    "    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices\n",
    "    inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)\n",
    "\n",
    "    # 10% of the time, we replace masked input tokens with random word\n",
    "    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced\n",
    "    random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long)\n",
    "    inputs[indices_random] = random_words[indices_random]\n",
    "\n",
    "    # The rest of the time (10% of the time) we keep the masked input tokens unchanged\n",
    "    return inputs, labels\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5122bc1e-6c06-4841-9e6b-0a13dd779871",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9c2c3c4-c0da-4823-9bf9-dc40981e5dcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Tunable hyperparameters\n",
    "seed = 0 # @param {type:\"integer\"}\n",
    "gpu = 0 # @param {type:\"integer\"}\n",
    "\n",
    "eval_data_file = 'data/wiki-cased/en.valid.raw' # @param {type:\"string\"}\n",
    "\n",
    "checkpoints_path = './data-efficient-scaling/' # @param {type:\"string\"}\n",
    "\n",
    "batch_size = 64 # @param {type:\"integer\"}\n",
    "\n",
    "# Fixed hyperparameters\n",
    "tokenizer_name = 'bert-base-uncased'\n",
    "model_type = 'bert'\n",
    "block_size = 126\n",
    "split_sent = True\n",
    "cache_dir = None\n",
    "\n",
    "\n",
    "configs = ['12L-64H', '12L-128H', '12L-192H', '12L-256H', '12L-320H', '12L-384H', '12L-448H', '12L-512H', '12L-576H', '12L-640H']\n",
    "ratios = [0.009, 0.005, 0.004, 0.003, 0.002, 0.001][::-1]\n",
    "\n",
    "\n",
    "# Get class names\n",
    "config_class, model_class, tokenizer_class = MODEL_CLASSES[model_type]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "894164a3-30cb-4bae-92be-b08f9f0cdf15",
   "metadata": {},
   "source": [
    "## Count model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef8df9be-b674-45b7-a3bd-2f51f15211dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def compute_num_params(conf):\n",
    "    conf_path = f'configs/bert-{conf}.json'\n",
    "    config = config_class.from_pretrained(conf_path, cache_dir=None)\n",
    "    model = model_class(config=config)\n",
    "    return sum(p.numel() for p in model.parameters())\n",
    "\n",
    "num_model_params = {conf: compute_num_params(conf) for conf in configs}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1ed95d3-7b91-49c1-9270-64dbde0f2927",
   "metadata": {},
   "source": [
    "## Evaluate checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953a0473-d3b5-4f7e-acef-8cda97537013",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def evaluate(eval_dataset, eval_batch_size, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device, prefix=\"\", mlm=True) -> Dict:\n",
    "\n",
    "    # Note that DistributedSampler samples randomly\n",
    "\n",
    "    def collate(examples: List[torch.Tensor]):\n",
    "        if tokenizer._pad_token is None:\n",
    "            return pad_sequence(examples, batch_first=True)\n",
    "        return pad_sequence(examples, batch_first=True, padding_value=tokenizer.pad_token_id)\n",
    "\n",
    "    eval_sampler = SequentialSampler(eval_dataset)\n",
    "    eval_dataloader = DataLoader(\n",
    "        eval_dataset, sampler=eval_sampler, batch_size=eval_batch_size, collate_fn=collate\n",
    "    )\n",
    "\n",
    "    # Eval!\n",
    "    eval_loss = 0.0\n",
    "    nb_eval_steps = 0\n",
    "    model.eval()\n",
    "\n",
    "    for batch in tqdm(eval_dataloader, desc=\"Evaluating\"):\n",
    "        inputs, labels = mask_tokens(batch, tokenizer) if mlm else (batch, batch)\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        # If some of the input is padded, then the attention mask is needed\n",
    "        attention_mask = (inputs != tokenizer.pad_token_id)  # word_tokens --> 1, pad_token --> 0\n",
    "        if attention_mask.all():\n",
    "            attention_mask = None\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(inputs, attention_mask=attention_mask, masked_lm_labels=labels) if mlm else model(inputs, labels=labels)\n",
    "            lm_loss = outputs['lm_loss']\n",
    "            eval_loss += lm_loss.mean().item()\n",
    "        nb_eval_steps += 1\n",
    "\n",
    "    eval_loss = eval_loss / nb_eval_steps\n",
    "    perplexity = torch.exp(torch.tensor(eval_loss)).item()\n",
    "\n",
    "    result = {\"perplexity\": perplexity}\n",
    "\n",
    "    return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e996eb3-2fe4-4d52-92ae-90ee6da5d384",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup CUDA & GPU\n",
    "torch.cuda.set_device(gpu)\n",
    "device = torch.device(\"cuda\", gpu)\n",
    "\n",
    "ppl_results = {}\n",
    "\n",
    "for ratio in ratios:\n",
    "    for conf in configs:\n",
    "        for k in [f'bert-{conf}-{ratio}D', f'bert-n2n-{conf}-{ratio}D']:\n",
    "\n",
    "            model_name_or_path = os.path.join(checkpoints_path, k)\n",
    "            \n",
    "            print(f\"Evaluating checkpoint at {model_name_or_path}\")\n",
    "\n",
    "            # Set seed\n",
    "            set_seed(seed)\n",
    "\n",
    "\n",
    "            # Get config\n",
    "            config = config_class.from_pretrained(model_name_or_path, cache_dir=cache_dir)\n",
    "\n",
    "            # Get tokenizer\n",
    "            tokenizer = tokenizer_class.from_pretrained(model_name_or_path, cache_dir=cache_dir)\n",
    "            assert block_size <= tokenizer.model_max_length\n",
    "\n",
    "            # Load model\n",
    "            model = model_class.from_pretrained(\n",
    "                model_name_or_path,\n",
    "                from_tf=bool(\".ckpt\" in model_name_or_path),\n",
    "                config=config,\n",
    "                cache_dir=cache_dir\n",
    "            )\n",
    "            model.to(device)\n",
    "\n",
    "\n",
    "            # Loop to handle MNLI double evaluation (matched, mis-matched)\n",
    "            eval_dataset = CoLDataset(eval_data_file, tokenizer_name, tokenizer, block_size, split_sent=split_sent, verbose=False)\n",
    "\n",
    "            result = evaluate(eval_dataset, batch_size, model, tokenizer, device)\n",
    "\n",
    "            ppl_results[k] = math.log(result['perplexity'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13ed496e-c0ae-4709-97dd-587871226d32",
   "metadata": {},
   "source": [
    "## Load or save results for next time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bf5fb59-0cdd-4184-9454-b92d30cb5b11",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### uncomment to read results from a saved one\n",
    "# if os.path.exists('ppl_results.pkl'):\n",
    "#     with open('ppl_results.pkl', 'rb') as f:\n",
    "#         # read information from file\n",
    "#         ppl_results = pickle.load(f)\n",
    "\n",
    "# save results\n",
    "with open('ppl_results.pkl', 'wb') as f:\n",
    "    # dump information to that file\n",
    "    pickle.dump(ppl_results, f)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dea5e5fa-2d5a-4830-a2bd-270518beadf4",
   "metadata": {},
   "source": [
    "## Plot curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "974805c0-9af7-4b66-8a67-72de9d069e77",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# setup figure\n",
    "NEW_SIZE = 20\n",
    "plt.rcParams[\"font.family\"] = \"DejaVu Sans\"\n",
    "plt.rcParams[\"font.size\"] = NEW_SIZE\n",
    "fig = plt.figure(figsize=(8, 8))\n",
    "[x.set_linewidth(2.) for x in plt.gca().spines.values()]\n",
    "\n",
    "\n",
    "# colormaps\n",
    "cols_scratch = matplotlib.cm.get_cmap('GnBu', len(ratios))(np.linspace(0.4, 1., len(ratios)))\n",
    "cols_n2n = matplotlib.cm.get_cmap('YlOrBr', len(ratios))(np.linspace(0.3, 1., len(ratios)))\n",
    "\n",
    "# plot curves\n",
    "for ratio, col_scratch, col_n2n in zip(ratios, cols_scratch, cols_n2n):\n",
    "    xs = []\n",
    "    ys_scratch = []\n",
    "    ys_n2n = []\n",
    "\n",
    "    for conf in configs:\n",
    "        \n",
    "        num_params = num_model_params[conf] / 1e6\n",
    "\n",
    "        xs.append(int(num_params))\n",
    "        ys_scratch.append(ppl_results[f'bert-{conf}-{ratio}D'])\n",
    "        ys_n2n.append(ppl_results[f'bert-n2n-{conf}-{ratio}D'])\n",
    "\n",
    "    ratio_txt = f'{float((ratio if ratio > 0 else 1.) * 100):.1f}%'\n",
    "    xs, ys_scratch, ys_n2n = np.array(xs), np.array(ys_scratch), np.array(ys_n2n)\n",
    "\n",
    "    plt.plot(xs, ys_scratch, linewidth=3.5, linestyle='dotted', alpha=0.8, color=col_scratch)\n",
    "    plt.scatter(xs, ys_scratch, s=150, color=col_scratch, alpha=1., linewidths=0, marker='*')\n",
    "\n",
    "    plt.plot(xs, ys_n2n, label=ratio_txt, linewidth=3.5, alpha=0.8, color=col_n2n)\n",
    "    plt.scatter(xs, ys_n2n, s=100, color=col_n2n, alpha=1., linewidths=0, marker='o')\n",
    "    \n",
    "leg = plt.legend(title=\"Data Frac.\", bbox_to_anchor=(1., 1), frameon=False, prop={'size': 18})\n",
    "plt.setp(leg.get_title(), fontsize=18)\n",
    "\n",
    "plt.xlabel('# Param. (M)')\n",
    "plt.ylabel('Log Perplexity')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f5a607c-57c2-44ce-8aac-ea1da8dc45e2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "eff_bert",
   "language": "python",
   "name": "eff_bert"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}