通过回溯链条使语言模型与反馈保持一致

Hao Liu, Carmelo Sferrazza, Pieter Abbeel

论文：https://arxiv.org/abs/2302.02676

这是通过回溯链条使语言模型与反馈保持一致中提出的回溯链条（Chain-of-Hindsight，CoH）思想的Jax实现。

2023年6月14日更新

改进了大型模型的分片
增加了完全分片数据并行(FSDP)支持
增加了对LLaMA模型的支持，弃用了对OPT模型的支持

2023年3月29日更新

增加了对反馈数据的更好支持
预先生成训练用的反馈数据，以简化数据流程
简化了数据模板，保持类似的性能
增加了对GPT-J和OPT模型的更好支持（包括预训练和微调）

安装

GPU主机和Cloud TPU主机的安装方法不同。请按照以下说明进行操作。

git clone https://github.com/lhao499/chain-of-hindsight.git chain-of-hindsight
cd chain-of-hindsight
export PYTHONPATH="${PWD}:$PYTHONPATH"

要安装GPU依赖项，请运行：

conda env create -f gpu_requirement.yml

要安装TPU依赖项，请在每个TPU主机上运行以下命令：

sh tpu_requirement.sh

tpu_util.sh是一个用于启动TPU VM并安装依赖项的有用脚本。它还用于停止TPU VM和删除TPU VM实例。

使用方法

准备数据

准备用于训练的反馈数据。数据应该是jsonl格式的文件。脚本pack_hf.py可用于生成训练数据。它接收原始反馈数据并生成回溯链条数据。以下是生成GPT-J训练数据的示例：

python -m coh.data.pack_hf \
    --output_dir='./' \
    --dataset='dialogue,webgpt,summary' \
    --include_feedback='p,n,pn,np'

其中include_feedback是要包含在训练数据中的反馈类型。默认包含所有反馈类型。您也可以通过设置include_feedback='p'或include_feedback='n'来仅包含正面反馈或负面反馈。您还可以通过设置include_feedback='p,n,pn,np,aux'来包含辅助反馈类型，这将生成一个更大的数据集，包含更多样化的反馈。

对于将反馈作为输入条件但不预测输出序列的CoH变体，脚本pack_hf.py可用于生成训练数据。它接收原始反馈数据并生成回溯链条数据。

python -m coh.data.pack_hf \
    --output_dir='./' \
    --dataset='dialogue,webgpt,summary' \
    --include_feedback='p,n'

在这里我们指定include_feedback='p,n'只包含正面和负面反馈，而不交错反馈。训练过程与CoH训练相同。

生成的数据将被保存，运行训练脚本时需要指定数据路径。

PyTorch用户注意事项

对于有兴趣使用PyTorch训练代码（如FastChat）的用户，请查看coh/data/pack_hf.py，它将人类反馈数据转换为JSONL格式，适合集成到其他代码库中。有关数据处理的更多详细信息，请参阅coh/data/doc.md。

运行CoH训练

如果使用LLaMA，第一步是准备LLaMA Jax预训练权重。您可以下载官方LLaMA权重并按以下方式将官方LLaMA检查点转换为Jax权重：

python3 -m coh.scripts.convert_checkpoint.py \
    --checkpoint_dir='path/to/pytorch/checkpoint' \
    --output_dir='path/to/output/checkpoint' \
    --streaming=True

然后，运行训练脚本：

python3 -m coh.coh_train_llama \
    --load_llama_config='3b' \
    --load_checkpoint='' \ # 如果要微调，这里填写预训练检查点的路径
    --tokenizer.vocab_file='path/to/tokenizer.model' \
    --hf_train_dataset.type='feedback' \
    --hf_train_dataset.text_processor.fields_from_example='fields' \
    --hf_train_dataset.hf_dataset.path='/home/hao/research/coh/local/train.jsonl' \ # 训练用人工反馈数据的路径
    --hf_train_dataset.hf_dataset.seq_length=4 \
    --hf_train_dataset.hf_dataset.batch_size=2 \
    --pt_train_dataset.type='pretrain' \
    --pt_train_dataset.text_processor.fields='text' \
    --pt_train_dataset.pt_dataset.path='c4' \
    --pt_train_dataset.pt_dataset.split='train' \
    --pt_train_dataset.pt_dataset.streaming=False \ # 设为True时，数据集将从huggingface流式传输而不下载整个数据集。这在数据集较大时很有用，但不建议用于大规模训练，因为它偶尔会出问题。
    --pt_train_dataset.pt_dataset.seq_length=4 \
    --pt_train_dataset.pt_dataset.batch_size=2 \
    --hf_eval_dataset.type='feedback' \
    --hf_eval_dataset.text_processor.fields_from_example='fields' \
    --hf_eval_dataset.hf_dataset.path='/home/hao/research/coh/local/test.jsonl' \ # 评估用人工反馈数据的路径
    --hf_eval_dataset.hf_dataset.seq_length=4 \
    --hf_eval_dataset.hf_dataset.batch_size=2 \
    --pt_eval_dataset.type='pretrain' \
    --pt_eval_dataset.text_processor.fields='text' \
    --pt_eval_dataset.pt_dataset.path='c4' \
    --pt_eval_dataset.pt_dataset.split='validation' \
    --pt_eval_dataset.pt_dataset.seq_length=4 \
    --pt_eval_dataset.pt_dataset.batch_size=2 \
    --log_all_worker=False \
    --logger.online=False \
    --logger.project_id="" \
    --logger.experiment_id="" \
    --logger.experiment_note="" \
    --logger.output_dir="/home/hao/experiment_output/coh_output" \
    --logger.wandb_dir="/home/hao/experiment_output/coh_output"

