LLMBox: A Comprehensive Library for Large Language Models (2024)

Abstract: To facilitate the research on large language models (LLMs), this paper presents a comprehensive and unified library, LLMBox, to ease the development, use, and evaluation of LLMs. This library is featured with three main merits: (1) a unified data interface that supports the flexible implementation of various training strategies, (2) a comprehensive evaluation that covers extensive tasks, datasets, and models, and (3) more practical consideration, especially on user-friendliness and efficiency. With our library, users can easily reproduce existing methods, train new models, and conduct comprehensive performance comparisons. To rigorously test LLMBox, we conduct extensive experiments in a diverse coverage of evaluation settings, and experimental results demonstrate the effectiveness and efficiency of our library in supporting various implementations related to LLMs. The detailed introduction and usage guidance can be found at https://github.com/RUCAIBox/LLMBox.

heading: 1 Introduction
text: Recent years have witnessed the rapid progress of large language models (LLMs) (Zhao et al., 2023). In the research community, great efforts have been devoted to the release of well-trained LLMs, the design of special tuning and inference methods, and the conduct of systematic capacity evaluation. However, the reproducibility and fair comparison of existing studies should still be emphasized, since they are mostly developed in different ways or frameworks. Without the standardized and unified implementation, it would take substantial efforts to reproduce or improve upon existing research work.
* Co-leading the project. † Equal Contribution. Ordered by name. B Corresponding author.
Considering the above issue, in this paper, we present a comprehensive library, called LLMBox, for easing the development, use, and evaluation of LLMs. In particular, our library focuses on building a comprehensive and unified framework (including training, inference, and evaluation) for better supporting LLM-based research and applications. Although there are already several opensource libraries for LLMs (Kwon et al., 2023; Gao et al., 2023a; hiyouga, 2023), they typically focus on a certain or some stage(s) of LLMs (either pre-training or fine-tuning) or conduct the training pipeline of LLMs in a separate way. Moreover, they can seldom support comprehensive and unified evaluation of various LLMs.
In order to better facilitate research on LLMs, LLMBox introduces a series of new features for the library design, which can be summarized into three major aspects below: • Unified data interface. We design a unified data interface to encapsulate different formats of training data, including both plain texts and instruction data. With this interface, LLMBox can flexibly support the implementation of various strategies, such as dynamic mixture proportion (Xie et al., 2023) and combined training with pre-training and instruction data (Zeng et al., 2022). Furthermore, we extensively support mainstream training methodologies, including parameter-efficient tuning (e.g., LoRA (Hu et al., 2022)) and alignment tuning (e.g., PPO (Schulman et al., 2017)). • Comprehensive evaluation. To support a comprehensive comparison of LLMs’ performance, our library encompasses 18 downstream tasks alongside 56 datasets. Beyond the common benchmarks such as MMLU (Hendrycks et al., 2021) and GSM8K (Cobbe et al., 2021), our framework also extends the support for probing LLMs’ advanced
388
capabilities: human alignment, hallucination detection, instruction following, etc. Furthermore, LLMBox integrates a variety of publicly available LLMs and commercial APIs, offering a convenient platform for holistic evaluation. • More practical considerations. To be userfriendly, LLMBox is designed to provide an easyto-use pipeline, enabling users to quickly start with minimal commands. We introduce a GPU calculator to help users determine the minimum GPU resources necessary for training. To be efficient, we propose a novel prefix caching strategy for inference and a packing strategy for training. Remarkably, given the LLaMA (7B) model, our library can perform inference on the entire MMLU benchmark within six minutes on a single A800 GPU and completes instruction tuning with 52K instances on eight A800 GPUs in ten minutes.
An additional feature is that LLMBox is closely aligned with our prior survey paper on LLMs (Zhao et al., 2023). This is particularly useful for beginners, enabling the learning of basic knowledge and practice of LLMs through integrating the survey paper and the associated library.
In what follows, we will first introduce the training framework of our library in Section 2, then detail the utilization and evaluation parts in Section 3, and showcase how to use our library in Section 4. After that, we will conduct the experiments to verify the reliability of our LLMBox in Section 5, and conclude the paper in Section 6.

heading: 2 Training
text: The training process is a crucial step for the development of LLMs. However, it typically needs massive detailed designs considering both efficiency and effectiveness, and also often faces intractable problems when adapting into new domains or meeting special needs. To facilitate easy training of LLMs, we integrate various training methods and resources in our library, to unify and simplify their usage. Besides, we provide suggestions for GPU usage tailored to different training requirements.

