trlx
Free4.3k
FreeFree tier
Inputs: textOutputs: text
About trlx
trlX is a distributed training framework from CarperAI designed for fine-tuning large language models using Reinforcement Learning via Human Feedback (RLHF). It supports training with either a provided reward function or a reward-labeled dataset, and offers trainers for causal and T5-based models up to 20B parameters via Hugging Face Accelerate, as well as NVIDIA NeMo-backed trainers for scaling beyond 20B parameters. The framework implements Proximal Policy Optimization (PPO) and Implicit Language Q-Learning (ILQL) algorithms, and includes a human-in-the-loop data collection library called CHEESE. trlX is open-source and available on GitHub.
Key Features
Distributed training framework for RLHF
Supports Proximal Policy Optimization (PPO) and Implicit Language Q-Learning (ILQL)
Hugging Face Accelerate-backed trainers for models up to 20B parameters
NVIDIA NeMo-backed trainers for models beyond 20B parameters
Train using a reward function or a reward-labeled dataset
CHEESE library for human-in-the-loop data collection
Supports causal and T5-based models (e.g., GPT-2, OPT, GPT-NeoX, Flan-T5)
Configurable hyperparameters and integration with Weights & Biases
Pros & Cons
Pros
- Open-source and freely available on GitHub
- Supports both small and very large language models (up to 20B+ parameters)
- Implements two state-of-the-art RL algorithms (PPO and ILQL)
- Seamless integration with Hugging Face models and tools
- Includes a companion library (CHEESE) for human annotation collection
- Detailed documentation and example notebooks (Colab)
Cons
- Requires significant computational resources for models beyond 20B parameters
- Primarily focused on text-based language models, not multimodal
- Setup and configuration can be complex for beginners
- Limited to reinforcement learning-based fine-tuning, not general-purpose training
Best For
Fine-tuning large language models with reinforcement learning from human feedbackSentiment control and style transfer in text generationTraining reward models for downstream RL applicationsSimulating agent behavior with Implicit Language Q-LearningPrompt-completion learning for instruction following
FAQ
What is trlX?
trlX is a distributed training framework from CarperAI for fine-tuning large language models with Reinforcement Learning via Human Feedback (RLHF). It supports PPO and ILQL algorithms and can handle models up to 20B parameters with Hugging Face Accelerate trainers and beyond with NVIDIA NeMo trainers.
What RL algorithms does trlX implement?
trlX currently implements Proximal Policy Optimization (PPO) and Implicit Language Q-Learning (ILQL).
How can I train a model with trlX?
You can train using a reward function, a reward-labeled dataset, or a prompt-completion dataset. Example: trainer = trlx.train('gpt2', reward_fn=lambda samples, **kwargs: [sample.count('cats') for sample in samples]).
What models are supported?
trlX supports causal and T5-based language models from Hugging Face, such as GPT-2, OPT-6.7B, GPT-NeoX-20B, and Flan-T5-XXL. For models beyond 20B parameters, NVIDIA NeMo trainers are used.
Is trlX free and open-source?
Yes, trlX is open-source and freely available under the terms of its license on GitHub.