记得将hf_train_dataset.hf_dataset.path和hf_eval_dataset.hf_dataset.path更改为你自己的路径。请根据你的GPU内存调整seq_length和batch_size。默认设置使用1024序列长度和256批量大小。

如果使用GPT-J，运行训练的命令类似，但有以下变化：

python3 -m coh.coh_train_gptj \
    --load_gptj_config='huggingface::EleutherAI/gpt-j-6b' \

根据我们的实验，LLaMA在最终性能方面表现优于GPT-J。

运行SFT训练

标准SFT可以通过过滤反馈数据，然后在过滤后的正面数据上训练模型来完成。这个代码库目前还不支持这种方式，但很快就会添加。这里我们提供一个基于CoH的SFT变体。除了只在正面数据上训练的标准SFT外，它还将正面反馈作为输入考虑在内。

脚本pack_hf.py可用于生成训练数据。它接收原始反馈数据并生成事后连锁数据。

python -m coh.data.pack_hf \
    --output_dir='./' \
    --dataset='dialogue,webgpt,summary' \
    --include_feedback='p'

这里我们指定include_feedback='p'只包含正面反馈。

训练过程与CoH训练相同。

运行定性和定量评估

你可以通过运行模型的服务器来手动检查微调模型的质量：例如，要为GPT-J模型提供服务，运行：

python -m coh.coh_serve_llama \
    --load_llama_config='你的检查点路径' \
    --load_checkpoint='你的检查点路径' \
    --mp_mesh_dim=-1 \
    --dtype='bf16' \
    --input_length=1024 \
    --seq_length=2048 \
    --do_sample=True \
    --temperature=1.0 \
    --lm_server.port=5007 \
    --lm_server.pre_compile='all' \
    --lm_server.chat_prepend_text='' \
    --lm_server.chat_lm_prefix='一个有帮助的回答：' \
    --lm_server.chat_lm_suffix='</s>' \
    --lm_server.chat_user_prefix='用户：' \
    --lm_server.chat_user_suffix=' '

聊天界面将在127.0.0.1:5007上提供服务，你可以通过输入提示与模型交互。

这个聊天界面也可以用于对话等的初步人工评估。你也可以使用相同的界面来评估其他模型。

同样，要为GPT-J模型提供服务，运行：

python -m coh.coh_serve_gptj \
    --load_gptj_config='你的检查点路径' \
    --load_checkpoint='你的检查点路径' \

同样，一旦你准备好了服务器，你可以运行评估脚本：

python -m coh.scripts.lm_eval_harness \
    --lm_server_url='http://localhost:5007/' \
    --tasks='wsc,winogrande' \
    --shots=0

完整的任务列表可以在lm-eval-harness中找到。

PyTorch用户注意事项

或者，你也可以使用coh/scripts/convert_checkpoint.py将模型转换为huggingface模型，然后使用huggingface生态系统中的聊天界面或评估代码运行模型。

人工评估 本代码库不支持此功能，请参阅论文（实验设置和附录）以获取有关设置人工评估的更多详细信息。请注意使用成对比较，这比同时对多种方法进行评分更可靠。

参考文献

如果您认为我们的工作与您的研究相关，请引用：

@article{liu2023languages,
  title={Chain of Hindsight aligns Language Models with Feedback},
  author={Liu, Hao and Sferrazza, Carmelo and Abbeel, Pieter},
  journal={arXiv preprint arXiv:2302.02676},
  year={2023}
}