docs/tutorials/pytorch/language-modeling/bert-base-uncased.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "17ec3098",
   "metadata": {},
   "source": [
    "# Introduction "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e00856",
   "metadata": {},
   "source": [
    "This tutorial demonstrates how to quantize a BERT model with both static and dynamic post training quantization based on [Intel® Neural Compressor](https://github.com/intel/neural-compressor) and benchmark the quantized models. For better int8 performance benefit, multi-instance benchmarking with 4 cores/instance is recommended."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0d4a9c5",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0cbd8bd4",
   "metadata": {},
   "source": [
    "# Prerequisite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86b20e2b",
   "metadata": {},
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1816be1",
   "metadata": {},
   "source": [
    "* Follow [installation](https://github.com/intel/intel-extension-for-transformers#installation) to install **intel-extension-for-transformers**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef103e44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install model dependency\n",
    "!pip install accelerate datasets>=1.8.0 sentencepiece!=0.1.92 torch>=1.10.0 wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e35d9563",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b424903",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import transformers\n",
    "from dataclasses import dataclass, field\n",
    "from datasets import load_dataset, load_metric\n",
    "from itertools import chain\n",
    "from intel_extension_for_transformers.transformers import metrics, OptimizedModel, QuantizationConfig\n",
    "from intel_extension_for_transformers.transformers.trainer import NLPTrainer\n",
    "from transformers import (\n",
    "    CONFIG_MAPPING,\n",
    "    MODEL_FOR_MASKED_LM_MAPPING,\n",
    "    AutoConfig,\n",
    "    AutoModelForMaskedLM,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorForLanguageModeling,\n",
    "    TrainingArguments,\n",
    "    is_torch_tpu_available,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils.versions import require_version\n",
    "from typing import Optional\n",
    "\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "MODEL_CONFIG_CLASSES = list(MODEL_FOR_MASKED_LM_MAPPING.keys())\n",
    "MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "048665d4",
   "metadata": {},
   "source": [
    "## Define arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aead749f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.\n",
    "    \"\"\"\n",
    "\n",
    "    model_name_or_path: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"The model checkpoint for weights initialization.\"\n",
    "            \"Don't set if you want to train a model from scratch.\"\n",
    "        },\n",
    "    )\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataTrainingArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
    "    \"\"\"\n",
    "\n",
    "    dataset_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    dataset_config_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The configuration name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    mlm_probability: float = field(\n",
    "        default=0.15, metadata={\"help\": \"Ratio of tokens to mask for masked language modeling loss\"}\n",
    "    )\n",
    "                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b035ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_args = ModelArguments(\n",
    "    model_name_or_path=\"bert-base-uncased\", \n",
    ")\n",
    "data_args = DataTrainingArguments(\n",
    "    dataset_name=\"wikitext\",\n",
    "    dataset_config_name=\"wikitext-2-raw-v1\",\n",
    ")\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./saved_results_static\",\n",
    "    do_eval=True,\n",
    "    do_train=True,\n",
    "    no_cuda=True,\n",
    "    per_device_eval_batch_size=1,\n",
    "    overwrite_output_dir=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "522f4153",
   "metadata": {},
   "source": [
    "## Download dataset from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22bcb0fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_datasets = load_dataset(\n",
    "    data_args.dataset_name, data_args.dataset_config_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d162019c",
   "metadata": {},
   "source": [
    "## Download fp32 model from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2620a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set seed before initializing model.\n",
    "set_seed(training_args.seed)\n",
    "\n",
    "# get fp32 model\n",
    "config = AutoConfig.from_pretrained(model_args.model_name_or_path)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)\n",
    "model = AutoModelForMaskedLM.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    config=config\n",
    ")\n",
    "model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2634a186",
   "metadata": {},
   "source": [
    "## Preprocessing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a873a324",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First we tokenize all the texts.\n",
    "if training_args.do_train:\n",
    "    column_names = raw_datasets[\"train\"].column_names\n",
    "else:\n",
    "    column_names = raw_datasets[\"validation\"].column_names\n",
    "text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n",
    "\n",
    "max_seq_length = tokenizer.model_max_length\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[text_column_name], return_special_tokens_mask=True)\n",
    "\n",
    "column_names = raw_datasets[\"train\"].column_names\n",
    "text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n",
    "    \n",
    "with training_args.main_process_first(desc=\"dataset map tokenization\"):\n",
    "    tokenized_datasets = raw_datasets.map(\n",
    "        tokenize_function,\n",
    "        batched=True,\n",
    "        remove_columns=column_names,\n",
    "        load_from_cache_file=True,\n",
    "        desc=\"Running tokenizer on every text in dataset\",\n",
    "    )\n",
    "\n",
    "# Main data processing function that will concatenate all texts from our dataset and generate chunks of max_seq_length.\n",
    "def group_texts(examples):\n",
    "    # Concatenate all texts.\n",
    "    concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n",
    "    # customize this part to your needs.\n",
    "    if total_length >= max_seq_length:\n",
    "        total_length = (total_length // max_seq_length) * max_seq_length\n",
    "    # Split by chunks of max_len.\n",
    "    result = {\n",
    "        k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    return result\n",
    "\n",
    "# Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a\n",
    "# remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value\n",
    "# might be slower to preprocess.\n",
    "\n",
    "with training_args.main_process_first(desc=\"grouping texts together\"):\n",
    "    tokenized_datasets = tokenized_datasets.map(\n",
    "        group_texts,\n",
    "        batched=True,\n",
    "        load_from_cache_file=True,\n",
    "        desc=f\"Grouping texts in chunks of {max_seq_length}\",\n",
    "    )\n",
    "\n",
    "if training_args.do_train:\n",
    "    if \"train\" not in tokenized_datasets:\n",
    "        raise ValueError(\"--do_train requires a train dataset\")\n",
    "    train_dataset = tokenized_datasets[\"train\"]\n",
    "\n",
    "if training_args.do_eval:\n",
    "    if \"validation\" not in tokenized_datasets:\n",
    "        raise ValueError(\"--do_eval requires a validation dataset\")\n",
    "    eval_dataset = tokenized_datasets[\"validation\"]\n",
    "\n",
    "    def preprocess_logits_for_metrics(logits, labels):\n",
    "        if isinstance(logits, tuple):\n",
    "            # Depending on the model and config, logits may contain extra tensors,\n",
    "            # like past_key_values, but logits always come first\n",
    "            logits = logits[0]\n",
    "        return logits.argmax(dim=-1)\n",
    "\n",
    "    metric = load_metric(\"accuracy\")\n",
    "\n",
    "    def compute_metrics(eval_preds):\n",
    "        preds, labels = eval_preds\n",
    "        # preds have the same shape as the labels, after the argmax(-1) has been calculated\n",
    "        # by preprocess_logits_for_metrics\n",
    "        labels = labels.reshape(-1)\n",
    "        preds = preds.reshape(-1)\n",
    "        mask = labels != -100\n",
    "        labels = labels[mask]\n",
    "        preds = preds[mask]\n",
    "        return metric.compute(predictions=preds, references=labels)\n",
    "\n",
    "\n",
    "# Data collator will take care of randomly masking the tokens.\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer,\n",
    "    mlm_probability=data_args.mlm_probability,\n",
    "    pad_to_multiple_of=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942d2967",
   "metadata": {},
   "source": [
    "# Quantization & Benchmark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a07a8735",
   "metadata": {},
   "source": [
    "## Static Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df04531d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Trainer\n",
    "set_seed(training_args.seed)\n",
    "trainer_ptq_static = NLPTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics if training_args.do_eval and not is_torch_tpu_available() else None,\n",
    "    preprocess_logits_for_metrics=preprocess_logits_for_metrics\n",
    "    if training_args.do_eval and not is_torch_tpu_available()\n",
    "    else None,\n",
    ")\n",
    "\n",
    "tune_metric = metrics.Metric(\n",
    "    name=\"eval_loss\", # Metric used for the tuning strategy.\n",
    "    is_relative=True, # Metric tolerance mode, True is for relative, otherwise for absolute.\n",
    "    criterion=\"0.25\", # Performance tolerance when optimizing the model.\n",
    "    greater_is_better=False \n",
    ")\n",
    "quantization_config = QuantizationConfig(\n",
    "    approach=\"PostTrainingStatic\",\n",
    "    metrics=[tune_metric],\n",
    ")\n",
    "\n",
    "# run quantization\n",
    "trainer_ptq_static.quantize(quant_config=quantization_config)\n",
    "\n",
    "# save quantized model\n",
    "trainer_ptq_static.save_model(\"./saved_results_static\")\n",
    "model.config.save_pretrained(\"./saved_results_static\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7e93de",
   "metadata": {},
   "source": [
    "## Run Benchmark after Static Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf54c034",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "results = trainer_ptq_static.evaluate()\n",
    "bert_task_acc_keys = ['eval_loss', 'eval_f1', 'eval_accuracy', 'eval_matthews_correlation',\n",
    "                              'eval_pearson', 'eval_mcc', 'eval_spearmanr']\n",
    "\n",
    "throughput_ptq_static = results.get(\"eval_samples_per_second\")\n",
    "eval_loss_ptq_static = results[\"eval_loss\"]\n",
    "print('Batch size = {}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval eval_loss Accuracy: {}\".format(eval_loss_ptq_static))\n",
    "print(\"Latency: {:.3f} ms\".format(1000 / throughput_ptq_static))\n",
    "print(\"Throughput: {} samples/sec\".format(throughput_ptq_static))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7e93de",
   "metadata": {},
   "source": [
    "## Run Benchmark after Static Post Training Quantization with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a6795aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.system('numactl --hardware')\n",
    "results = os.system('bash ../multi_instance.sh --model=saved_results_static --core_per_instance=4 --data_type=int8')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a51f6ca",
   "metadata": {},
   "source": [
    "## Dynamic Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a57aa0d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Trainer\n",
    "set_seed(training_args.seed)\n",
    "training_args.output_dir = \"./saved_results_dynamic\"\n",
    "trainer_ptq_dynamic = NLPTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics if training_args.do_eval and not is_torch_tpu_available() else None,\n",
    "    preprocess_logits_for_metrics=preprocess_logits_for_metrics\n",
    "    if training_args.do_eval and not is_torch_tpu_available()\n",
    "    else None,\n",
    ")\n",
    "\n",
    "tune_metric = metrics.Metric(\n",
    "    name=\"eval_loss\", \n",
    "    is_relative=True,\n",
    "    criterion=\"0.25\", # why performance tolerance\n",
    "    greater_is_better=False\n",
    ")\n",
    "quantization_config = QuantizationConfig(\n",
    "    approach=\"PostTrainingDynamic\",\n",
    "    metrics=[tune_metric],\n",
    ")\n",
    "\n",
    "# run quantization\n",
    "trainer_ptq_dynamic.quantize(quant_config=quantization_config)\n",
    "\n",
    "# save quantized model\n",
    "trainer_ptq_dynamic.save_model(\"./saved_results_dynamic\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dd5b041",
   "metadata": {},
   "source": [
    "## Run Benchmark after Dynamic Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "689e288d",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "results = trainer_ptq_dynamic.evaluate()\n",
    "\n",
    "throughput_ptq_dynamic = results.get(\"eval_samples_per_second\")\n",
    "eval_loss_ptq_dynamic = results[\"eval_loss\"]\n",
    "print('Batch size = {}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval eval_loss Accuracy: {}\".format(eval_loss_ptq_dynamic))\n",
    "print(\"Latency: {:.3f} ms\".format(1000 / throughput_ptq_dynamic))\n",
    "print(\"Throughput: {} samples/sec\".format(throughput_ptq_dynamic))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7e93de",
   "metadata": {},
   "source": [
    "## Run Benchmark after Dynamic Post Training Quantization with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea631f92",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "results = os.system('bash ../multi_instance.sh --model=saved_results_dynamic --core_per_instance=4 --data_type=int8')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44cca2a1",
   "metadata": {},
   "source": [
    "## Run Benchmark for FP32 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda9f985",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Initialize the Trainer\n",
    "set_seed(training_args.seed)\n",
    "trainer = NLPTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics if training_args.do_eval and not is_torch_tpu_available() else None,\n",
    "    preprocess_logits_for_metrics=preprocess_logits_for_metrics\n",
    "    if training_args.do_eval and not is_torch_tpu_available()\n",
    "    else None,\n",
    ")\n",
    "results = trainer.evaluate()\n",
    "\n",
    "throughput_fp32 = results.get(\"eval_samples_per_second\")\n",
    "eval_loss_fp32 = results[\"eval_loss\"]\n",
    "print('Batch size = {}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval eval_loss Accuracy: {}\".format(eval_loss_fp32))\n",
    "print(\"Latency: {:.3f} ms\".format(1000 / throughput_fp32))\n",
    "print(\"Throughput: {} samples/sec\".format(throughput_fp32))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7e93de",
   "metadata": {},
   "source": [
    "## Run Benchmark for FP32 Model with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "571317cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "results = os.system('bash ../multi_instance.sh --model=bert-base-uncased --core_per_instance=4 --data_type=fp32')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.6 64-bit",
   "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.8.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "f54fd8d6160ddfbc370985ee3ad2925997e28943a671b1747496a6859c59cd26"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}