trlx - Optimize Large Language Models with Advanced Reinforcement Learning Techniques

Transformer Reinforcement Learning X (trlX)

Transformer Reinforcement Learning X (trlX) is a powerful distributed training framework that focuses on fine-tuning large language models using reinforcement learning. What makes trlX stand out is its ability to work with either a provided reward function or a reward-labeled dataset to enhance model performance.

Key Features

Support for Large Models: trlX is equipped to handle extremely large models, up to 20 billion parameters. For models exceeding this size, it utilizes NVIDIA NeMo-backed trainers, which leverage efficient parallelism techniques for scalability.
Compatibility with Hugging Face Models: It offers support for Hugging Face models through Accelerate-backed trainers, enabling the fine-tuning of popular language models such as facebook/opt-6.7b, EleutherAI/gpt-neox-20b, and google/flan-t5-xxl.
Implemented Algorithms: Currently, it incorporates two robust reinforcement learning algorithms - Proximal Policy Optimization (PPO) and Implicit Language Q-Learning (ILQL), both supported by Accelerate and NeMo Trainers.

How to Get Started

Installation

To install trlX, follow these steps:

git clone https://github.com/CarperAI/trlx.git
cd trlx
pip install torch --extra-index-url https://download.pytorch.org/whl/cu118
pip install -e .

Usage Examples

For practical insights, you can delve into the examples or consider trying the provided colab notebooks, such as for Simulacra and Sentiment analysis using GPT2 with ILQL.

Training Your Model

trlX allows users to train their models with versatility by using different data types:

Using a Reward Function:

trainer = trlx.train('gpt2', reward_fn=lambda samples, **kwargs: [sample.count('cats') for sample in samples])

Using a Reward-Labeled Dataset:

trainer = trlx.train('EleutherAI/gpt-j-6B', samples=['dolphins', 'geese'], rewards=[1.0, 100.0])

Using a Prompt-Completion Dataset:

trainer = trlx.train('gpt2', samples=[['Question: 1 + 2 Answer:', '3'], ...])

Fine-Tuning and Configuration

Users can configure hyperparameters for optimized training, and adjustments can be made to manage memory usage effectively. The ability to save the trained model to a Hugging Face pretrained language model is also provided, making it ready for easy deployment.

Advanced Features

Distributed Training: trlX supports distributed training through tools like Accelerate and NeMo-Megatron, offering seamless scalability for larger training tasks.
Hyperparameter Sweep: Utilize Ray Tune to optimize hyperparameters efficiently across your training jobs.
Benchmark and Logging: Users can benchmark their frameworks against the trlX main branch or adjust logging verbosity to manage the flow of information during training.

Community and Contribution

The trlX project welcomes contributions from the AI community and draws its inspiration from projects like trl. Those interested in contributing can follow the provided guidelines and refer to the detailed documentation.

Citation and Acknowledgments

For academic contexts, trlX can be cited with the provided BibTeX reference, acknowledging the comprehensive work of the authors in developing this advanced reinforcement learning framework.

In summary, trlX is a versatile and high-capacity framework well-equipped for the robust training and fine-tuning of massive language models using cutting-edge reinforcement learning techniques.