View on GitHub

LLM Finetuning

The LLM fine-tuning package developed by the UVA ILP group

This package is to provide a codebase to easily fine-tune some open-source LLMs and reproduce the results.

So far, the codebase supports

For detailed instruction, please refer to the Notion site: LLM Fine-tuning


A Simple Example

Here is a simple example of using the package for fine-tuning or prediction on the SST2 dataset. You can also find this demo code in demo.py.

import fire

from lora import Llama_Lora


def main(
    task: str = "eval",
	base_model: str = "meta-llama/Llama-2-7b-hf",
):
    m = Llama_Lora(
        base_model=base_model,
    )
    if task == "train":
        m.train(
            train_file = "data/sst2/train.json",
            val_file = "data/sst2/val.json",
            output_dir = "./ckp_sst_llama2_lora",
            train_batch_size = 32,
            num_epochs = 1,
        )
    elif task == "eval":
        m.predict(
            input_file = "data/sst2/val.json",
            max_new_tokens = 32,
            verbose = True,
        )
    else:
        raise ValueError(f"Unrecognized task: {task}")


if __name__ == "__main__":
    fire.Fire(main)

For prediction, please use

python demo.py

For fine-tuning, please use

python demo.py --task train

Contributors

In addition, this package has been tested by many members from the UVA ILP group via their research projects.

This project was inspired by the Alpaca-LoRA project, from which we learned a lot!