ACL'22: Structured Pruning Learns Compact and Accurate Models

Hi @xiamengzhou , many thanks to your contribution. I have small questions in your paper, in your paper you said that

FNN pruning introduce a Zint

And in your paper there is a Eq, but what is diag, why do we have to put Zint into a diagonal matrix? Do diag(Zint) is df*df size?

And you also says that

Coarse-grained and Fine- grained units (§3.1) with a layerwise distillation objective transferring knowledge from unpruned to pruned models (§3.2)

However, distilling intermediate layers during the pruning process is challenging as the model struc- ture changes throughout training. （previous method）

So are we pruning a student model during distillation？

Many thanks!!

CoFi is a great work which may benefit the research in related areas.

However, I have found the numbers of the task performance on other sparsities are not available. Could you please provide these numbers in detail?

Besides, metrics besides accuracy scores on GLUE would also be appreciated.

I do not understand how this loss works ---- since $\lambda_1$ and $\lambda_2$ are 0 as default, I find that sometimes the loss maybe a negative number sometimes.

Hello,

I am trying to run your codebase. I am having some issues however:

1. Set of python requirements cannot be installed due to incompatibilities. Are these requirements strict or can they be relaxed?

2. After relaxing the above evaluation runs fine, but training requires a --distillation_path. Could you provide an example on how to use this argument?

3. To overcome 2, I set variable additional_args.do_distill to False. this results in an epoch being trained but crashing at the end. Model loss succesfully reduces but reg loss and lag loss is 0. The error at the end is a failure in assertion: " assert "head" in self.types" in the l0 module

Could you help me or provide pointers on resolving the above?

Thank you

The original papers mentioned: Specifically, let T denote a set of teacher layers that we use to distill knowledge to the student model.'' And the code in trainer provides[2, 5, 8, 11]'' only, which is part of settings in Appendix. Any suggestions of selection of such teacher layer sets for distillation,? 4 layers at most? which 4 layers are proper? how do we specify task-aware settings? i.e., There are 12 layers for Students, why we only choose to select from given 4 layers? How about 5, 6, 12 layers for T,? I think it is critical for reproduce results, where I barely reproduce any results to match the reported scores now?

Great job. However, when I train to the 3rd epoch in the QNLI task, I encounter the following problem, but the CoLA or Squad tasks do not encounter this problem. Do you have any suggestions? I will be very grateful!

The error may appear in the following code block in the file: https://github.com/princeton-nlp/CoFiPruning/blob/main/trainer/trainer.py,Lines 680-685

lagrangian_loss = None
if self.start_prune:
        lagrangian_loss, _, _ = \
                 self.l0_module.lagrangian_regularization(
                        self.global_step - self.prepruning_finetune_steps)
        loss += lagrangian_loss

Hello,

In the following line: https://github.com/princeton-nlp/CoFiPruning/blob/022847ae88f49fa7b8fc58f9c0613492fd1230cc/trainer/trainer.py#L599

existing_layers tensor is in cpu and the result of indexes<last_aligned_layer is in gpu. This throws an error as a result

Is this a bug? maybe first move existing_layers to gpu?

Hello, Thank you for providing code. But I have a question on how to reproduce the results the 95% sparsity on MNLI with the following commands:

TASK=MNLI
SUFFIX=sparsity0.95
EX_CATE=CoFi
PRUNING_TYPE=structured_heads+structured_mlp+hidden+layer
SPARSITY=0.95
DISTILL_LAYER_LOSS_ALPHA=0.9
DISTILL_CE_LOSS_ALPHA=0.1
LAYER_DISTILL_VERSION=4
DISTILLATION_PATH=dynabert/MNLI
CUDA_VISIBLE_DEVICES=1 bash scripts/run_CoFi.sh $TASK $SUFFIX $EX_CATE $PRUNING_TYPE $SPARSITY $DISTILLATION_PATH $DISTILL_LAYER_LOSS_ALPHA $DISTILL_CE_LOSS_ALPHA $LAYER_DISTILL_VERSION

And I get following results with accuracy 78.20 on MNLI:

wandb: Run history:
wandb:                   eval/loss ▃▁▂▂▂▃██▆▆▅▅▅▅▆▅▅▄▄▅▄▅▅▅▄▄▄▄▄▄▄▄▄▄▄▄▅▄▄▄
wandb:              train/accuracy ▆█▇██▇▁▁▃▄▄▄▄▅▄▅▅▅▅▅▅▅▅▅▅▅▅▅▅▆▆▆▆▆▆▆▆▆▆▆
wandb:     train/expected_sparsity ▁▃▄▆████████████████████████████████████
wandb:           train/global_step ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb:           train/hidden_dims █████▁▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃
wandb:              train/lag_loss ▆▆▇▆▆█▁▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆
wandb:         train/learning_rate █████▇▇▇▇▇▆▆▆▆▆▆▅▅▅▅▅▄▄▄▄▄▃▃▃▃▃▃▂▂▂▂▂▁▁▁
wandb:                  train/loss ▂▁▆▂▂▇▃█▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▆▅▆▆▆
wandb: train/pruned_model_sparsity ▁▃▄▆████████████████████████████████████
wandb:         train/pruned_params ▁▃▄▆████████████████████████████████████
wandb:              train/reg_loss ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb:      train/remaining_params █▆▅▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb:       train/target_sparsity ▁▂▄▆████████████████████████████████████
wandb:
wandb: Run summary:
wandb:                   eval/loss 0.66644
wandb:              train/accuracy 0.78197
wandb:     train/expected_sparsity 0.94999
wandb:           train/global_step 0
wandb:           train/hidden_dims 764
wandb:              train/lag_loss 1e-05
wandb:         train/learning_rate 0.0
wandb:                  train/loss 0.40625
wandb: train/pruned_model_sparsity 0.95561
wandb:         train/pruned_params 81243440
wandb:              train/reg_loss 0.0
wandb:      train/remaining_params 3774160
wandb:       train/target_sparsity 0.95

