Efficient LLMs Pipeline

This repository combines three efficient techniques into one pipeline for selecting training data, fine-tuning with LoRA selected model layers, and pruning the resulting model.

• Covariance-based data selection (CPDS): Repository coming soon.
• Layer-wise LoRA with CKA: https://github.com/c2d-usp/Layer-wise-LoRA-with-CKA
• Mixture of Pruners (MoP): https://github.com/c2d-usp/Efficient-LLMs-with-MoP

Pipeline overview

The pipeline runs four steps:

1. Selects a subset of the instruction dataset with CPDS.
2. Uses CKA to choose which transformer layers to fine-tune, then trains those layers with LoRA.
3. Prunes the fine-tuned model with MoP.
4. Runs recovery fine-tuning on the pruned model.

Computation saved: The pipeline reduced fine-tuning time by 91.60% and GPU memory use by 13.90%, while pruning 20% of the model weights and speeding up inference by 1.219x.

Repository structure

.
├── run_pipeline.sh                      # Runs the complete pipeline
├── setup_environments.sh                # Creates the three Python environments
├── data_selection_with_cpds/
│   ├── covariance_selection.py          # Selects the training subset
│   ├── finetuning.py
│   ├── datasets/
│   │   └── alpaca_gpt4_data.json
│   ├── requirements.txt
│   └── run.sh
├── layer_wise_lora_finetuning/
│   ├── select_layers_decoder_alpaca.py  # Selects layers with CKA
│   ├── finetuning.py                    # Fine-tunes with layer-wise LoRA
│   ├── requirements.txt
│   └── run.sh
├── pruning_with_mixture_of_pruners/
│   ├── llama2/                          # MoP pruning and recovery code
│   ├── requirements.txt
│   ├── run.sh
│   └── LICENSE
└── outputs/                             # Selected data, checkpoints, and models

The setup script also creates three local environments under .venvs/: cpds, layer_wise_lora, and mop.

Run

From the repository root:

chmod +x setup_environments.sh run_pipeline.sh
./setup_environments.sh
./run_pipeline.sh

Paths, GPU selection, and the main hyperparameters can be changed at the top of run_pipeline.sh.

By default, the pipeline selects 10% of the data, fine-tunes selected layers of Llama-2-7B, performs seven MoP iterations (20 % Compression), and runs recovery fine-tuning. Results are written to outputs/default/. The final recovered model is saved in:

outputs/default/mop/prune_iteration_7_ft_2ep

Acknowledgments

We thank Instituto de Ciência e Tecnologia Itaú (ICTi) for the technical support, resources, and financial aid in the development of the research project. The authors would also like to thank the Programa de Bolsas Itaú (PBI) of the Centro de Ciência de Dados (C2D), supported by Itaú Unibanco S.A.