{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "rEJBSTyZIrIb" }, "source": [ "# Faster Multi-label Text Classification on IPU with Packed BERT - Fine-tuning and Inference" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "This notebook builds on the process in the notebook on fine-tuning BERT on a text classification task `natural-language-processing/other-use-cases/text_classification.ipynb`. Here we show how to implement packing for BERT for multi-label classification. [Packing](https://www.graphcore.ai/posts/introducing-packed-bert-for-2x-faster-training-in-natural-language-processing) is an optimisation method originally used for 2x faster BERT pre-training, which can now also provide massive throughput increases for **fine-tuning** and **batched inference**! \n", "\n", "**So, what *is* packing?** The basic idea of \"packing\" a dataset is to utilise the requirement for constant-shaped inputs into a model. Instead of padding it with empty, unused space, we can recycle this unused space and fill it with more inputs! The architecture of transformer models like BERT supports this, and lets us optimally use this space to process multiple sequences within one input.\n", "\n", "**And here is why you might want to use it:** Having a single input that contains multiple sequences leads to multiple sequences being processed in parallel in a single pass within a single iteration inside a batch, increasing the \"effective\" batch size of the model by a considerable factor in many cases, and most importantly, increasing model throughput for training and batched inference significantly.\n", "\n", "The [GoEmotions](https://ai.googleblog.com/2021/10/goemotions-dataset-for-fine-grained.html) dataset will be fine-tuned using packing. This notebook outlines how to easily enable packing for BERT when performing fine-tuning and inference on a text-classification task in 🤗 Graphcore Optimum, resulting in an impressive 5-9x faster training and inference run-time for the dataset. \n", "\n", "You can read more about packing in the original paper on [efficient sequence packing without cross-contamination](https://arxiv.org/abs/2107.02027)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "![GoEmotions dataset (Source: GoogleBlog)](../images/go_emotions.png)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "The dataset consists of 58k comments labelled for 27 different emotion categories (and a 28th \"neutral\" category). This dataset is used for multi-label, multi-class classification. The dataset format and categories can be viewed on the [Hugging Face Hub](https://huggingface.co/datasets/go_emotions)." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "| Domain | Tasks | Model | Datasets | Workflow | Number of IPUs | Execution time |\n", "|---------|-------|-------|----------|----------|--------------|--------------|\n", "| Natural language processing | Multi-label text classification | PackedBERT | GoEmotions | Fine-tuning, Inference | POD4 or POD16 | |" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Environment setup\n", "\n", "The best way to run this demo is on Paperspace Gradient's cloud IPUs because everything is already set up for you.\n", "\n", "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Dependencies and configuration\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "In order to improve usability and support for future users, Graphcore would like to collect information about the\n", "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n", "\n", "- User progression through the notebook\n", "- Notebook details: number of cells, code being run and the output of the cells\n", "- Environment details\n", "\n", "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Also make sure all the packages required for this notebook are installed." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "%pip install optimum-graphcore==0.7\n", "%pip install scikit-learn \"datasets>=2.7.0\" evaluate tokenizers matplotlib scipy huggingface_hub\n", "\n", "%pip install graphcore-cloud-tools[logger]@git+https://github.com/graphcore/graphcore-cloud-tools\n", "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Values for machine size and cache directories can be configured through environment variables or directly in the notebook:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "import os\n", "n_ipu = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n", "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache/\") + \"/packed_bert_mlseqcls/\"" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "We use both data parallelism and pipeline parallelism (see this [tutorial on efficient data loading](https://github.com/graphcore/examples/tree/master/tutorials/tutorials/pytorch/efficient_data_loading/walkthrough.ipynb) for more information). Therefore the global batch size, which is the actual number of samples used for the weight update, is determined using four factors:\n", "\n", " global batch size = micro_batch_size * gradient accumulation steps * device iterations * replication factor\n", "\n", "Replication factor is determined by dividing `n_ipu`, the number of IPUs to use, by `ipus_per_replica`.\n", "\n", "Depending on your model and the IPU Pod machine you are using, you might need to adjust these batch-size-related arguments.\n", "\n", "By default this notebook is configured to run on 4 IPUs.\n", "\n", "We also define `max_seq_length` which is the maximum length a sequence can be, and all sequences will be padded to this length. Therefore, `max_seq_length` should not be larger than the maximum length of the model. Given the small size of the sequences in `go-emotions`, we can reduce the model maximum input size to `max_seq_length = 256`. \n", "\n", "Set these parameters and the rest of the notebook should run smoothly:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "zVvslsfMIrIh", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "model_checkpoint = \"bert-base-uncased\" # Default uncased pre-trained BERT checkpoint\n", "ipu_config_name = \"Graphcore/bert-base-uncased\" # Default Graphcore IPU config initialisation for pre-trained BERT\n", "max_seq_length = 256 # The maximum sequence length allowed for sequences in the model.\n", "number_ipus_per_replica = 4\n", "number_of_replicas = n_ipu // number_ipus_per_replica\n", "micro_batch_size = 2 \n", "# The size of the machine and the length of the pipeline impact the number of \n", "# samples that need to be processed between gradient updates\n", "gradient_accumulation_steps = 39 // number_of_replicas\n", "device_iterations = 32\n", "model_task = 'go_emotions'\n", "num_labels = 28" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Gradients are not calculated during validation, so gradient accumulation is not applicable, and the global batch size for validation can be defined separately as:\n", "```\n", "global_validation_batch_size=device_iterations*replication_factor*max_seq_per_pack\n", "```\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### Sharing your model with the community" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "You can share your model with the 🤗 community. You do this by completing the following steps:\n", "\n", "1. Store your authentication token from the 🤗 website. [Sign up to 🤗](https://huggingface.co/join) if you haven't already.\n", "2. Execute the following cell and input your username and authentication token." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from huggingface_hub import notebook_login\n", "\n", "notebook_login()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Git-LFS must also be installed to manage large files:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "!apt install git-lfs" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "whPRbBNbIrIl" }, "source": [ "## Loading the dataset" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "W7QYTpxXIrIl" }, "source": [ "We will use the [🤗 Datasets](https://github.com/huggingface/datasets) library to download the data and get the metric we will use for evaluation (to compare our model to the benchmark). This can be easily done with the functions `load_dataset` and `load_metric`. " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "IreSlFmlIrIm", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from datasets import load_dataset\n", "import evaluate" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "dataset = load_dataset(model_task)\n", "metric = evaluate.load(\"roc_auc\", \"multilabel\")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "RzfPtOMoIrIu" }, "source": [ "The `dataset` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test sets." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "GWiVUF0jIrIv", "outputId": "35e3ea43-f397-4a54-c90c-f2cf8d36873e", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "dataset" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "u3EtYfeHIrIz" }, "source": [ "To access an actual element, you need to select a split (\"train\" in the example) and then specify an index:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "X6HrpprwIrIz", "outputId": "d7670bc0-42e4-4c09-8a6a-5c018ded7d95", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "dataset[\"train\"][0]" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "WHUmphG3IrI3" }, "source": [ "To get a sense of what the data looks like, the following function will show some samples picked randomly from the dataset." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "i3j8APAoIrI3", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "import datasets\n", "import random\n", "import pandas as pd\n", "from IPython.display import display, HTML\n", "\n", "def show_random_elements(dataset, num_examples=10):\n", " assert num_examples <= len(dataset), \"Can't pick more elements than there are in the dataset.\"\n", " picks = []\n", " for _ in range(num_examples):\n", " pick = random.randint(0, len(dataset)-1)\n", " while pick in picks:\n", " pick = random.randint(0, len(dataset)-1)\n", " picks.append(pick)\n", " \n", " df = pd.DataFrame(dataset[picks])\n", " for column, typ in dataset.features.items():\n", " if isinstance(typ, datasets.ClassLabel):\n", " df[column] = df[column].transform(lambda i: typ.names[i])\n", " display(HTML(df.to_html()))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "SZy5tRB_IrI7", "outputId": "ba8f2124-e485-488f-8c0c-254f34f24f13", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "show_random_elements(dataset[\"train\"])" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "lnjDIuQ3IrI-" }, "source": [ "The metric is an instance of [`datasets.Metric`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Metric):" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5o4rUteaIrI_", "outputId": "18038ef5-554c-45c5-e00a-133b02ec10f1", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "metric" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "jAWdqcUBIrJC" }, "source": [ "You can call its `compute` method with your predictions and labels directly and it will return a dictionary with the metric(s) value:" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "n9qywopnIrJH" }, "source": [ "## Preprocessing the data" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "YVx71GdAIrJH" }, "source": [ "Before we can feed the text samples to our model, we need to preprocess them. This is done by a 🤗 Transformers `Tokenizer` which will (as the name indicates) tokenize the inputs (including converting the tokens to their corresponding IDs in the pre-trained vocabulary), putting them into a format the model expects, as well as generating the other inputs that the model requires.\n", "\n", "To do all of this, we instantiate our tokenizer with the `AutoTokenizer.from_pretrained` method, which will ensure:\n", "\n", "- We get a tokenizer that corresponds to the model architecture we want to use.\n", "- We download the vocabulary used when pre-training this specific checkpoint.\n", "\n", "That vocabulary will be cached, so it's not downloaded again the next time we run the cell." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "eXNLu_-nIrJI", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from transformers import AutoTokenizer\n", " \n", "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "Vl6IidfdIrJK" }, "source": [ "We pass `use_fast=True` to the call above to use one of the fast tokenizers (backed by Rust) from the 🤗 Tokenizers library. Fast tokenizers are available for almost all models, but if you get an error with the previous call then simply set `use_fast` to False." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "rowT4iCLIrJK" }, "source": [ "You can directly call this tokenizer on one sentence or a pair of sentences:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "a5hBlsrHIrJL", "outputId": "acdaa98a-a8cd-4a20-89b8-cc26437bbe90", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "tokenizer(\"Hello, this one sentence!\", \"And this sentence goes with it.\")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "qo_0B1M2IrJM" }, "source": [ "Depending on the model you selected, you will see different keys in the dictionary returned by the cell above. They don't matter for what we're doing here, but they are required by the model we will instantiate later. You can learn more about keys in this [tutorial on preprocessing](https://huggingface.co/transformers/preprocessing.html).\n", "\n", "To preprocess our dataset, we will need the names of the columns containing the sentence(s). In this case, the column is called `'text'` and it is indexed as such in the tokenization function." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "2C0hcmp9IrJQ" }, "source": [ "We can then write the function that will preprocess our samples. We feed the samples to the `tokenizer` with three arguments.`truncation=True` will ensure that an input longer than the maximum length will be truncated to the maximum length. `max_length=max_seq_length` sets the maximum length of a sequence.\n", "\n", "**Note that since we use packing later, we don't set any padding in the tokenizer.**" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "vc0BSBLIIrJQ", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "# no padding for packing\n", "def preprocess_function(examples):\n", " return tokenizer(examples['text'], truncation=True, max_length=max_seq_length)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "For multi-label classification, we also need to convert our labels from integer values indicating a category to an N-hot binary format (where N is the maximum number of labels). This makes sure we have constant-sized labels, and all of our labels (one input can have multiple target labels) are present for training. The conversion looks something like this:\n", "\n", "```python\n", "unprocessed_labels = [3,21] # Where 3 and 21 are label categories\n", "preprocessed_labels = id_to_N_hot([3,21])\n", "preprocessed_labels = [0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0]\n", "```" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "\n", "The following function processes one example and converts it to N-hot. The `.map()` functionality available in the `datasets` library allows the function to be applied easily to the entire dataset." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "import numpy as np\n", "\n", "def id_to_N_hot(example):\n", " indexes = example['labels']\n", " label = np.zeros((num_labels,), dtype=int)\n", " for idx in indexes:\n", " label[idx] = 1\n", " example['labels'] = label\n", " return example" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "zS-6iXTkIrJT" }, "source": [ "To apply this function on all the sentences (or pairs of sentences) in our dataset, we just use the `map` method of our `dataset` object we created earlier. This will apply the function to all the elements of all the splits in `dataset`, so our training, validation and testing data will be preprocessed in a single command." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "DDtsaJeVIrJT", "outputId": "aa4734bf-4ef5-4437-9948-2c16363da719", "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "encoded_dataset = dataset.map(id_to_N_hot)\n", "encoded_dataset = encoded_dataset.map(preprocess_function, batched=True)\n", "\n", "len(encoded_dataset['validation'])" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "voWiw8C7IrJV" }, "source": [ "Even better, the results are automatically cached by the 🤗 Datasets library to avoid spending time on this step the next time you run your notebook. The 🤗 Datasets library is able to detect when the function you pass to `map` has changed (and thus to not use the cached data). For instance, it will detect if you change the task in the first cell and rerun the notebook. 🤗 Datasets warns you when it uses cached files. You can pass `load_from_cache_file=False` in the call to `map` to not use the cached files and force the preprocessing to be applied again.\n", "\n", "Note that we passed `batched=True` to encode the text samples together into batches. This is to leverage the full benefit of the fast tokenizer we loaded earlier, which will use multi-threading to treat the text samples in a batch concurrently." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Packing the dataset" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "To implement packing, we need to pack our dataset first. Each new element will be a pack containing at most `max_seq_per_pack` sequences." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "max_seq_per_pack = 6" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "The problem type for this task is multi_label_classification and this also needs to be defined for the packed model to work." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "problem_type = 'multi_label_classification'" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### Packing algorithm" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "In order to pack efficiently, we will use a histogram-based algorithm. The shortest-pack-first histogram packing (SPFHP) was presented in the Graphcore blog post [introducing Packed BERT for a training speedup in natural language processing](https://www.graphcore.ai/posts/introducing-packed-bert-for-2x-faster-training-in-natural-language-processing). We have adapted the [code](https://github.com/graphcore/tutorials/tree/master/blogs_code/packedBERT) from the blog post for this notebook. The full process of packing the dataset consists of four steps:\n", "\n", "1. Create a histogram of the sequence lengths of the dataset.\n", "2. Generate the \"strategy\" for the dataset using one of the state-of-the-art packing algorithms. The strategy maps out the order and indices of the sequences that need to be packed together.\n", "3. Use this strategy to create the actual dataset, concatenating the tokenized features together for each column in the dataset, including the labels.\n", "4. Finally, pass these new columns into a custom PyTorch dataset, ready to be passed to the PopTorch dataloader!\n", "\n", "These steps have been simplified through the easy-to-use `utils.packing` package available in Graphcore Optimum. You can simply generate the packed dataset after the usual tokenization and preprocessing by passing all necessary packing configuration to the `PackedDatasetCreator` class, and generate the ready-to-use PyTorch dataset with `.create()`.\n", "\n", "Within the function, there are some column names used by default. The expected default columns for text classification include:\n", "* `input_ids`\n", "* `attention_mask`\n", "* `token_type_ids`\n", "* `labels`\n", "\n", "These should all be generated automatically when tokenizing any classification dataset for BERT. However, the labels key, as it is not encoded, may have a different name. For this dataset, the column key for the labels for this dataset is `label`, since the dataset creator expects `labels`, we can pass this to the argument `custom_label_key`, so the class can find our labels. \n", "\n", "The `PackedDatasetCreator` requires different instantiations for different datasets, so it must be called separately for each of our dataset splits. We can set either `training`, `validation` or `inference` to `True` as needed." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from utils.packing.dataset_creator import PackedDatasetCreator\n", "\n", "train_data_packer = PackedDatasetCreator(\n", " tokenized_dataset = encoded_dataset['train'],\n", " max_sequence_length = max_seq_length,\n", " max_sequences_per_pack = max_seq_per_pack,\n", " training = True,\n", " num_labels = num_labels,\n", " problem_type = problem_type,\n", " algorithm = 'SPFHP',\n", " custom_label_key = 'labels'\n", ")\n", "\n", "val_data_packer = PackedDatasetCreator(\n", " tokenized_dataset = encoded_dataset['validation'],\n", " max_sequence_length = max_seq_length,\n", " max_sequences_per_pack = max_seq_per_pack,\n", " validation = True,\n", " num_labels = num_labels,\n", " problem_type = problem_type,\n", " algorithm = 'SPFHP',\n", " custom_label_key = 'labels'\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "This will create the strategy and initialise the necessary parameters for packing the dataset. We can see that the ideal speed-up we have achieved is approximately 5.7x the original dataset, which corresponds directly to the average packing factor: the average number of sequences within one pack.\n", "\n", "The `PackedDatasetCreator` class also has some other features we do not use here for training, such as `pad_to_global_batch_size`, a feature useful for performing batched inference on a large samples when we do not want to lose any of the samples. When creating data iterators using the `poptorch.Dataloader`, it applies 'vertical' padding to the dataset, adding filler rows to bring the dataset up to a value divisible by the global batch size, and allows for the largest possible batch sizes to be used without any loss of data.\n", "\n", "You can also view the histogram generated in the first step of the packing process, to observe whether the distribution of sequence lengths in the dataset will benefit from packing - as a general rule, as long as the average length of the sequences in the dataset is 50% or less of the maximum sequence length, packing will offer at least a 2x throughput benefit, in other words: `throughput_increase ≈ max_seq_len/mean_seq_len`\n", "\n", "Many datasets have distributions with much smaller average lengths, and will benefit much more. We can easily observe this distribution by retrieving and plotting the histogram from the data class:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "train_histogram = train_data_packer.histogram\n", "\n", "plt.hist(train_histogram, bins = [k for k in range(0,max_seq_length,10)]) \n", "plt.title(\"Sequence length histogram\") \n", "plt.xlabel('Sequence lengths')\n", "plt.ylabel('Frequency')\n", "plt.show()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Now we need to create the actual packed dataset (step 3 in the packing process outlined above).\n", "\n", "In this stage, we take the strategy for mapping the sequences by size into packs that were generated by the packing algorithm, and use this to extract the sequences from the tokenized dataset, inserting them into packs for each column in the dataset. Any remaining space in a pack after the sequences have been concatenated is padded to bring all sequences up to the maximum sequence length.\n", "\n", "Some key features unique to packed datasets are worth mentioning here:\n", "\n", "- The specific attention mask (`attention_mask`) that is generated contains a unique index for each sequence of the pack and `0` for the remaining padding tokens. This, essentially, tells the model where to look from the perspective of a single token, ignoring any encoded information (such as a different sequence) that is not relevant to that token.\n", " - Example of 3 sequences: `attention_mask = [1,1,1,1,1,1,2,2,2,2,2,3,3,3,3,3,0,...,0,1,2,3]`\n", "\n", "\n", "- The [CLS] tokens of each sequence must be moved to the end of the pack.\n", " - For instance: `[CLS,a,b,c] + [CLS, d,e,f] + [CLS, g,h,i] -> [a,b,c,d,e,f,g,h,i,...,CLS,CLS,CLS]`\n", " \n", "\n", "- `position_ids` for a pack contains the concatenated `position_ids` of each sequences \n", " - For instance given 3 sequences: `[0,1,2,3,4] + [0,1,2,3] + [0,1,2] -> [1,2,3,4,1,2,3,1,2,...,0,0,0]` (note: the CLS tokens' position ID '0' are also moved to the end of the pack)\n", " \n", "- `labels` and `token_type_ids` are also packed to correspond to the `input_ids` pack.\n", "\n", "\n", "To create a dataloader-ready packed dataset, all you need to do is call the `create()` method:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "packed_train_dataset = train_data_packer.create()\n", "packed_val_dataset = val_data_packer.create()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let's visualize one sample of the new `packed_train_dataset`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "packed_train_dataset[133]" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Fine-tuning the model" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Now that our data is ready, we can download the pre-trained model and fine-tune it." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### Implementing Packed BERT\n", "\n", "A few model modifications are required to make packing work with BERT.\n", "We extend the existing `BertForSequenceClassification` class to `PipelinedPackedBertForSequenceClassification` which incorporates the required changes to the pooler and the model output. The crux of these changes is to modify the generic sequence classification model to handle unpacking multiple sequences in the output stage, treating them as a larger batch size for classification, as well as masking any padding created by packing.\n", "\n", "First let's load a default BERT configuration using `AutoConfig`. The config includes a new parameter we must set, `max_sequences_per_pack`, which informs the model of the maximum number of sequences it will need to unpack in the model output. It also allows us to clearly define the `num_labels` and `problem_type` for this model.\n", "\n", "It is essential we define the problem type here, as switching between the methods used by different types of classification requires that it be defined within the custom model." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from transformers import AutoConfig\n", "\n", "config = AutoConfig.from_pretrained(model_checkpoint)\n", "config.max_sequences_per_pack = max_seq_per_pack\n", "config.num_labels = num_labels\n", "config.problem_type = problem_type" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Now we can instantiate the model class with the config, loading the weights from the model checkpoint." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true, "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "torch.manual_seed(43)\n", "np.random.seed(43)\n", "\n", "from models.modeling_bert_packed import PipelinedPackedBertForSequenceClassification\n", "\n", "\n", "model = PipelinedPackedBertForSequenceClassification(config).from_pretrained(\n", " model_checkpoint, config=config)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "The warning tells us we are throwing away some weights and randomly initializing others. This is normal in this case, because we are removing the head used to pre-train the model on a masked language modelling objective and replacing it with a new head for sequence classification, which we don't have pre-trained weights for, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do.\n", "\n", "We can first test the model on a CPU and observe that the output logits have the size `[batch_size, max_seq_per_pack, 2] = [1, 6, 28]` with this notebook's default values, and the 28 labels for the dataset. The logits are reshaped into this form in the model output, to be the same shape as the labels, for ease of post-processing." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "# test the model on CPU\n", "from transformers.data.data_collator import default_data_collator\n", "\n", "loader = torch.utils.data.DataLoader(packed_train_dataset,\n", " batch_size=1,\n", " shuffle=True,\n", " drop_last=True,\n", " collate_fn=default_data_collator)\n", "data = next(iter(loader))\n", "labels = data['labels']\n", "\n", "print('labels: ', labels.shape)\n", "o = model(**data)\n", "print('outputs (loss, logits): ', o[0], o[1].shape)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Now, let's prepare the model for an IPU." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "First, we set the model to half-precision:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true, "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "model.half()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "For validation, we need to define a function to compute the metrics from the predictions, which will use `metric` which we loaded earlier. Preprocessing here involves a step to mask the labels and predictions we are not using, set to a `-100` value when creating the dataset, with a boolean mask. Then, the predictions are passed into a `softmax` function to determine the probabilities of each class, as this is a multi-label task. \n", "\n", "These predictions and labels are passed into the metric function to compute the accuracy during evaluation." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "model_name = model_checkpoint.split(\"/\")[-1]\n", "from scipy.special import softmax\n", "from tqdm import tqdm\n", "def compute_metrics(eval_pred):\n", " predictions, labels = eval_pred\n", " \n", " labels = labels.reshape(-1, labels.shape[-1])\n", " predictions = predictions.reshape(-1, predictions.shape[-1])\n", " \n", " # Remove the padding labels\n", " mask = (labels != -100)[:,0]\n", " \n", " labels = labels[mask,:]\n", " predictions = predictions[mask,:]\n", " pred_scores = softmax(predictions.astype(\"float32\"), axis=1) \n", "\n", " auc = metric.compute(\n", " prediction_scores=pred_scores, references=labels, multi_class=\"ovr\"\n", " )[\"roc_auc\"]\n", "\n", " return {\"roc_auc\": auc}" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Next, we need to define `IPUConfig`, which is a class that specifies attributes and configuration parameters to compile and put the model on the device. We initialize it with a config name or path, which we set earlier. Then we use it to set the mode attribute `model.ipu_config` " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from optimum.graphcore import IPUConfig, IPUTrainer, IPUTrainingArguments\n", "\n", "ipu_config = IPUConfig.from_pretrained(\n", " ipu_config_name,\n", " executable_cache_dir = executable_cache_dir,\n", " gradient_accumulation_steps=gradient_accumulation_steps,\n", " device_iterations = device_iterations,\n", " replication_factor=1,\n", " inference_device_iterations = 64,\n", " inference_replication_factor = 1\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "`IPUTrainingArguments` define any custom parameter modification we want to do, such as the initial learning rate for the model. It also allows other options, such as dataloader parameters, micro batch sizes and an automatic push to the Hugging Face Hub (if credentials were set up earlier) to happen at given intervals.\n", "\n", "These arguments are passed to `IPUTrainer` which wraps the model training and evaluation process into a simple single-line process, doing all of the heavy lifting for us, for example training and evaluation loops, device assignment, optimiser definition and data-loading.\n", "\n", "Note that only some arbitrary hyperparameter tuning was performed for this task. Other tasks and datasets may require further tuning to get the most optimal results." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true, "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from transformers import default_data_collator\n", "metric_name = \"roc_auc\"\n", "\n", "args = IPUTrainingArguments(\n", " \"./\"+f\"{model_name}-{model_task}\",\n", " per_device_train_batch_size=2,\n", " per_device_eval_batch_size=2,\n", " num_train_epochs=5,\n", " learning_rate=2e-4,\n", " adam_epsilon=1e-6,\n", " loss_scaling=16.0,\n", " warmup_ratio=0.1,\n", " weight_decay=0,\n", " lr_scheduler_type = \"cosine\",\n", " metric_for_best_model=metric_name,\n", " dataloader_drop_last=True,\n", " dataloader_mode=\"async_rebatched\",\n", " logging_steps=1,\n", " n_ipu=n_ipu,\n", " gradient_accumulation_steps=gradient_accumulation_steps,\n", " push_to_hub=False \n", ")\n", "\n", "trainer = IPUTrainer(\n", " model,\n", " ipu_config,\n", " args,\n", " train_dataset=packed_train_dataset,\n", " eval_dataset=packed_val_dataset,\n", " data_collator=default_data_collator,\n", " compute_metrics=compute_metrics\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Then, to train the model we can simply call the `train()` method:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "trainer.train()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "***About the performance:*** `IPUTrainer` doesn't take into account that we have packed data samples when computing the speed metrics. It treats a \"sample\" as a single input to the model, which means one **pack**.\n", "\n", "So the actual throughput estimation can be obtained by multiplying `samples_per_second` by the average packing factor (the average number of samples per pack) of the dataset. These were obtained in the `packing_algorithm` section: `5.68` for the `go-emotions` training set and `5.83` for validation set.\n", "\n", "Next, we can evaluate the model by simply calling the `evaluate()` method:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "trainer.evaluate()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "You can now upload the result of the training to the Hugging Face Hub if you successfully authenticated at the beginning of this notebook. Simply uncomment and execute the following cell: " ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true, "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "#trainer.push_to_hub()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "You can also save the model locally:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "trainer.save_model(\"./\"+f\"{model_name}-{model_task}\")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "You have now successfully fine-tuned and evaluated your speed-optimised model for text classification using packing!" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Fast batched inference" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "Packing can also be used for inference, particularly for performing inference for workloads. This section demonstrates how to perform faster, batched inference with a large number of samples using a super-easy custom pipeline which batches and packs your input data, performs inference and returns post-processed predictions. \n", "\n", "For the pipeline, we need to import it, and initialise a few essential parameters.\n", "\n", "The `model` is the model checkpoint and we are going to use the locally saved checkpoint generated from training `go_emotions`. `executable_cache_dir`, `problem_type`, `max_seq_length` must be specified. To return predictions organised by class names, the class names for your output must be passed to `label_categories`. \n", "\n", "The pipeline will automatically determine your model's IPU config, given that the checkpoint was trained using Optimum Graphcore, which will be the case for the model fine-tuned in this notebook.\n", "\n", "In this example, we pre-load `IPUConfig` and modify some of the default parameters to get the best performance out of inference and leverage the benefits of IPU parallelism. The micro-batch size can also be specified, for which the default is 1.\n", "\n", "When training, the packing factor affects the convergence in the same way as a large increase in batch size would do. However, for inference, we are free to use a bigger packing factor to speed it up. Let's try it with `max_seq_per_pack = 12`.\n", "\n", "**Note:** Packing brings huge benefits for performing inference on large amounts of data. For small scale inference tasks, such as those which more suit sequential inference on a single un-batched input, the generic Optimum Graphcore `TextClassificationPipeline` class may be preferred. This won't affect fine-tuning and the weights generated from fine-tuning using packing will work just the same!" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let's list the class names for the GoEmotions dataset." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "class_names = [\n", " \"admiration\",\n", " \"amusement\",\n", " \"anger\",\n", " \"annoyance\",\n", " \"approval\",\n", " \"caring\",\n", " \"confusion\",\n", " \"curiosity\",\n", " \"desire\",\n", " \"disappointment\",\n", " \"disapproval\",\n", " \"disgust\",\n", " \"embarrassment\",\n", " \"excitement\",\n", " \"fear\",\n", " \"gratitude\",\n", " \"grief\",\n", " \"joy\",\n", " \"love\",\n", " \"nervousness\",\n", " \"optimism\",\n", " \"pride\",\n", " \"realization\",\n", " \"relief\",\n", " \"remorse\",\n", " \"sadness\",\n", " \"surprise\",\n", " \"neutral\",\n", "]" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let's initialise the `PackedBertTextClassificationPipeline`." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "from pipeline.packed_bert import PackedBertTextClassificationPipeline\n", "\n", "from optimum.graphcore import IPUConfig\n", "\n", "model = \"./\"+f\"{model_name}-{model_task}\"\n", "# model = 'your_username/{model_name}-{task}' # to load from Hugging Face Hub\n", "\n", "inference_boosted_ipu_config = IPUConfig.from_pretrained(model, \n", " inference_device_iterations=32,\n", " inference_replication_factor=4,\n", " ipus_per_replica=1,\n", " layers_per_ipu=[12]\n", " )\n", "\n", "pipeline = PackedBertTextClassificationPipeline(\n", " model = model,\n", " executable_cache_dir = executable_cache_dir,\n", " problem_type='multi_label_classification',\n", " max_seq_per_pack=12,\n", " max_seq_length=max_seq_length,\n", " ipu_config=inference_boosted_ipu_config,\n", " micro_batch_size=8,\n", " label_categories=class_names\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "The pipeline expects a **list of strings** directly passed to it. There is no need to tokenize, preprocess, pack or post-process the data to use the inference pipeline.\n", "\n", "As a test, we can load the entire `sst2` dataset and perform packed inference using `.predict()` on the text column to generate predictions. \n", "\n", "Datasets with multiple sentences can simply be passed as `predict(,)`" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [], "vscode": { "languageId": "python" } }, "outputs": [], "source": [ "import datasets\n", "dataset = datasets.load_dataset('go_emotions','simplified')\n", "preds = pipeline.predict(dataset['train']['text'])\n", "\n", "print(preds.keys())\n", "print(f\"Number of predictions: {len(preds['predictions'])}\")\n", "print(f\"Preprocessing time: {preds['preprocessing_time']}s\")\n", "print(f\"Postprocessing time: {preds['postprocessing_time']}s\")\n", "print(f\"Throughput: {preds['throughput']} samples/s\")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "There is minimal overhead from tokenizing and packing the dataset, but the speed benefits are evident. After increasing the maximum sequences to 12, we can observe a much higher packing factor of 9.14.\n", "\n", "Running the above pipeline, we achieve a throughput approximately 45000 samples per second, demonstrating the huge time benefit you can achieve by using packing!" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Next steps\n", "\n", "You may want to try out the following notebooks using BERT with packing for:\n", "\n", "* Single-label text classification - `packedBERT_single_label_text_classification.ipynb`\n", "* Question answering - `packedBERT_question_answering.ipynb`\n", "\n", "Also, check out the full list of [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to get more of a feel for how IPUs perform on other tasks." ] } ], "metadata": { "colab": { "name": "Text Classification on GLUE", "provenance": [] }, "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "vscode": { "interpreter": { "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6" } } }, "nbformat": 4, "nbformat_minor": 4 }