LMFlow: An Extensible Toolkit for Finetuning and Inference of Large Foundation Models (2306.12420v2)

Abstract: Foundation models have demonstrated a great ability to achieve general human-level intelligence far beyond traditional approaches. As the technique keeps attracting attention from the AI community, an increasing number of foundation models are becoming publicly accessible. However, a significant shortcoming of most of these models lies in their performance in specialized-domain and task-specific applications, necessitating domain- and task-aware fine-tuning to develop effective scientific LLMs. As the number of available foundation models and specialized tasks keeps growing, the job of training scientific LLMs becomes highly nontrivial. In this paper, we initiate steps to tackle this issue. We introduce an extensible and lightweight toolkit, LMFlow, which aims to simplify the domain- and task-aware finetuning of general foundation models. LMFlow offers a complete finetuning workflow for a foundation model to support specialized training with limited computing resources. Furthermore, it supports continuous pretraining, instruction tuning, parameter-efficient finetuning, alignment tuning, inference acceleration, long context generalization, model customization, and even multimodal finetuning, along with carefully designed and extensible APIs. This toolkit has been thoroughly tested and is available at https://github.com/OptimalScale/LMFlow.

Overview of the LMFlow Toolkit for Finetuning Large Foundation Models

The paper introduces LMFlow, a toolkit designed to streamline the finetuning and inference processes for large foundation models, particularly focusing on LLMs. The authors address the challenge of adapting these models to specialized tasks, acknowledging that despite the general capabilities of foundation models, domain-specific finetuning remains indispensable.

Key Features and Contributions

LMFlow is positioned as an extensible and lightweight toolkit, with the following salient features:

• Comprehensive Finetuning Workflow: It supports continuous pretraining, instruction tuning, and reinforcement learning with human feedback (RLHF). This comprehensive approach enables users to perform domain adaptation, task adaptation, and alignment tuning effectively.
• Efficient Resource Utilization: The toolkit is designed to work efficiently with limited computational resources. For instance, it allows the personalization of a 7-billion-parameter model using a single Nvidia 3090 GPU in a matter of hours.
• Low-Rank Adaptation (LoRA): By incorporating LoRA, LMFlow offers parameter-efficient finetuning, reducing the number of trainable parameters while maintaining model performance.
• Novel Reinforcement Learning Approach: The paper introduces a new algorithm, Reward rAnked FineTuning (RAFT), which simplifies the RLHF pipeline. RAFT allows for the continuation of training using SFT-like techniques, providing stability and computational efficiency over traditional PPO methods.

Numerical Results and Claims

In task tuning, LMFlow demonstrates notable improvements in the medical domain using models such as the LLaMA series. For example, the LLaMA-33B model, with LoRA finetuning, achieved significant performance enhancements in medical QA tasks. Furthermore, the Robin models, derived from extensive instruction tuning, performed competitively on the Huggingface Open LLM Leaderboard, indicating the effectiveness of the toolkit in instruction-following tasks.

Implications and Future Directions

Practically, LMFlow enables researchers and developers to rapidly adapt large models to diverse tasks using limited resources. Theoretically, the introduction of RAFT hints at a more stable and resource-efficient method for aligning models to human preferences.

Future work may include expanding LMFlow's capabilities to other domains and integrating more sophisticated techniques for instruction and alignment tuning. As LLMs evolve, the need for flexible and efficient finetuning tools like LMFlow will likely increase, underscoring its potential impact in the AI community.

By democratizing access to advanced finetuning capabilities, LMFlow stands to significantly influence how LLMs are adapted and applied across various specialized applications.