CoLLiE

Introducing CoLLiE: A Comprehensive Toolkit for Training Large Language Models

CoLLiE, short for Collaborative Tuning of Large Language Models in an Efficient Way, is a comprehensive toolbox designed to help users train large language models from scratch. CoLLiE is built to cater to both beginners and seasoned professionals by offering an intuitive interface and highly customizable features.

Why Choose CoLLiE?

CoLLiE distinguishes itself by providing an all-in-one solution for large model training. This toolkit facilitates data preprocessing, model fine-tuning, model saving, and monitoring various metrics throughout the training process. CoLLiE supports multiple mainstream models such as MOSS, InternLM, LLaMA, and ChatGLM, allowing users to switch seamlessly between different models. Moreover, CoLLiE comes with extensive documentation, helping novices to quickly get started and empowering experienced users to customize their setup according to their specific needs. Whether you're just starting out or you're a seasoned professional, CoLLiE offers solutions that fit your needs.

Key Features

CoLLiE is based on DeepSpeed and PyTorch, offering collaborative and efficient tuning methods for large language models. Here are four main features:

• Parallel Strategies

  • Data Parallelism (DP)
  • Pipeline Parallelism (PP)
  • Tensor Parallelism (TP)
  • Zero Redundancy Optimizer (ZeRO)

• Efficient Fine-tuning

  • Techniques like LOMO and LoRA
  • Utilization of Flash Attention for faster computations

• Elegant Design and User-Friendly Interface

With these features, CoLLiE accelerates training speed, improves training quality, and reduces training costs, making it a versatile tool for model training.

Supported Models

• MOSS Series
• InternLM Series
• LLaMA Series
• ChatGLM Series

Installation

Before installing CoLLiE, ensure that your environment meets these requirements:

• PyTorch version 1.13 or higher
• CUDA version 11.6 or higher
• Linux Operating System

CoLLiE can be installed through PyPI with the following command:

pip install collie-lm

Alternatively, you can install it from the source:

git clone https://github.com/OpenLMLab/collie
python setup.py install

Getting Started

CoLLiE provides a straightforward way to start your training endeavors. The process involves importing necessary packages, setting paths, configuring CoLLiE, loading datasets and pre-trained models, setting up optimizers, adding monitors, and initializing a trainer. For instance, when training the Moss model, CoLLiE enables the use of optimizers such as LOMO and provides functionality such as ZeRO for memory efficiency. Detailed step-by-step instructions and scripts are available for reference.

Exciting Features and Plugins

CoLLiE offers a range of plug-and-play features, such as the Monitor and asynchronous DataProvider:

• Monitor: Tracks various metrics during training and stores results in a CSV format for analysis.
• Asynchronous DataProvider: Allows real-time data provision during Human Evaluation with ease.

Users can leverage these plugins to enhance their training process.

Community and Contribution

CoLLiE is supported by an active community. Users are encouraged to contribute to the project, explore available resources, and exchange ideas with other contributors. Further information and examples can be found in the project's documentation and example section.

CoLLiE stands out with its blend of comprehensive features, efficient strategies, and user-centric design, ensuring it meets varied user requirements effectively. Whether you’re embarking on your first foray into large model training or enhancing existing projects, CoLLiE is equipped to provide you with the tools necessary for success.