heading: 2.1 LLM Training
text: In our LLMBox, we develop a unified architecture to encapsulate important training methods in developing LLMs, and implement efficient training strategies to support training on limited computing resource. The overall framework of LLMBox is illustrated in Figure 1.
Key Training Methods. In our LLMBox, we integrate massive functionalities to support the following four training processes: • Pre-training. Our LLMBox supports pretraining LLMs from scratch or continually pretraining using corpora in specific languages or specialized domains. For continually pre-training, LLMBox supports expanding the vocabulary to facilitate the adaptation of LLMs to new fields. • Instruction tuning. LLMBox integrates ten commonly-used datasets for supporting instructiontuning, covering NLP task (e.g., FLAN v2 (Chung et al., 2022)), daily chat (e.g., ShareGPT (Eccleston, 2023)), and synthetic datasets (e.g., Alpaca52K (Taori et al., 2023)). Additionally, we integrate three methods to synthesize or rewrite instructions, namely Self-Instruct (Wang et al., 2023a), Evol-Instruct (Xu et al., 2023), and topic diversifying (YuLan-Team, 2023). Based on the above datasets, we specially design unified dataset class, which can automatically preprocess these datasets into a unified format for training LLMs, and provide flexible interfaces for users to adjust the settings about the data (e.g., data mixture proportion). • Human alignment. To enhance the alignment of LLMs with human values, we incorporate both the widely-used RLHF method PPO (Schulman et al., 2017) and non-RL approach DPO (Rafailov et al., 2023). Besides, LLMBox also integrates several preference datasets, including HH-RLHF (Bai et al., 2022) and SHP (Ethayarajh et al., 2022).
Efficient Training Strategies. We also integrate several widely-used efficient training strategies or libraries, to support training LLMs with limited computing resources. • LoRA and QLoRA. LLMBox integrates the lightweight module LoRA (Hu et al., 2022) to facilitate the different training methods of LLMs in resource-constrained environments. We also encapsulate QLoRA (Dettmers et al., 2023) in LLMBox, which performs quantization on LoRA for further reducing its used GPU memory. • DeepSpeed. Our LLMBox library is based on the distributed training library DeepSpeed (Rasley et al., 2020), which includes a range of training optimization strategies for efficient training LLMs, including zero redundancy optimizer (ZeRO) (Rajbhandari et al., 2020), gradient checkpointing (Chen et al., 2016), FlashAttention (Dao et al., 2022), etc. • Packing. We implement the packing strategy (Raffel et al., 2020; Touvron et al., 2023b)
to enhance training efficiency. During pre-training, we concatenate all tokens into a long sentence and then split it to multiple sentences with the max length. For instruction-tuning, we concatenate all instructions as a long multi-turn conversation, and then break it into multiple conversations approaching to the maximum length constraint. Through minimizing paddings, we can optimize memory efficiency while maintaining model performance.

heading: 2.2 Training Suggestions
text: In practice, it is necessary for users to estimate the hardware requirements for training LLMs. Based on our LLMBox, we meticulously analyze GPU memory consumption throughout the model training process, by fully considering the impacts of parameters, gradients, optimizer states, and activation value (Rajbhandari et al., 2020; Ren et al., 2021; Korthikanti et al., 2023). We further introduce a “GPU memory calculator” to aid users in determining the minimal GPU requirements across LLMs with different parameter scales.
By merging the above strategies to reach efficiency1, the memory consumption of each GPU can be roughly estimated by the equation:
16p
n + (12 + 2l)bsh+ 12bsv, (1)
where p represents the total number of parameters, and n, l, b, s, h, v stand for the number of GPUs, number of layers, batch size, sequence
1For the training settings, we utilize data parallelism, ZeRO-3, gradient checkpointing, and FlashAttention.
length, hidden size, and vocabulary size, respectively. Taking the training of LLaMA-2 (7B) (l = 32, s = 4096, h = 4096, v = 32000) as an example, we employ two A100 (80G) GPUs (n = 2) with a batch size of b = 8. By using Eq. 1 with the above configuration, we can estimate an approximate GPU memory usage of 71.42GB per unit. As shown in Table 1, we extrapolate the minimum GPU requirements using Eq. 1 for different model sizes across varying training settings, to help users for selecting proper GPU resources. For other special training settings, we invite users to utilize the calculator available on our library2.
2https://github.com/RUCAIBox/LLMBox/blob/main/ training/gpu_calculator.py