By the way, I found some issues during reproducing:

1. In evaluation.py:77, datasets["validation"] should be datasets["validation_matched"] for MNLI.
2. Label_map in MNLI, dynabert and princeton-nlp/CoFi-MNLI-s95 use different label map compare to MNLI in datasets.load_dataset. And directly evaluating with python evaluation.py MNLI princeton-nlp/CoFi-MNLI-s95 will get wrong result.
3. Pruning is unavailable for trained models in evaulation.py. For example, model is not purnned according to zs.pt with python evaluation.py MNLI ./out/MNLI/CoFi/MNLI_sparsity0.95.

Hi! I am trying to apply CoFi pruning to my own model, and I noticed that there might exist some edge cases where removing the already-pruned parts in my model will cause some changes in the outputs. I think this will happen when all the dims of the intermediate layer are removed.

I found that when intermediate_zs are all zero, the intermediate.dense in the pruned model is set to None https://github.com/princeton-nlp/CoFiPruning/blob/5423094e7b318806462f2a7bdca5384d078e5eed/utils/cofi_utils.py#L229-L231 , and the FFN parts will then be skipped https://github.com/princeton-nlp/CoFiPruning/blob/5423094e7b318806462f2a7bdca5384d078e5eed/models/modeling_bert.py#L364-L365

But before pruning, intermediate.dense is not None, and these zero outputs will still pass through CoFiBertOutput.dense which add a bias to the output https://github.com/princeton-nlp/CoFiPruning/blob/5423094e7b318806462f2a7bdca5384d078e5eed/models/modeling_bert.py#L562-L566 , so the FFN parts are not skipped.

Should I change some part of my code to skip the FFN parts when intermediate_zs are all zero during training?

in the file: https://github.com/princeton-nlp/CoFiPruning/blob/main/trainer/trainer.py line 279 sepcifies following statement :

 if self.start_prune:
    zs = self.l0_module.forward(training=True)
    self.fill_inputs_with_zs(zs, inputs)

only when this runs, we can get the gradient for the params in self.l0_optimizer. Only when the condiction satisfied as below (line 268):

  if self.prepruning_finetune_steps > 0 and self.global_step == self.prepruning_finetune_steps:
      self.start_prune = True

However, line 301 just directly update the params without checking whether the grads are ready:

  if self.l0_module is not None and self.l0_optimizer is not None:
      self.l0_optimizer.step()
      self.lagrangian_optimizer.step()

therefore, the adamw yields bugs for beta1/beta2 referred before define in its step method. As the the grad of the params are all None, the adamw implementation will skip define the hyper-params via the self.group dict.

Hello,@xiamengzhou !The result on the Squad task dataset is 79.74, which is quite different from the result (82.6) in the paper. Can you further announce the detailed parameters? The teacher model F1 value is 88.43. I will be very grateful!

Hi, First of all, thanks a lot for open-sourcing your code and models!

I've been trying to use your code to generate predictions with CoFi models (with --do_predict on for example test-split of GLUE tasks) but unfortunately the prediction loop always fails with CUDA OOM exception (even on the 80GB A100 GPU). Could you also please try and let me know if I did something wrong?

hello,@xiamengzhou! When i use your script to finetune the pruned model, there is an issue. But i hava no idea about it. What`s wrong with my code?

TASK=MRPC
SUFFIX=sparsity0.95
EX_CATE=CoFi
SPARSITY=0.95
DISTILL_LAYER_LOSS_ALPHA=0.9
DISTILL_CE_LOSS_ALPHA=0.1
LAYER_DISTILL_VERSION=4
SPARSITY_EPSILON=0.01
DISTILLATION_PATH=/home/tt6232/KdQuant/teacher-model/bert-base-uncased/

PRUNED_MODEL_PATH=./out/$TASK/$EX_CATE/${TASK}_${SUFFIX}/best
PRUNING_TYPE=None # Setting the pruning type to be None for standard fine-tuning.
LEARNING_RATE=3e-5

bash scripts/run_CoFi.sh $TASK $SUFFIX $EX_CATE $PRUNING_TYPE $SPARSITY $DISTILLATION_PATH $DISTILL_LAYER_LOSS_ALPHA $DISTILL_CE_LOSS_ALPHA $LAYER_DISTILL_VERSION $PRUNED_MODEL_PATH $LEARNING_RATE &

Hi! In your code you calculate the Lc

https://github.com/princeton-nlp/CoFiPruning/blob/main/trainer/trainer.py#L682

And you use expected_size to calculate expected_sparsity , but does it match the equation in your paper?

https://github.com/princeton-nlp/CoFiPruning/blob/main/models/l0_module.py#L267

Actually you said that sˆ is the expected model sparsity calculated from z , but the lagrangian_regularization() do not have inputs or z Many thanks!

Hi @xiamengzhou , thanks for your contribution. But in your code, you use Model.from_pretrained to load the model architecture, and the files you have already provided. But if I want to prune my own, original model, for instance T5 model, using your method in the paper. Which code should I check? Many thanks:)

load_pruned_model in the cofi_utils file seems to take a model as a first argument, however load_model(..) calls load_pruned_model by passing a string. In this case the program crashes as the string doesn't have a "config" property for example