transformercompression

To add support for a new model, it must be in Hugging Face Hub format (or local storage via --model and --model-path). You must implement a new model adapter and update hf_utils.get_model_and_tokenizer.

Implementation Steps

1. Implement ModelAdapter: Define how to interact with the model instance (e.g., accessing layers).
2. Implement LayerAdapter: Define how to interact with transformer layers (e.g., accessing attention/MLP components and updating the forward method arguments).
3. Implement a Compressed Transformer Layer: Create a class that subclasses the original transformer layer and provides an adapted forward() method. This method must handle how skip connection orthogonal matrices (self.*_shortcut_Q) are used.
   • For sequential blocks (e.g., OPT), handle accordingly.
   • For parallel blocks (e.g., Phi-2), handle accordingly.
   • If a skip connection does not need modification, self.*_shortcut_Q will be None.

Slicing Workflow

Once the adapter is implemented, the compression process follows these steps:

1. Replace modules with compressed equivalents via slicegpt.layernorm_fusion.replace_layers.
2. Fuse layer norms and add rotations to skip connections via slicegpt.layernorm_fusion.fuse_modules.
3. Rotate inputs and slice layers via slicegpt.rotate.rotate_and_slice.

Use the run_slicegpt.py script from the experiments folder to compress a model. This process applies orthogonal transformations and slices weight matrices to reduce the embedding dimension.

Key Arguments:

• --model: The Hugging Face model ID or local path.
• --save-dir: Directory where the sliced model will be saved.
• --sparsity: The target sparsity level (e.g., 0.25).
• --device: The device to use (e.g., cuda:0).
• --eval-baseline: Whether to evaluate the original model as a baseline.
• --no-wandb: Disable Weights & Biases logging.
• --hf-token: Required for models requiring Hugging Face authentication (alternatively, set the HF_TOKEN environment variable).

Note: Consult the script for the full set of available options.

python run_slicegpt.py \
       --model microsoft/phi-2 \
       --save-dir dir/to/save/sliced_model/in \
       --sparsity 0.25 \
       --device cuda:0 \
       --eval-baseline \
       --no-wandb

After slicing, you can perform Recovery Fine-Tuning (RFT) to restore model performance.

1. Install dependencies:

pip install -e .[experiment,finetune]

2. Run fine-tuning: Use run_finetuning.py to replicate paper experiments.

Key Arguments:

• --model: The original model (use --model-path instead of --sliced-model-path to fine-tune the original model).
• --sliced-model-path: Path to the previously sliced model.
• --sparsity: Must be specified when using --sliced-model-path to avoid default values.
• --finetune-dataset: The dataset for fine-tuning (e.g., alpaca).
• --ppl-eval-dataset: The dataset for perplexity evaluation.
• --lora-*: Parameters for LoRA (e.g., --lora-alpha, --lora-r, --lora-dropout).
• --lora-target-option: Specifies target modules (e.g., attn_head_and_mlp).

Note: You can use bo_finetuning.py to run Bayesian optimization over RFT hyperparameters.

python run_finetuning.py \
       --model microsoft/phi-2 \
       --sliced-model-path path/to/sliced \
       --save-dir dir/to/save/finetuned_model/in \
       --sparsity 0.25 \
       --device cuda:0 \
       --ppl-eval-dataset alpaca \
       --finetune-dataset alpaca \
       --finetune-train-nsamples 8000 \
       --finetune-train-seqlen 1024 \
       --finetune-train-batch-size 3 \
       --lora-alpha 10 \
       --lora-r 32 \
       --lora-dropout 0.05 \
       --lora-target-option attn_head_and_mlp \
       --eval-steps 16 \
       --save-steps 16 \
       --no-wandb

The SlicingScheduler is an abstract base class used to determine the slicing dimensions for various model components (embeddings, attention, MLP, and the LM head). When implementing a new scheduler, you must provide implementations for the following protected methods:

• _get_input_embedding_dimensions() -> dict[int, int]: Returns the input embedding dimensions.
• _get_attention_input_dimension(idx: int) -> int: Returns the attention input dimension for a specific layer.
• _get_attention_output_dimension(idx: int) -> int: Returns the attention output dimension for a specific layer.
• _get_mlp_input_dimension(idx: int) -> int: Returns the MLP input dimension for a specific layer.
• _get_mlp_output_dimension(idx: int) -> int: Returns the MLP output dimension for a specific layer.
• _get_head_dimension() -> int: Returns the LM head dimension.

Note that the public methods (e.g., get_embedding_dimensions(), get_attention_input_dimension()) are final and handle the logic of updating a SlicingConfig object with the results returned by your implementation.

The SlicingConfig dataclass holds the dimensions for the sliced model components. It can be instantiated directly or loaded from a dictionary or JSON string.

Key Fields:

• hidden_size: The model's hidden dimension.
• layers_num: Number of layers.
• do_slice_head: Boolean indicating if the head should be sliced.
• parallel_blocks: Boolean indicating if attention and MLP blocks are parallel.
• embedding_dimensions: dict[int, int] mapping embedding indices to dimensions.
• attention_input_dimensions, attention_output_dimensions, mlp_input_dimensions, mlp_output_dimensions: Dictionaries mapping layer indices to target dimensions.
• head_dimension: Target dimension for the head.
• const_dimension: Used for loading models sliced with constant sparsity.

Methods:

• SlicingConfig.from_dict(d): Creates a config from a dictionary (handles numeric string keys by converting them to integers).
• SlicingConfig.from_json_string(json_str): Creates a config from a JSON string.
• to_dict(): Returns a dictionary representation.
• to_json_string(): Returns a JSON string representation.

# Creating from a dictionary
data = {
    "hidden_size": 2560,
    "layers_num": 32,
    "embedding_dimensions": {"0": 1280}
}
config = SlicingConfig.from_dict(data)

# Creating from a JSON string
config = SlicingConfig.from_json_string('{"hidden_size": 2560, "layers_num": 32}')