heading: 3 Utilization and Evaluation
text: After training, we can develop suitable prompting strategies to use LLMs and assess their effectiveness. Users can reuse existing models, APIs or the models trained by LLMBox. The framework of our utilization pipeline is depicted in Figure 1.

heading: 3.1 Utilization Methods
text: We include quantization deployment strategies for using LLMs alongside two prompting methods: incontext learning (ICL) and chain-of-thought (CoT). • Quantization. To enhance memory efficiency during inference, LLMBox incorporates two quantization techniques: bitsandbytes (Dettmers et al., 2022) and GPTQ (Frantar et al., 2023). Both methods facilitate 8-bit and 4-bit quantization and GPTQ additionally supports 3-bit quantization. • In-context learning. We design a unified dataset class to organize diverse examples for fewshot learning. Furthermore, we implement three advanced ICL strategies, including KATE for example selection (Liu et al., 2022), GlobalE for example order arrange (Lu et al., 2022), and APE for instruction designing (Zhou et al., 2023c). • Chain-of-thought. Moreover, LLMBox incorporates several CoT prompting methods, such as program-aided (PAL) CoT (Gao et al., 2023b) and least-to-most CoT (Zhou et al., 2023a). We develop a flexible framework to facilitate selfconsistency (Wang et al., 2023a) and repeated sampling (Nijkamp et al., 2023), which are beneficial for tasks involving mathematics and coding.

heading: 3.2 Evaluation Methods
text: In LLMBox, we implement the evaluation of LLM performance through three distinct mechanisms: • Free-form generation: This is the basic evaluation method for generative LLMs and is applicable across all tasks. Models are required to generate responses to queries in an auto-regressive manner. We integrate common decoding strategies, including greedy search, temperature sampling, top-p sampling, repetition penalties, etc. • Completion perplexity: This method is widely adopted for assessing multi-choice tasks in base LLMs. It involves comparing the perplexity (PPL) of each completion conditioned on the context and choose the one with the lowest average PPL. Additionally, we incorporate the use of normalized PPL as introduced in GPT-3 (Brown et al., 2020). • Option probability: Similar to the multi-choice
formats in human examination, we feed a context with all the options to LLMs and require them to select the option letter (e.g., A). This approach is commonly utilized in chat-based models.
Significantly, we introduce prefix caching mechanism that caches the hidden states of common prefix texts to speed up the inference process. This strategy is organized at two levels: (1) we store the states of few-shot examples and compute them just once for all instances, e.g., 5-shot examples in MMLU (Hendrycks et al., 2021) and 8-shot examples in GSM8K (Cobbe et al., 2021); (2) we cache the states of identical contexts of different options when calculating completion perplexity. The effectiveness of this method is verified in Section 5.2.

heading: 3.3 Supported Models
text: We integrate a variety of LLMs to keep pace with the swift advancements in this field. Given that LLMBox is based on the Transformers library (Wolf et al., 2020), it is compatible with a vast majority of publicly available models. We list some included models as follows: • General models: LLaMA (Touvron et al., 2023a) and Mistral (Jiang et al., 2023); • Chinese models: Qwen (Bai et al., 2023) and Baichuan (Yang et al., 2023); • Multilingual models: BLOOM (Le Scao et al., 2022); • Chat models: LLaMA-2 Chat (Touvron et al., 2023b) and Vicuna (Chiang et al., 2023); • Code models: CodeGen (Nijkamp et al., 2023) and StarCoder (Li et al., 2023c); • Mathematical models: Llemma (Azerbayev et al., 2024) and DeepSeekMath (Shao et al., 2024). We also incorporate commercial APIs including OpenAI3 and Anthropic Claude4.

