DOCSEARCHbyAlgolia
Submit Docs

trlX Documentation

repository·main·Indexed 26 days ago

https://github.com/carperai/trlx

A distributed training framework for fine-tuning large language models using reinforcement learning (RL) with reward functions or reward-labeled datasets. trlX supports algorithms such as Proximal Policy Optimization (PPO), Implicit Language Q-Learning (ILQL), and Supervised Fine-Tuning (SFT). It provides integration with Hugging Face Accelerate and NVIDIA NeMo, including support for NeMo Megatron models and Triton Inference Server for external reward model hosting.

Tokens
7.5K
Snippets
13
Records
25
Agent score
89%

What's inside trlX

  1. Overview of trlX Trainers

    main

    trlX provides several trainer implementations categorized by their underlying execution framework. Developers can choose between abstract base trainers, Accelerate-based trainers for distributed training via Hugging Face Accelerate, or NeMo-based trainers for NVIDIA NeMo integration.

    Trainer Categories

    • Abstract Trainers: Base classes defining the core reinforcement learning trainer interface.
    • Accelerate Trainers: Implementations optimized for use with the accelerate library, supporting various RL algorithms:
      • AcceleratePPOTrainer: For Proximal Policy Optimization.
      • AccelerateILQLTrainer: For Implicit Q-Learning.
      • AccelerateSFTTrainer: For Supervised Fine-Tuning.
    • NeMo Trainers: Implementations designed for NVIDIA NeMo integration:
      • NeMoPPOTrainer
      • NeMoILQLTrainer
      • `NeMoSFTTrainer"

  2. Overview of trlX reinforcement learning algorithms and backends

    main

    trlX is a library designed for training large language models using reinforcement learning. It supports two primary RL algorithms:

    • PPO (Proximal Policy Optimization): Used for online training.
    • ILQL (Implicit Language Q-Learning): Used for offline training.

    For distributed training, trlX supports the following backends:

    • Huggingface Accelerate
    • NVIDIA NeMo

  3. Use trlx Pipelines for data accumulation and conversion

    main
    Pipelines in trlx are responsible for accumulating training data and converting it into the appropriate formats required for different reinforcement learning algorithms. The library provides several specialized pipeline components including base classes and algorithm-specific implementations like PPO and ILQL.

  4. Setup Reward Model using Triton Inference Server

    main

    If you prefer to host the reward model externally using Triton Inference Server instead of letting the training script instantiate it locally, follow these steps:

    1. Convert the model to a Triton-compatible format and create a model_store.
    2. Build the Triton Docker image (if not using an existing one).
    3. Start the Triton Server pointing to your model_store.
    4. Launch training by setting the TRITON_HOST environment variable.
    # 1. Convert model
python to_triton.py --base_model EleutherAI/gpt-j-6B --checkpoint Dahoas/gptj-rm-static --revision 676bfd4d

# 2. Build singularity image
singularity build --sandbox tritonserver-pyt.sif docker://nvcr.io/nvidia/tritonserver:22.08-pyt-python-py3

# 3. Start Triton Server
SINGULARITYENV_CUDA_VISIBLE_DEVICES=7 singularity run --nv --bind model_store:/model_store tritonserver-pyt.sif tritonserver --model-repository=/model_store &

# 4. Launch training
export TRITON_HOST=localhost:8001/gptj-rm-static
accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml ppo_hh.py
    # convert the model and create a config and a folder `model_store` structured for Triton
python to_triton.py --base_model EleutherAI/gpt-j-6B --checkpoint Dahoas/gptj-rm-static --revision 676bfd4d

# convert the docker image (skip this if you use docker instead)
singularity build --sandbox tritonserver-pyt.sif docker://nvcr.io/nvidia/tritonserver:22.08-pyt-python-py3

# start Triton Server pointing to the `model_store` containing the reward model
SINGULARITYENV_CUDA_VISIBLE_DEVICES=7 singularity run --nv --bind model_store:/model_store tritonserver-pyt.sif tritonserver --model-repository=/model_store &

# Launch training:

# set model's url and replace the name after the slash if you use a different checkpoint
export TRITON_HOST=localhost:8001/gptj-rm-static
accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml ppo_hh.py

  5. Setup NeMo Megatron environment

    main

    Setting up NeMo Megatron requires specific versions of NeMo and Apex. Currently, only NeMo r1.15.0 is supported.

    Prerequisites

    1. Install conda (or mamba/micromamba).
    2. If on an HPC cluster, use srun to access a GPU compute node:
      srun --pty bash -i

    Installation Steps

    1. Create Conda Environment: Use the provided env.yaml (see conda env export in the source) to create the environment:

      conda env create -f env.yaml

    2. Install NeMo:

      git clone https://github.com/NVIDIA/NeMo/
cd NeMo
git checkout r1.15.0
pip install '.[all]'

    3. Install Apex:

      git clone https://github.com/NVIDIA/apex/
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" --global-option="--distributed_adam" --global-option="--deprecated_fused_adam" ./
    srun --pty bash -i

  6. Use pretrained NeMo models

    main

    To use NeMo models in .nemo format (e.g., NeMo Megatron-GPT-20B), follow these steps:

    1. Download and extract the model:

      tar xvf nemo_gpt20B_bf16_tp4.nemo

      This extracts the model weights and the model configuration.

    2. Configure the training parameters:

      • Set train.trainer_kwargs.pretrained_model to the path of the directory containing the extracted parameters.
      • Ensure that the model hyperparameters in train.trainer_kwargs.megatron_cfg match the ones found in the extracted model config.
    tar xvf nemo_gpt20B_bf16_tp4.nemo

  7. Install NVIDIA NeMo (Optional Backend)

    main

    trlX supports NVIDIA NeMo as an optional distributed backend. To install NeMo version v1.17.0, clone the NeMo repository and checkout the specific commit d3017e4.

    $ git clone https://github.com/NVIDIA/NeMo/
$ cd NeMo
$ git checkout d3017e4
$ pip install -e '.[all]'

  8. Configure trlX training using configuration objects

    main

    Training in trlX requires passing a set of specific configuration objects. Depending on your setup, you will need to instantiate and provide the following configs:

    • TrainConfig: General training configuration.
    • ModelConfig: Configuration for the model architecture.
    • TokenizerConfig: Configuration for the tokenizer.
    • OptimizerConfig: Configuration for the optimizer.
    • SchedulerConfig: Configuration for the learning rate scheduler.
    • MethodConfig: Algorithm-specific configuration (e.g., for PPO, ILQL, or SFT).

    For specific algorithms, use their dedicated configuration classes:

    • PPOConfig: For Proximal Policy Optimization.
    • ILQLConfig: For Implicit Language Q-Learning.

  9. Launch distributed training with Accelerate

    main

    To run distributed training using Hugging Face Accelerate, you must first configure your training environment. This configuration step only needs to be executed once per training node.

    1. Run the configuration wizard:
      $ accelerate config
    2. Launch your training script:
      $ accelerate launch examples/ppo_sentiments.py

    Alternatively, you can use pre-defined configuration files from the trlX repository:

    $ accelerate launch --config_file configs/accelerate/zero2-bf16.yaml examples/ppo_sentiments.py
    $ accelerate config
$ accelerate launch examples/ppo_sentiments.py

# Or using a specific config file:
$ accelerate launch --config_file configs/accelerate/zero2-bf16.yaml examples/ppo_sentiments.py

  10. Train with NVIDIA NeMo

    main

    To use NVIDIA NeMo for training, your model must be in the NeMo format.

    1. Convert Llama models to NeMo format: You can use the provided conversion script. Specify the source model path, output folder, total tensor parallel (total_tp) degree, and a name.

    2. Start training: You can either execute the python script manually for each GPU or use the provided Slurm sbatch script (configured for -ntasks-per-node=8).

  11. Install trlX

    main

    To install trlX, clone the repository and install the dependencies using pip. It is recommended to install torch with the appropriate CUDA index first.

    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 .