LLM Finetuning Frameworks: A Review
Published on March 28, 2024.
When it comes to finetuning LLMs, a variety of frameworks exist to easily manage training. I tried out the most popular tools, and here are my takeaways.
Axolotl
Axolotl is a open source LLM training framework created by the OpenAccess AI Collective. Licensed under the Apache 2.0 license, Axolotl is a powerful and customizable framework for finetuning LLMs.
Unlike most of the other frameworks, Axolotl does not offer a GUI (graphical user interface). Instead, it accepts advanced configuration through YAML. While it doesn't have an integrated dashboard or logging, it supports both Weights and Biases, a proprietary logging platform, and MLFlow, an open source alternative.
In addition, Axolotl is constantly updated with new features. For example, Axolotl is one of the few frameworks I tried that supports the new GaLore training method.
Overall, although Axolotl provides increased customizability, it has a steeper learning curve than other options and is less beginner-friendly. Personally, I found Axolotl to be the best framework for my needs.
H2O LLM Studio
H2O LLM Studio is a framework by H2O.ai that allows easy finetuning of models. It's open sourced under the Apache license and beginner-friendly, with simple installation and an easy-to-navigate web-based GUI. Although it does not provide as much flexibility when customizing training runs as Axolotl (for example, H2O only supports training on a single dataset while Axolotl supports training on multiple), it's simple web-based GUI makes it much easier to use.
Overall, H2O LLM Studio is a good option when running basic/simpler finetunes, but lacks the customizability to make it feasible for more complex runs.
AutoTrain Advanced
AutoTrain Advanced, released by Hugging Face, is a no-code GUI framework for finetuning LLMs. Like many other frameworks, it's open-sourced and Apache licensed.
AutoTrain Advanced is the successor to Hugging Face's proprietary AutoNLP (later AutoTrain) service, first launched in 2021. Open sourced in 2023, it has now completely replaced AutoTrain/AutoNLP.
The main advantage of AutoTrain is it's integration with Hugging Face. Not only does AutoTrain automatically upload trained models to the Hugging Face Hub, but it also seamlessly integrates with Hugging Face Spaces.
Configuration of training is done through a simple dashboard, which allows selecting the training type (i.e. Default/SFT/DPO) and selecting a base model from a dropdown (the Web UI doesn't support custom models). The dashboard also features JSON editor, which can be used for slightly more advanced configuration.
While AutoTrain's simple web interface makes it easy to configure basic options, more advanced options, such as prompt templating, must be modified through a JSON configuration file. This can be difficult, since there is currently no available guide for the options in the JSON configuration.
Although the Web UI lacks easy customizability, AutoTrain also offers a CLI that can be used locally which provides more options.
Overall, AutoTrain seems like a good option for basic projects, but is more difficult to use for more complex training runs.
LLaMA Factory
LLaMA Factory is another open sourced, Apache-licensed framework for finetuning LLMs. Similar to AutoTrain Advanced, LLaMA Factory offers both a CLI and GUI.
LLaMA Factory's GUI allows you to select a base model and configure training. It supports both SFT (Supervised Fine-Tuning) and DPO (Direct Performance Optimization), but does not allow selecting a custom dataset without adding it to an index in a JSON file.
When it comes to customizability, LLaMA Factory features standard configuration options. While H2O LLM Studio allows for more customization, LLaMA Factory supports more features such as DoRA and GaLore.
Overall, while LLaMA Factory provides a cleaner and more organized interface than H2O LLM Studio, as well as additional features, H2O LLM Studio allows for more customization.
Oobabooga
Oobabooga is an open source framework for text generation. While it's primarily designed for inference, it also features a built-in LoRA training feature. While its user interface is quite intuitive and easy-to-use, it only supports LoRA and doesn't support methods such as DPO.
While Oobabooga supports conversational datasets and chat templating, it only supports a subset of chat templates (Alpaca, Llama 2, and Vicuna) and dataset formats.
Overall, Oobabooga's training feature is a useful addition to the Web UI but may not be suitable to use in production.
Conclusion
Personally, Axolotl was the favorite framework I tried. It was powerful and flexible, with many useful features and advanced customizability. I found H2O LLM Studio the easiest to use, with a simple dashboard and intuitive setup. AutoTrain Advanced was promising, but I found the JSON configuration difficult to navigate. LLaMA Factory, while providing an easy-to-use GUI, made it difficult to add a custom dataset. Oobabooga, while simple and easy to use, only supported LoRA.
Note: Updated on 04/01/24 with updated information about AutoTrain Advanced.