heading: 3.4 Supported Tasks
text: Currently, LLMBox integrates 18 diverse tasks and corresponding 56 datasets with hundreds of subsets. The broad range of supported datasets within LLMBox enables to evaluate various models. For instance, users can employ English benchmarks, language modeling, and knowledge reasoning datasets for evaluating foundational pre-trained LLMs. In the case of chat-based models, users can utilize datasets focused on open-ended dialogue, human alignment, and instruction following. We list some included tasks and datasets as follows:
3https://openai.com/ 4https://www.anthropic.com/
• English benchmarks: MMLU (Hendrycks et al., 2021) and BBH (Srivastava et al., 2023); • Chinese benchmarks: CMMLU (Li et al., 2023a) and C-Eval (Huang et al., 2023); • Multilingual benchmarks: TyDi QA (Clark et al., 2020) and MGSM (Shi et al., 2023); • Language modeling: LAMBADA (Paperno et al., 2016); • Open-ended dialogue: MT-Bench (Zheng et al., 2023) and AlpacaEval (Li et al., 2023d); • Machine translation: general translation task in WMT5 of every year and Flores-200 (Costajussà et al., 2022); 8 • Text summarization: CNN/Daily Mail (See et al., 2017) and XSum (Narayan et al., 2018); • Code synthesis: HumanEval (Chen et al., 2021) and MBPP (Austin et al., 2021); • Closed-book question answering: Natural Questions (Kwiatkowski et al., 2019) and TriviaQA (Joshi et al., 2017); • Reading comprehension: SQuAD 2.0 (Rajpurkar et al., 2018) and RACE (Lai et al., 2017); • Knowledge reasoning: HellaSwag (Zellers et al., 2019) and ARC (Clark et al., 2018); • Symbolic reasoning: Tables of Penguins (Herzig et al., 2020) and Colored Objects (Srivastava et al., 2023); • Mathematical reasoning: GSM8K (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021); • Human Alignment: TruthfulQA (Lin et al., 2022) and CrowS Pairs (Nangia et al., 2020); • Hallucination detection: HaluEval (Li et al., 2023b); • Instruction following: IFEval (Zhou et al., 2023b); • Environment Interaction: ALFWorld (Shridhar et al., 2021) and WebShop (Yao et al., 2022); • Tool Manipulation: Gorilla (Patil et al., 2023).

heading: 4 Library Usage
text: In this section, we present the application of our library across four distinct research scenarios, illustrated through example code snippets.
Continually Pre-Training Language-Specific Models. As introduced in Section 2, we facilitate the continual pre-training of existing English-based LLMs for quick acquisition of new languages. Figure 2 (a) illustrates the process of tuning a Chinese LLM from LLaMA-2. Users are required only to
5https://www2.statmt.org/
prepare Chinese plain texts, such as those from Wikipedia, into a file named chinese.txt. Subsequently, LLMBox integrates new Chinese tokens into the vocabulary and trains the model.
Adapting LLMs to Specialized Domains. LLMBox facilitates the adaptation of LLMs to various specialized domains via instruction tuning, covering domains such as medicine, law, and finance. We present a script in Figure 2 (b) to train a medical LLM. We implement a convenient dataset mixture approach to sample instances from raw medical texts, medical instruction data, and general conversation data. This enables users to adjust the proportion to make a balance between medical knowledge, medical tasks, and conversational skills, thereby crafting an effective medical assistant.
Comprehensively Evaluating LLMs. We cover a broad range of tasks and various models within LLMBox to implement comprehensive evaluation. As illustrated in Figure 2 (c), (d), and (e), we present three exemplary command lines. Users are only required to designate the model and dataset names via the -m and -d options to achieve an efficient and accurate assessment of model performance. Furthermore, LLMBox supports multiple utilization methods, such as in-context learning (-shots), self-consistency (--sample_num), and quantitation (--load_in_4bit).

heading: 100 / 0 50.6 15.0
text: Designing Novel Prompting Methods. Since the implementation of each dataset in LLMBox is unified, it offers the flexibility to add new datasets or design various prompting methods without affecting other modules. Figure 2 (f) overviews the design of our Dataset class. When adding a new dataset, users are only required to implement three functions: load_dataset to load evaluation and example datasets; format_instance to format each instance with instruction or few-shot examples; and reference to define the ground truth. In the core function format_instance, users can develop innovative prompting methods tailored for each evaluation instance using example datasets.

heading: 5 Experiment
text: In the section, we conduct extensive experiments to verify the effectiveness and efficiency.

