AutoTrain Advanced now supports Experiment Tracking

Training large language models on custom datasets has become more accessible with the AutoTrain Advanced library by Hugging Face. This low-code solution empowers AI researchers like you to fine-tune models such as Llama 2, Falcon, Mistral, or any other Large Language Models (LLMs) on your local machine, all for free. Under the hood, AutoTrain Advanced leverages powerful libraries like torch, transformers, peft, and trl to streamline the fine-tuning process. This not only simplifies the workflow but also accelerates the training, making it one of the fastest options available.

Announcement

After the merge of my recent Pull Request, one notable feature addition of AutoTrain Advanced is its support for reporting results and training logs. By integrating with popular experiment tracking platforms such as "azure_ml", "clearml", "codecarbon", "comet_ml", "dagshub", "flyte", "mlflow", "neptune", "tensorboard", and "wandb" the library offers flexibility and choice in monitoring and analyzing your experiments.

How to use

To leverage this reporting functionality, you can use the --log argument. For instance, specifying --log wandb will seamlessly log your results to Weights & Biases. Similarly, you can choose other platforms like TensorBoard by providing the corresponding argument. If you prefer minimal logging and only want CLI output, you can omit the --log argument.

Demo

To demonstrate these changes we will fine-tune an LLM on a math dataset and try achieving State-Of-The-Art result in pass@1 on the GSM8k Benchmarks. The A100 GPU used for this fine-tuning process is generously provided by Weights & Biases. I am thankful to Soumik Rakshit from team W&B for constant support in this integration. The experiment can be tracked using Weights & Biases here.

Preparing the dataset

AutoTrain Advanced expects your CSV custom dataset in a certain format to work properly. Your training file must contain a "text" column on which the training will be done. For best results, the "text" column should have data in the ### Human: Question?### Assistant: Answer. format. A great example for the kind of dataset AutoTrain Advanced expects would be timdettmers/openassistant-guanaco. However, if you observe the MetaMathQA dataset, there are 3 columns - "query", "response" and "type". We will preprocess this dataset by removing the "type" column and combining the content of the "query" and "response" columns under one "text" column with the ### Human: Query?### Assistant: Response. format. The resulting dataset is rishiraj/guanaco-style-metamath and it will be used for training.

Adjusting hyperparameters

AutoTrain Advanced comes with a host hyperparameters we can tune to get the best model. While the default hyperparameters are a great start for everyone, I made a few changes there that are suitable for our use case. Here are the hyperparameters I used:

learning_rate = 2e-5
num_epochs = 3
batch_size = 4
block_size = 1024
trainer = "sft"
warmup_ratio = 0.03
weight_decay = 0.
gradient_accumulation = 4
lora_r = 16
lora_alpha = 32
lora_dropout = 0.05
logging_steps = 10
log = "wandb"

Run command

If you want to push your trained model to a private repo in your Hugging Face Account, you can do so by setting the hf_token variable. You can find your Hugging Face write token here. Here is the run command I used:

autotrain llm \
--train \
--model "HuggingFaceH4/zephyr-7b-alpha" \
--project-name "zephyr-math" \
--data-path data/ \
--text-column text \
--lr str(learning_rate) \
--batch-size str(batch_size) \
--epochs str(num_epochs) \
--block-size str(block_size) \
--warmup-ratio str(warmup_ratio) \
--lora-r str(lora_r) \
--lora-alpha str(lora_alpha) \
--lora-dropout str(lora_dropout) \
--weight-decay str(weight_decay) \
--gradient-accumulation str(gradient_accumulation) \
--logging-steps str(logging_steps) \
--log str(log) \
--fp16 \
--use-peft \
--use-int4 \
--merge-adapter \
--push-to-hub \
--token str(hf_token) \
--repo-id "rishiraj/zephyr-math"

Results

Check out the W&B Report for a detailed overview of the finetuned model including its Benchmark scores on a variety of tests like the ARC, HellaSwag, MMLU, TruthfulQA. I also included a comparison with other open-source LLMs on GSM8k Pass@1 and MATH Pass@1. rishiraj/zephyr-math is the LLM (released under Apache License 2.0) fully fine-tuned on the MetaMathQA dataset and based on the powerful HuggingFaceH4/zephyr-7b-alpha model.

Conclusion

AutoTrain Advanced by Hugging Face emerges as a versatile tool for fine-tuning Large Language Models. Its combination of simplicity, efficiency, and extensive logging support makes it a valuable asset for AI researchers and developers. By allowing seamless integration with popular experiment tracking platforms, it facilitates effective monitoring and analysis of model training.

Experimentation is at the core of AI research, and tools like AutoTrain Advanced contribute significantly to advancing the field by providing accessible yet powerful solutions for model fine-tuning. Follow me on my HF profile if you found this read useful.