Transformer Head Pruner

[xiaowu0162]

This pr adds a pruner for pruning attention heads in transformers.
To-do:

• basic: module matching, name-based weight grouping, group-aware maskers using weight norm as criteria
• graph-based weight grouping
• maskers relying on activation
• maskers relying on gradient
• global sort in maskers
• support iterative pruning from scratch / integration with pruning scheduler for iterative features
• example
• docs

[zheng-ningxin]

Is there a better way to get the index of the attention head? the first integer may be not strong.

[xiaowu0162]

This layer_idx is the layer index of the BERT encoder. Here I include these lines of code only to show the user how they may take advantage of the pruned_heads dict inside pruner to get the pruned heads for each group, and then match each group to the original layer, and finally call the built-in transformers _prune_heads() function to do model speedup. This is meant to be a temporary workaround before we can properly handle speedup for transformers.

[xiaowu0162]

If our speedup code after refactor can handle transformer, then I will replace these lines with our speedup methods (maybe in a separate pr)

[xiaowu0162]

Since the users are aware of the naming of their own model to prune, I think they can also use their own rules to match layers to groups

[zheng-ningxin]

Feel like that the TransformerHeadPruner is too heavy. I prefer to locate this pruner as a one-shot pruner, which means we do not need handle with the num_iteration, optimizer, trainer, criterion, things. That's much clearer. All those finetuning related things we can offload to the outer search algorithms. We can discuss with Quanlu @QuanluZhang .

[xiaowu0162]

Yes, we can further discuss on that. One challenge is that this does not fit well in our current compression V1 framework (since the current iterative and dependency aware pruner are limited to convolutions), and compression V2 is not ready yet. My initial thought was to first integrate all these logic in one pruner (because of empirically good performance compared to one-shot pruning), and then factor out when compression V2 is ready.