heading: 5.1 Effectiveness Evaluation
text: The essential attribute of an open-source library is its ability to reproduce results effectively. To confirm this, we choose several representative training and utilization scenarios for testing the outcomes derived from LLMBox.
Training results. We train LLaMA-2 (Touvron et al., 2023b) with the mixture of instruction tuning data FLAN (Chung et al., 2022) and Alpaca52K (Taori et al., 2023) and evaluate its performance. We adjust the proportions of these datasets and assess the impact on performance using the MMLU benchmark (Hendrycks et al., 2021) and the chat-oriented AlpacaEval (Dubois et al., 2023).
The experiments are conducted with a batch size of 128 and a constant learning rate of 1× 10−5. The model undergoes training for a total of 1200 steps, and we report the peak performance observed on the evaluation datasets. The results in Table 3 indicate that FLAN improves the model’s performance on NLP tasks, whereas Alpaca-52K significantly enhances its performance in daily chat. Moreover, when mixing both instruction datasets yields a balanced improvement across both tasks, aligning with findings from prior research (Wang et al., 2023b).
Utilization results. Firstly, we examine the performance of LLaMA-2 (Touvron et al., 2023b) across various supported tasks. We totally evaluate nine tasks, including MMLU (5-shot, accuracy) (Hendrycks et al., 2021), BBH (3-shot, accuracy) (Srivastava et al., 2023), HumanEval (0-shot, pass1) (Chen et al., 2021), Natural Questions (NQs, 5-shot, EM) (Kwiatkowski et al., 2019), HellaSwag (0-shot, accuracy) (Zellers et al., 2019), ARCChanllge (ARC-C, 0-shot, accuracy) (Clark et al., 2018), WinoGrande (0-shot, accuracy) (Sakaguchi et al., 2021), BoolQ (0-shot, accuracy) (Clark et al., 2019), and GSM8K (8-shot, accuracy) (Cobbe et al., 2021). The results in Table 2 demonstrates that our LLMBox library faithfully reproduces the results reported in their original papers. Furthermore, we verify the performance of LLMBox across a variety of models. We utilize HellaSwag, MMLU, and GSM8K to evaluate the per-
formance of ten English LLMs, including GPTNeoX (Black et al., 2022), OPT (Zhang et al., 2022), BLOOM (Le Scao et al., 2022), LLaMA2 (Touvron et al., 2023b), Pythia (Biderman et al., 2023), MPT (Team, 2023b), Phi-2 (Javaheripi et al., 2023), Mistral (Jiang et al., 2023), Falcon (Almazrouei et al., 2023), Gemma (Google, 2024). We employ HellaSwag, C-Eval (Huang et al., 2023), and GSM8K to evaluate the performance of eight Chinese LLMs, including ChatGLM3 (Zeng et al., 2022), Chinese-LLaMA-2 (Cui et al., 2023), InternLM-2 (Team, 2023a), Baichuan2 (Baichuan, 2023), Qwen-1.5 (Bai et al., 2023), Aquila-2 (BAAI, 2023), Deepseek (DeepSeek-AI, 2024), Yi (Young et al., 2024). The results of these evaluations are detailed in Tables 4 and 5. We can find that our LLMBox is also compatible with various English and Chinese LLMs.

heading: 5.2 Efficiency Evaluation
text: The implementation efficiency is also a key factor to deploy LLMs. In addition to accurately reproducing results, we have optimized LLMBox for training and utilization efficiency. From the results in Table 6, it is evident that our prefix caching approach substantially decreases the inference time compared to the traditional Transformers implementation. As the number of examples increases (from 5-shot setting in MMLU to 8-shot setting in GSM8K), the efficiency gains from our method become increasingly pronounced. Remarkably, with the application of our prefix caching technique to the MMLU benchmark, LLMBox requires merely six minutes to process over ten thousand instances, achieving a 60% reduction in processing time compared to the vLLM toolkit. In the future, we aim to incorporate this prefix caching strategy into vLLM to further enhance the inference efficiency.

heading: 6 Conclusion
text: This paper presented LLMBox, a comprehensive library for conducting research on training, utilizing, and evaluating large language models. For training, we designed a unified data interface to support the implementation of various training strategies. For utilization and evaluation, we implemented typical approaches to use LLMs (including quantization, ICL, and CoT prompting), covered 18 tasks and 56 datasets, and included a number of popular opensourced LLMs and closed-source APIs. We further conducted extensive experiments to verify the effectiveness and efficiency of LLMBox. Our library provides a unified framework to compare, reproduce, and develop LLMs and supporting methods for academic purposes, which would be of important value to promote the research on LLMs.

heading: Acknowledgement
text: This work was partially supported by National Natural Science Foundation of China under Grant No. 62222215, Beijing Natural Science Foundation under Grant No.4222027 and L233008. Xin Zhao is the corresponding author.

heading: Future Work
text: In the future, we aim to incorporate this prefix caching strategy into vLLM to further enhance the inference efficiency