DOCSEARCHbyAlgolia
Submit Docs

CoFiPruning

repository·main·Indexed 41 days ago

https://github.com/princeton-nlp/cofipruning

A task-specific, structured pruning approach that jointly prunes coarse-grained units (layers, MLP dimensions) and fine-grained units (heads, hidden dimensions) to create compact, efficient models. Includes tools for pruning, fine-tuning via scripts/run_CoFi.sh, and evaluation of sparsity and inference time, with support for loading pruned models through HuggingFace using CoFiBertForSequenceClassification.

Tokens
1.4K
Snippets
6
Records
7
Agent score
2%

What's inside CoFiPruning

  1. Train CoFiPruning models

    main

    Training is performed using the scripts/run_CoFi.sh script. The script supports single-GPU training and can be used for both the initial pruning phase and subsequent fine-tuning of a pruned model.

    Pruning Phase

    To prune a model, provide the target sparsity and distillation parameters.

    Fine-tuning Phase

    To fine-tune a model that has already been pruned, set --pruning_type to None and provide the --pretrained_pruned_model path and a --learning_rate.

    # Example: Training (pruning)
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
SPARSITY_EPSILON=0.01

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 $SPARSITY_EPSILON

# Example: Fine-tuning after pruning
PRUNED_MODEL_PATH=$proj_dir/$TASK/$EX_CATE/${TASK}_${SUFFIX}/best
PRUNING_TYPE=None
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 $SPARSITY_EPSILON [PRUNED_MODEL_PATH] $LEARNING_RATE

  4. Use CoFiPruning models with HuggingFace

    main

    You can load and use the released CoFiPruning models via the HuggingFace interface by importing CoFiBertForSequenceClassification from models.modeling_bert. This allows you to use pretrained pruned models for sequence classification tasks.

    from models.modeling_bert import CoFiBertForSequenceClassification
model = CoFiBertForSequenceClassification.from_pretrained("princeton-nlp/CoFi-MNLI-s95") 
output = model(**inputs)

  5. Reference: `scripts/run_CoFi.sh` arguments

    main

    Arguments for the scripts/run_CoFi.sh training script:

    ArgumentDescription
    --task_nameSequence classification or extractive QA task name (e.g., MNLI). Use --train_file and --validation_file for other HuggingFace tasks.
    --ex_name_suffixExperiment name (used for output directory).
    --ex_cateExperiment category name (used for output directory).
    --pruning_typePruning units to use. Options: structured_heads, structured_mlp, hidden, layer. Default: structured_heads+structured_mlp+hidden+layer. Set to None for standard fine-tuning.
    --target_sparsityTarget sparsity of the pruned model.
    --distillation_pathDirectory of the teacher model.
    --distillation_layer_loss_alphaWeight for layer distillation.
    --distillation_ce_loss_alphaWeight for cross entropy distillation.
    --layer_distill_versionVersion of layer distillation (e.g., 4 for small datasets, 6 for more stable results via random teacher layer selection).
    --sparsity_epsilonEpsilon to relax sparsity constraint. If $>0$, models are saved at target_sparsity - sparsity_epsilon. Recommended 0.01 for 0.95 sparsity.
    --pretrained_pruned_model(Fine-tuning only) Directory of the pruned model.
    --learning_rate(Fine-tuning only) Learning rate for the fine-tuning stage.

  7. Load a pre-trained model using Transformers

    main

    To use pre-trained models within a notebook environment, import AutoModelForSequenceClassification from the transformers library. You can load a specific model architecture (e.g., roberta-base) using the .from_pretrained() method.

    import transformers
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("roberta-base")