A novel decoding method for LLM inference that skips the non-essenial layers using a learned RL policy.
Explore the docs »
View Demo
·
Report Bug
·
Request Feature
Table of Contents
skipr is an advanced method designed to optimize the performance of large language models (LLMs) by dynamically skipping layers based on a reinforcement learning (RL) policy with a reward model. This approach improves inference speed and computational efficiency while maintaining model accuracy.
- Dynamic Layer Skipping: Utilizes an RL policy to decide which layers to skip during inference, reducing computational overhead.
- Reward-Based Optimization: Employs a reward model to evaluate the effectiveness of skipped layers, ensuring that performance is not compromised.
- Adaptive Inference: Continuously learns and adapts to different tasks and contexts, optimizing layer usage for each specific scenario.
- Easy Integration: Designed to integrate seamlessly with existing LLM frameworks and architectures.
The skipr initializes with a pre-trained LLM and a reward model. The reward model evaluates the impact of skipping layers on the overall performance.
An RL policy is used to determine which layers to skip based on the reward model's feedback. We use a policy gradient method to maximize the reward, the reward function being
During inference, the tool dynamically skips layers as determined by the RL policy, allowing for faster inference. The reward model provides feedback on the performance of skipped layers, enabling the RL policy to adjust its decisions and improve over time.
This project use poetry for dependency management, and is required for running skipr locally.
-
Clone the repo
git clone https://github.com/dhmnr/skipr.git
-
Install Dependencies
poetry install
-
Train the classifier for BERT and later tain the policy network or run the script
run.shexport TASK_NAME=rte export EPOCHS=1000 poetry run glue \ --model_name_or_path google-bert/bert-base-cased \ --task_name $TASK_NAME \ --do_train \ --do_eval \ --max_seq_length 128 \ --per_device_train_batch_size 32 \ --learning_rate 2e-5 \ --num_train_epochs $EPOCHS \ --output_dir ./models/$TASK_NAME/classifier \ --overwrite_output_dir \ --mode classifier poetry run glue \ --model_name_or_path ./models/$TASK_NAME/classifier \ --task_name $TASK_NAME \ --do_train \ --do_eval \ --max_seq_length 128 \ --per_device_train_batch_size 32 \ --learning_rate 2e-4 \ --num_train_epochs $EPOCHS \ --output_dir ./models/$TASK_NAME/policy \ --overwrite_output_dir \ --mode policy \ --alpha 0.07
Contributions are what make the open source community such an amazing place to learn, inspire, and create. Any contributions you make are greatly appreciated.
If you have a suggestion that would make this better, please fork the repo and create a pull request. You can also simply open an issue with the tag "enhancement". Don't forget to give the project a star! Thanks again!
- Fork the Project
- Create your Feature Branch (
git checkout -b feature/AmazingFeature) - Commit your Changes (
git commit -m 'Add some AmazingFeature') - Push to the Branch (
git push origin feature/AmazingFeature) - Open a Pull Request
Distributed under the GNU GPLv3 License. See LICENSE for more information.
Project Link: https://github.com/dhmnr/skipr
Use this space to list resources you find helpful and would like to give credit to. I've included a few of my favorites to kick things off!