快速开始

安装

要使用OpenRLHF，首先启动docker容器（推荐），然后在docker容器内pip install openrlhf：

# Launch the docker container
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN -v $PWD:/openrlhf nvcr.io/nvidia/pytorch:24.07-py3 bash
pip uninstall xgboost transformer_engine flash_attn pynvml -y

# pip install
pip install openrlhf

# If you want to use vLLM acceleration (To install vLLM 0.7.2)
pip install openrlhf[vllm]
# latest vLLM is also supported
pip install openrlhf[vllm_latest]

# pip install the latest version
pip install git+https://github.com/OpenRLHF/OpenRLHF.git

# Or git clone
git clone https://github.com/OpenRLHF/OpenRLHF.git
cd OpenRLHF
pip install -e .

注意

我们推荐使用 vLLM 0.7.2+。 我们还提供了 vLLM 的 Dockerfiles 和 Nvidia-Docker 的一键安装脚本。

准备数据集

OpenRLHF 在我们的数据集类中提供了多种数据处理方法。 例如在 Prompt Dataset 中：

def preprocess_data(data, input_template=None, input_key="input", apply_chat_template=None) -> str:
   if apply_chat_template:
      chat = data[input_key]
      if isinstance(chat, str):
            chat = [{"role": "user", "content": chat}]
      prompt = apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
   else:
      prompt = data[input_key]
      if input_template:
            prompt = input_template.format(prompt)
   return prompt

  • 我们可以使用 --input_key 来指定输入数据集的 JSON key name，例如 --prompt_data {name or path} (PPO) 或 --dataset {name or path}，并使用 --apply_chat_template 来利用 Huggingface Tokenizer 中的 chat_template。

  • 如果你不想使用 --apply_chat_template，你可以使用 --input_template 代替，或者提前离线预处理数据集。

  • OpenRLHF 还支持使用 --prompt_data_probs 0.1,0.4,0.5 (PPO) 或 --dataset_probs 0.1,0.4,0.5 混合多个数据集。

聊天模板如何工作：

dataset = [{"input_key": [
   {"role": "user", "content": "Hello, how are you?"},
   {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
   {"role": "user", "content": "I'd like to show off how chat templating works!"},
]}]

tokenizer.apply_chat_template(dataset[0]["input_key"], tokenize=False)

"<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s> [INST] I'd like to show off how chat templating works! [/INST]"

如何指定训练和测试数据集？

您可以使用data_type@data_dir格式来指定它。例如，数据集可以设置为--dataset json@./data。

data
├── test.jsonl
└── train.jsonl

注意

默认情况下，我们使用train和test作为分割来区分来自Huggingface的训练和测试数据集。 JSON key选项取决于具体的数据集。 参见Reward Dataset和SFT Dataset

预训练模型

OpenRLHF的模型检查点与HuggingFace模型完全兼容。您可以使用--pretrain、--reward_pretrain和--critic_pretrain指定模型名称或路径。 我们在HuggingFace OpenRLHF上提供了一些预训练的检查点和数据集。

不使用Ray的PPO

然后你可以使用我们在examples目录中提供的启动脚本，或者使用以下命令开始训练：

deepspeed --module openrlhf.cli.train_ppo \
   --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
   --reward_pretrain OpenRLHF/Llama-3-8b-rm-mixture \
   --save_path ./checkpoint/llama-3-8b-rlhf \
   --save_steps -1 \
   --logging_steps 1 \
   --eval_steps -1 \
   --micro_train_batch_size 2 \
   --train_batch_size 128 \
   --micro_rollout_batch_size 4 \
   --rollout_batch_size 1024 \
   --max_epochs 1 \
   --prompt_max_len 1024 \
   --generate_max_len 1024 \
   --zero_stage 2 \
   --bf16 \
   --actor_learning_rate 5e-7 \
   --critic_learning_rate 9e-6 \
   --init_kl_coef 0.01 \
   --prompt_data OpenRLHF/prompt-collection-v0.1 \
   --input_key context_messages \
   --apply_chat_template \
   --max_samples 100000 \
   --normalize_reward \
   --adam_offload \
   --flash_attn \
   --gradient_checkpointing \
   --use_wandb {wandb_token}

Nvidia-Docker 一键安装脚本

# remove old docker
sudo apt-get autoremove docker docker-ce docker-engine docker.io containerd runc
dpkg -l |grep ^rc|awk '{print $2}' |sudo xargs dpkg -P
sudo apt-get autoremove docker-ce-*
sudo rm -rf /etc/systemd/system/docker.service.d
sudo rm -rf /var/lib/docker

# install docker
curl https://get.docker.com | sh \
&& sudo systemctl --now enable docker

# install nvidia-docker
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
      && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
      && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker

sudo groupadd docker
sudo usermod -aG docker $USER
newgrp docker
docker ps