Transformer Head Pruner (#3884)

docs/en_US/Compression/CompressionReference.rst

@@ -91,6 +91,8 @@ Pruners
 ..  autoclass:: nni.algorithms.compression.pytorch.pruning.lottery_ticket.LotteryTicketPruner
     :members:
 
+..  autoclass:: nni.algorithms.compression.pytorch.pruning.transformer_pruner.TransformerHeadPruner
+    :members:
 
 Quantizers
 ^^^^^^^^^^

docs/en_US/Compression/Overview.rst

@@ -35,7 +35,7 @@ The algorithms include pruning algorithms and quantization algorithms.
 Pruning Algorithms
 ^^^^^^^^^^^^^^^^^^
 
-Pruning algorithms compress the original network by removing redundant weights or channels of layers, which can reduce model complexity and mitigate the over-fitting issue.
+Pruning algorithms compress the original network by removing redundant weights or channels of layers, which can reduce model complexity and mitigate the over-fitting issue.
 
 .. list-table::
    :header-rows: 1
@@ -73,6 +73,8 @@ Pruning algorithms compress the original network by removing redundant weights o
      - Automatic pruning by iteratively call SimulatedAnnealing Pruner and ADMM Pruner `Reference Paper <https://arxiv.org/abs/1907.03141>`__
    * - `AMC Pruner <../Compression/Pruner.rst#amc-pruner>`__
      - AMC: AutoML for Model Compression and Acceleration on Mobile Devices `Reference Paper <https://arxiv.org/pdf/1802.03494.pdf>`__
+   * - `Transformer Head Pruner <../Compression/Pruner.rst#transformer-head-pruner>`__
+     - Pruning attention heads from transformer models either in one shot or iteratively.
 
 
 You can refer to this `benchmark <../CommunitySharings/ModelCompressionComparison.rst>`__ for the performance of these pruners on some benchmark problems.

docs/en_US/Compression/Pruner.rst

@@ -28,6 +28,7 @@ We provide several pruning algorithms that support fine-grained weight pruning a
 **Others**
 
 * `Lottery Ticket Hypothesis <#lottery-ticket-hypothesis>`__
+* `Transformer Head Pruner <#transformer-head-pruner>`__
 
 Level Pruner
 ------------
@@ -724,3 +725,95 @@ User configuration for Sensitivity Pruner
 **PyTorch**
 
 ..  autoclass:: nni.algorithms.compression.pytorch.pruning.SensitivityPruner
+
+Transformer Head Pruner
+-----------------------
+
+Transformer Head Pruner is a tool designed for pruning attention heads from the models belonging to the `Transformer family <https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf>`__. The following image from `Efficient Transformers: A Survey <https://arxiv.org/pdf/2009.06732.pdf>`__ gives a good overview the general structure of the Transformer.
+
+.. image:: ../../img/transformer_structure.png
+   :target: ../../img/transformer_structure.png
+   :alt: 
+
+Typically, each attention layer in the Transformer models consists of four weights: three projection matrices for query, key, value, and an output projection matrix. The outputs of the former three matrices contains the projected results for all heads. Normally, the results are then reshaped so that each head performs that attention computation independently. The final results are concatenated back before fed into the output projection. Therefore, when an attention head is pruned, the same weights corresponding to that heads in the three projection matrices are pruned. Also, the weights in the output projection corresponding to the head's output are pruned. In our implementation, we calculate and apply masks to the four matrices together.
+
+Note: currently, the pruner can only handle models with projection weights written as separate ``Linear`` modules, i.e., it expects four ``Linear`` modules corresponding to query, key, value, and an output projections. Therefore, in the ``config_list``, you should either write ``['Linear']`` for the ``op_types`` field, or write names corresponding to ``Linear`` modules for the ``op_names`` field.
+
+The pruner implements the following algorithm:
+
+.. code-block:: bash
+
+    Repeat for each pruning iteration (1 for one-shot pruning):
+       1. Calculate importance scores for each head in each specified layer using a specific criterion.
+       2. Sort heads locally or globally, and prune out some heads with lowest scores. The number of pruned heads is determined according to the sparsity specified in the config.
+       3. If the specified pruning iteration is larger than 1 (iterative pruning), finetune the model for a while before the next pruning iteration.
+
+Currently, the following head sorting criteria are supported:
+
+    * "l1_weight": rank heads by the L1-norm of weights of the query, key, and value projection matrices.
+    * "l2_weight": rank heads by the L2-norm of weights of the query, key, and value projection matrices.
+    * "l1_activation": rank heads by the L1-norm of their attention computation output.
+    * "l2_activation": rank heads by the L2-norm of their attention computation output.
+    * "taylorfo": rank heads by l1 norm of the output of attention computation * gradient for this output. Check more details in `this paper <https://arxiv.org/abs/1905.10650>`__ and `this one <https://arxiv.org/abs/1611.06440>`__.
+
+We support local sorting (i.e., sorting heads within a layer) and global sorting (sorting all heads together), and you can control by setting the ``global_sort`` parameter. Note that if ``global_sort=True`` is passed, all weights must have the same sparsity in the config list. However, this does not mean that each layer will be prune to the same sparsity as specified. This sparsity value will be interpreted as a global sparsity, and each layer is likely to have different sparsity after pruning by global sort.
+
+In our implementation, we support two ways to group the four weights in the same layer together. You can either pass a nested list containing the names of these modules as the pruner's initialization parameters (usage below), or simply pass a dummy input and the pruner will run ``torch.jit.trace`` to group the weights (experimental feature). However, if you would like to assign different sparsity to each layer, you can only use the first option, i.e., passing names of the weights to the pruner (see usage below). Also note that weights belonging to the same layer must have the same sparsity.
+
+In addition to the following usage guide, we provide a more detailed example of pruning BERT for tasks from the GLUE benchmark. Please find it in this :githublink:`page <examples/model_compress/pruning/transformers>`.
+
+Usage
+^^^^^
+
+Suppose we want to prune a BERT with Huggingface implementation, which has the following architecture (obtained by calling ``print(model)``). Note that we only show the first layer of the repeated layers in the encoder's ``ModuleList layer``. 
+
+.. image:: ../../img/huggingface_bert_architecture.png
+   :target: ../../img/huggingface_bert_architecture.png
+   :alt: 
+
+**Usage Example: one-shot pruning, assigning sparsity 0.5 to the first six layers and sparsity 0.25 to the last six layers (PyTorch code)**. Note that
+
+* Here we specify ``op_names`` in the config list to assign different sparsity to different layers.
+* Meanwhile, we pass ``attention_name_groups`` to the pruner so that the pruner may group together the weights belonging to the same attention layer.
+* Since in this example we want to do one-shot pruning, the ``num_iterations`` parameter is set to 1, and the parameter ``epochs_per_iteration`` is ignored. If you would like to do iterative pruning instead, you can set the ``num_iterations`` parameter to the number of pruning iterations, and the ``epochs_per_iteration`` parameter to the number of finetuning epochs between two iterations.
+* The arguments ``trainer`` and ``optimizer`` are only used when we want to do iterative pruning, or the ranking criterion is ``taylorfo``. Here these two parameters are ignored by the pruner.
+* The argument ``forward_runner`` is only used when the ranking criterion is ``l1_activation`` or ``l2_activation``. Here this parameter is ignored by the pruner.
+
+.. code-block:: python
+
+   from nni.algorithms.compression.pytorch.pruning import TransformerHeadPruner
+   attention_name_groups = list(zip(["encoder.layer.{}.attention.self.query".format(i) for i in range(12)],
+                                    ["encoder.layer.{}.attention.self.key".format(i) for i in range(12)],
+                                    ["encoder.layer.{}.attention.self.value".format(i) for i in range(12)],
+                                    ["encoder.layer.{}.attention.output.dense".format(i) for i in range(12)]))
+   kwargs = {"ranking_criterion": "l1_weight",
+             "global_sort": False,
+             "num_iterations": 1,
+             "epochs_per_iteration": 1,    # this is ignored when num_iterations = 1
+             "head_hidden_dim": 64,
+             "attention_name_groups": attention_name_groups,
+             "trainer": trainer,
+             "optimizer": optimizer,
+             "forward_runner": forward_runner
+             }
+    config_list = [{
+        "sparsity": 0.5,
+        "op_types": ["Linear"],
+        "op_names": [x for layer in attention_name_groups[:6] for x in layer]      # first six layers
+    },
+    {
+        "sparsity": 0.25,
+        "op_types": ["Linear"],
+        "op_names": [x for layer in attention_name_groups[6:] for x in layer]      # last six layers
+    }
+    ]
+   pruner = TransformerHeadPruner(model, config_list, **kwargs)
+   pruner.compress()
+
+   
+User configuration for Transformer Head Pruner
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+**PyTorch**
+
+..  autoclass:: nni.algorithms.compression.pytorch.pruning.TransformerHeadPruner

examples/model_compress/pruning/transformers/run.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+# Usage: ./run.sh gpu_id glue_task
+
+export CUDA_VISIBLE_DEVICES=$1
+TASK_NAME=$2                                  # "cola", "sst2", "mrpc", "stsb", "qqp", "mnli", "qnli", "rte", "wnli"
+PRETRAINED_MODEL="bert-base-uncased"          # "distilbert-base-uncased", "roberta-base", "bert-base-cased", ...
+
+# parameters for pruning
+SPARSITY=0.5
+RANKING_CRITERION=l1_weight                   # "l1_weight", "l2_weight", "l1_activation", "l2_activation", "taylorfo"
+NUM_ITERATIONS=1                              # 1 for one-shot pruning
+EPOCHS_PER_ITERATION=1
+
+# other training parameters, no need to change
+MAX_LENGTH=128
+BATCH_SIZE=32
+LR=2e-5
+N_EPOCHS=3
+
+time=$(date "+%Y%m%d%H%M%S")
+OUTDIR="models_${PRETRAINED_MODEL}_${TASK_NAME}_$time/"
+
+TASK_LIST=("cola" "sst2" "mrpc" "stsb" "qqp" "mnli" "qnli" "rte" "wnli")
+if [[ ${TASK_LIST[*]} =~ (^|[[:space:]])$TASK_NAME($|[[:space:]]) ]]; then
+    mkdir $OUTDIR
+    python transformer_pruning.py \
+	   --sparsity $SPARSITY \
+	   --ranking_criterion $RANKING_CRITERION \
+	   --num_iterations $NUM_ITERATIONS \
+	   --epochs_per_iteration $EPOCHS_PER_ITERATION \
+	   --speed_up \
+	   --model_name $PRETRAINED_MODEL \
+	   --task_name $TASK_NAME \
+	   --max_length $MAX_LENGTH \
+	   --batch_size $BATCH_SIZE \
+	   --learning_rate $LR \
+	   --num_train_epochs $N_EPOCHS \
+	   --output_dir $OUTDIR \
+	   2>&1 | tee "$OUTDIR/output.log"
+else
+    echo "Unsupported task $TASK_NAME."
+fi