C o LL i E : Collaborative Training of Large Language Models in an Efficient Way (2023)

Abstract: Large language models (LLMs) are increasingly pivotal in a wide range of natural language processing tasks. Access to pre-trained models, courtesy of the open-source community, has made it possible to adapt these models to specific applications for enhanced performance. However, the substantial resources required for training these models necessitate efficient solutions. This paper introduces CoLLiE, an efficient library that facilitates collaborative training of large language models using 3D parallelism, parameter-efficient fine-tuning (PEFT) methods, and optimizers such as Lion, Adan, Sophia, and LOMO. With its modular design and comprehensive functionality, CoLLiE offers a balanced blend of efficiency, ease of use, and customization. CoLLiE has proven superior training efficiency in comparison with prevalent solutions in pre-training and fine-tuning scenarios. Furthermore, we provide an empirical evaluation of the correlation between model size and GPU memory consumption under different optimization methods, as well as an analysis of the throughput. Lastly, we carry out a comprehensive comparison of various optimizers and PEFT methods within the instruction-tuning context. CoLLiE is available at https://github.com/OpenLMLab/collie.

## introduction large language models (llms) have demonstrated remarkable abilities across various natural language processing tasks and showcased potential as intelligent assistants. thanks to the vibrant opensource community, multiple excellent large language models' weights are accessible, including opt @xcite , bloom @xcite , llama @xcite , etc. despite the impressive general capabilities of pre-trained llms, training for particular application scenarios can lead to even more outstanding performance. as shown in figure 1 , the training process can be divided into two stages: 1. further pre-training, which supplements specific domain knowledge and expands the vocabulary to enhance tokenization efficiency; 2. instruction-tuning, which adapts the model to downstream tasks and improves its instruction-following ability. with the scaling of language models, the resources required for training have increased substantially, making it infeasible to train the entire model on a single gpu. model parallelism addresses this issue by partitioning the model across different gpus, distributing the training workload among these gpus. this can be achieved through three methods: tensor parallelism (tp, @xcite ), pipeline parallelism (pp, @xcite al. (2019); narayanan et al. (2019) ), and stage 3 of zero redundancy optimizer @xcite ). in addition, during the instruction-tuning stage, there are resourceefficiency and training-effectiveness trade-off approaches @xcite : parameter-efficient fine-tuning (peft) methods @xcite . these methods selectively choose or add a few parameters for training, effectively reducing the gpu memory required to train large language models. in this context, we introduce collie, an easyto-use library for collaborative training of large language models in an efficient way. the library not only integrates the previously mentioned three parallelism strategies and peft methods, but also implements efficient optimizers such as lion @xcite , adan @xcite , sophia @xcite , and lomo @xcite . we have restructured multiple mainstream open-source models to support tp and pp and incorporated flashattention @xcite to further boost efficiency, while retaining interfaces similar to huggingface models within colliemodel class. training efficiency is one of the most distinctive feature of collie, boasting a significantly higher training throughput compared to current popular solutions. collie also offers a wide range of functionalities, including data preprocessing, model training, checkpoint saving and monitoring and evaluation during training process. collie's modular design allows for flexible combinations of parallelism strategies, peft methods, and training hyperparameters, which can be configured simply by modifying the collieconfig class. furthermore, collie is purposefully designed with extensibility, providing customizable functionalities. in summary, collie offers a comprehensive solution that caters to the needs of both beginners and experienced professionals. our contributions can be summarized as follows: peft methods there has been a rise in using parameter-efficient fine-tuning (peft) techniques to adapt models for instruction-tuning by adjusting partial parameters. one of the early success is adapter tuning @xcite , which inserts trainable neural modules into transformers layers while keeping the original model unchanged. in line with adapter tuning, lora @xcite reparameterizes the dense layers and only updates low rank matrices while introducing no latency during inference. prefix-tuning (li and liang, 2021) trains a task specific prefix prepended to each layer of the transformer encoder and achieves comparable performance with full parameter fine-tuning on generative tasks. similarly, prompt-tuning @xcite simplifies the additional prefix to the input embeddings, and only updates the parameters corresponding to the prompts. while the peft library @xcite ) implements these algorithms at the model level, it relies on huggingface models and lacks a comprehensive functionality, particularly the necessary integration with model parallelism to facilitate training of extremely large models. parallelism strategies parallelism strategies refer to the methodology of utilizing multiple gpus to execute training or inference tasks. data parallelism involves distributing the input data to different gpus for computation. however, each gpu stores an identical copy of the optimizer state and model weights, which limits the maximum model size that can be trained with data parallelism to that of a single gpu. to mitigate this redundancy, @xcite proposes a parallelism strategy in the three stages of zero, evenly partitioning the optimizer states, gradients, and weights across different gpus. tensor parallelism also partitions the weights evenly, while it varies the approach to partition and communicate. specifically, whereas zero-3 gathers the weight matrices, tensor parallelism all reduces the intermediate computational results. pipeline parallelism partitions the model by layers across gpus, requiring communication only between the layers at the split points. this strategy yields the least communication overhead. existing toolkits, such as huggingface's trainer @xcite and lmflow @xcite , choose zero-3 as parallel method. zero-3 is preferred because it does not impose specific requirements on the model structure, allowing direct usage of huggingface models. however, it exhibits lower throughput compared to the combination of tp and pp in scenarios involving large batch size pre-training or constrained communication. collie supports the hybrid application of data parallelism, tensor parallelism, and pipeline parallelism, collectively termed as 3d parallelism, with the parallel sizes adjustable via collieconfig. ## collie in this section, we will introduce the implementation and features of collie. figure 2(a ) presents an overview of the collie's overall structure, centered around the trainer class. the collieconfig, colliemodel, colliedataset, and optimizer classes serve as inputs to the trainer. collie also provides a set of convenient plugins for the trainer, including callback, monitor, evaluator, and probe, enabling users to customize the training process. depending upon the configurations and selected plugins, the trainer performs the training process, saves model checkpoints, and records system metrics, including loss, throughput, and evaluation results. as shown in figure 2 (b), based on py-torch @xcite and deepspeed @xcite , collie employs a collaborative approach using various techniques to facilitate the more efficient training of large language models. specifically, collie integrates flashattention to enhance efficiency. it implements zero-powered dp, tp, and pp to support 3d parallelism. additionally, lomo and peft methods are incorporated to realize memory-efficient fine-tuning approaches. appendix a provides a brief tour demonstrating how to use collie for training. ## conclusion we have introduced collie, a library for collaboratively training large language models in an efficient way. collie offers efficient models with flashattention and structurally supportive for 3d parallelism. moreover, collie provides a comprehensive and customizable trainer to assist users throughout the training process, supporting various training methods. we have tested the relationship between the gpu memory requirements and model parameter sizes as a reference for users. in terms of throughput, collie is significantly more efficient than huggingface's parallel solutions. the effectiveness of different training methods are also empirically assessed on instruction-tuning tasks.