torchtune: The new PyTorch LLM fine-tuning library
A library to fine-tune a set of tested and proved LLMs on downstream tasks made by the PyTorch team.
Created on April 15|Last edited on April 17
Comment
The PyTorch team has been working hard to put together a very flexible and tested codebase to fine-tune their own set of LLama based LLMs and other popular models like Mistral and Gemma. The repo is under active development and is accepting contributions. Of course the team has integrated Weights & Biases logging to keep track of those precious metrics.
Showing first 10 runs
Showing first 10 runs
Using torchtune
First you will need to install torchtune, visit their github page and follow the instructions. As of today, torchtune is tested against latest PyTorch and is installable in editable mode. The latter is encouraged as you may want to change the configs and recipes to suit your own use case.
The paradigm is using the library as a CLI tool that enables you to pair a recipe and a config as follows:
torchtune repo provides a set of proven and tested recipes in the torchtune/recipes folder along with the configs to use for different models:
As many current fine-tuning libraries (axolotl, trl), the configs are in the form a YAML file that specifies the hyper-parameters to use for training. It also defines the inputs and outputs: dataset, tokenizer, model checkpoints, logging, metrics, etc...
Running a recipe with Weights & Biases
You can make use of the Weights & Biases integration by modifying the config file and selecting the appropriate WandBLogger class.
# Inside the config.yaml file# Loggingmetric_logger:_component_: torchtune.utils.metric_logging.WandBLogger. # tells torchtune to log to W&Bproject: torchtune-mistral # the name of the W&B project to log tolog_every_n_steps: 5 # log every 5 steps to W&B
Then, you can run the command as you normally would and the first time you will be asked for your W&B API key to connect your training to W&B. Got to https://wandb.ai/authorize and paste you API key, and you are good to go!
Let's take the excellent Mistral 7B model for a spin and let's perform a LoRA fine-tune on the Alpaca dataset (similar to what we did on the LLM Finetuning series).
We add the snippet of code above into the config, we then need to download the model. You can grab the Mistral model from the Hugging Face hub using the tune command:
tune download mistralai/Mistral-7B-v0.1 --hf-token <HF_TOKEN> --output-dir /tmp/Mistral-7B-v0.1
Note: Currently torchtune only supports models in *.bin format, *.safetensor support is coming.
Now that we have the model, we can launch training:
tune run lora_finetune_single_device --config mistral/7B_lora_single_device.yaml epochs=1
Another thing one can do is override parameters at launch time, in our case, we will run for only 1 epoch.
Run set
31
Experimenting with different hyper parameters using Parallel Coordinates Plots
We can visualize different experiment at once and check how the hyper parameters impact the target metric on visual way using the Parallel Coordinates Plot, for instance, we can assess how different combinations of lora configurations impact the final loss function:
Run set
34
Conclusions
We have a new shiny library that is coming directly from the PyTorch team. They are probably gearing up for Llama 3 by providing users a first-hand fine-tuning experience. The library is evolving quickly and adding support of new models every day. The codebase is very modular and enables user to override any part of the preprocessing steps. Also the training loop is exposed on a single file (the recipe file) so you can quickly add whatever logic you may need. We love this modularity and look forward to see what people build on top of this tool.
Add a comment
Iterate on AI agents and models faster. Try Weights & Biases today.