The growing interest in Large Language Models (LLMs) has led to a surge in tools and wrappers designed to streamline their training process.
Popular options include FastChat from LMSYS (used to train Vicuna) and Hugging Face’s transformers/trl libraries (used in my previous article). In addition, each big LLM project, like WizardLM, tends to have its own training script, inspired by the original Alpaca implementation.
In this article, we will use Axolotl, a tool created by the OpenAccess AI Collective. We will use it to fine-tune a Code Llama 7b model on an evol-instruct dataset comprised of 1,000 samples of Python code.
Why Axolotl?
The main appeal of Axolotl is that it provides a one-stop solution, which includes numerous features, model architectures, and an active community. Here’s a quick list of my favorite things about it:
- Configuration: All parameters used to train an LLM are neatly stored in a yaml config file. This makes it convenient for sharing and reproducing models. You can see an example for Llama 2 here.
- Dataset Flexibility: Axolotl allows the specification of multiple datasets with varied prompt formats such as alpaca (
{"instruction": "...", "input": "...", "output": "..."}), sharegpt:chat ({"conversations": [{"from": "...", "value": "..."}]}), and raw completion ({"text": "..."}). Combining datasets is seamless, and the hassle of unifying the prompt format is eliminated. - Features: Axolotl is packed with SOTA techniques such as FSDP, deepspeed, LoRA, QLoRA, ReLoRA, sample packing, GPTQ, FlashAttention, xformers, and rope scaling.
- Utilities: There are numerous user-friendly utilities integrated, including the addition or alteration of special tokens, or a custom wandb configuration.