trlX is a distributed training framework for fine-tuning large language models using Reinforcement Learning from Human Feedback (RLHF). It supports models up to 20B parameters via Hugging Face Accelerate and larger models using NVIDIA NeMo, offering PPO and ILQL algorithms.

How It Works

trlX leverages Hugging Face Accelerate for efficient distributed training of models up to 20B parameters, and NVIDIA NeMo for scaling beyond that. It supports training with either a custom reward function or a reward-labeled dataset, abstracting the complexities of RL algorithms like PPO and ILQL.

Quick Start & Requirements

• Install via pip install -e . after cloning the repository.
• Requires PyTorch with CUDA 11.8 (pip install torch --extra-index-url https://download.pytorch.org/whl/cu118).
• Supports distributed training via accelerate launch.
• Documentation: 📖 Documentation
• Colab Notebooks: Simulacra (GPT2, ILQL), Sentiment (GPT2, ILQL)

Highlighted Details

• Supports fine-tuning causal and T5-based LLMs up to 20B parameters with Accelerate.
• Integrates NVIDIA NeMo for efficient parallelism on models >20B parameters.
• Implements Proximal Policy Optimization (PPO) and Implicit Language Q-Learning (ILQL).
• Offers a human-in-the-loop data collection library called CHEESE.

Maintenance & Community

• Developed by CarperAI.
• Inspired by Hugging Face's trl library.
• Citable publication available: trlX: A Framework for Large Scale Reinforcement Learning from Human Feedback

Licensing & Compatibility

• The repository does not explicitly state a license in the provided README.

Limitations & Caveats

• The README does not specify a license, which may impact commercial use or closed-source integration.
• Setup for NeMo-based training requires following separate NeMo README instructions.