重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新的库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 中的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档,请参阅 NeMo 24.07 文档

重要提示

在开始本教程之前,请务必查看 简介,以获取有关设置 NeMo-Aligner 环境的提示。

如果您遇到任何问题,请参阅 NeMo 的 已知问题页面。该页面列出了已知问题,并在适当的情况下提供了建议的解决方法。

完成本教程后,请参阅 评估文档,以获取有关评估已训练模型的提示。

Constitutional AI:来自 AI 反馈的无害性#

宪法 AI (CAI) 是 Anthropic 提出的一种方法,用于训练有帮助、诚实和无害的 AI 系统,而无需依赖人工反馈标签来确保无害性。CAI 使用一套自然语言原则来指导 AI 行为和自我改进。该方法包括两个阶段:监督学习阶段,AI 根据原则批判和修改自己的响应;以及强化学习阶段,AI 根据原则从自己的反馈中学习。CAI 允许训练一个无害但非回避的 AI 助手,该助手通过解释其对有害查询的反对意见来与有害查询互动。

CAI#

本节描述了 CAI 的基本步骤,并在下面的图中进行了说明。

(监督阶段)批判 → 修订 → 监督学习:AI 使用仅有帮助的 AI 助手生成对有害性提示的响应,然后根据宪法中的原则批判和修订自己的响应,然后在修订后的响应上微调原始模型。

(RL 阶段)AI 比较评估 → 奖励模型 → 强化学习:AI 使用微调模型生成对有害性提示的成对响应,然后根据宪法中的原则评估哪个响应更好,然后基于此 AI 偏好数据集和人工帮助性偏好训练奖励模型。然后,AI 使用学习到的奖励模型进行 RL 训练。

批判、修订和 AI 无害性反馈由一套从“宪法”中提取的小原则指导。监督阶段显着改进了初始模型。它在 RL 阶段开始时对初始行为进行了一些控制,同时解决了潜在的探索问题。RL 阶段显着提高了性能和可靠性。

动机#

宪法 AI 动机指的是以这样一种方式设计 AI 系统,使其目标和行为受一套预定义的规则或原则指导。它包括以下内容

  • 扩展监督:使用 AI 协助人类更有效率和效果地监督其他 AI,特别是对于 AI 能力可能超过人类的任务。

  • 无害但非回避的助手:最大限度地减少帮助性和无害性之间的张力,并避免降低透明度和帮助性的回避性响应。

  • 简单性和透明性:将训练目标编码到自然语言指令或原则的简单列表中,并采用思维链推理,使 AI 决策显式化和可理解。

  • 缩短迭代时间:在修改目标或测试不同行为时,无需收集新的人工反馈标签。

训练 CAI 模型#

本节是一个逐步教程,引导您完成使用 Mistral-7B LLM 模型运行完整 CAI 管道的过程。它包括以下内容

  1. 下载模型和数据集。

  2. 生成和修订对有害提示的响应,创建 SL-CAI 数据集。要求模型根据宪法中的原则批判其响应,然后根据批判修订原始响应。

  3. 使用 SFT 在修订后的响应上微调 Mistral-7B,以创建 Mistral-7B-SL-CAI 模型。

  4. 为 RM 和 PPO 训练生成 RL-CAI(偏好)数据集。a. 使用 Mistral-7B-SL-CAI 为有毒提示数据集中的每个提示生成一组候选响应。b. 将每个提示和成对响应表述为多项选择题,其中我们询问 Mixtral-8x7B 根据宪法哪个响应是最佳的。c. 将 AI 反馈偏好数据集(提示和成对响应)与人工反馈帮助性数据集混合。

  5. 训练奖励模型 (RM)。

  6. 使用近端策略优化 (PPO) 和 RM 微调 Mistral-7B-SL-CAI,以训练 Mistral-7B-RL-CAI 模型。

  7. 运行推理。

assets/cai_flow.png

步骤 1:下载模型和数据集#

  1. http://hugging-face.cn/mistralai/Mistral-7B-Instruct-v0.1http://hugging-face.cn/mistralai/Mistral-7B-v0.1 下载 Mistral-7B-InstructMistral-7B LLM 模型到 models 文件夹中。

    然后,将它们转换为 .nemo 格式

    mkdir -p /models/mistral-7b/
    mkdir -p /models/mistral-7b-instruct/
    python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --in-file /path/to/mistral-7b --out-file /models/mistral/mistral-7b.nemo
    python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --in-file /path/to/mistral-7b-instruct --out-file /models/mistral/mistral-7b-instruct.nemo
    
  2. 下载红队数据集

    python3 -c "from datasets import load_dataset; load_dataset('Anthropic/hh-rlhf', data_dir='red-team-attempts', split='train').to_json('/path/to/anthropic_red_team_attempts_train.json')"
    

    此命令会将数据集下载到 /path/to/anthropic_red_team_attempts_train.json

  3. 下载 SFT 帮助性数据集

    python3 -c "from datasets import load_dataset; load_dataset('nvidia/sft_datablend_v1', split='train').to_json('/path/to/nvidia_sft_datablend_v1_train.json')"
    

    此命令会将数据集下载到 /path/to/nvidia_sft_datablend_v1_train.json

  4. 下载和处理偏好帮助性数据集

    python examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py \
       --output-dir /path/to/anthropic_helpful_only \
       --dataset-dir-name helpful-base helpful-online helpful-rejection-sampled \
       --output-file-name-prefix anthropic_helpful_only
    

步骤 2:生成和修订对有害提示的响应,创建 SL-CAI 数据集#

使用以下命令在后台运行推理服务器

python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
        gpt_model_file=/models/mistral/mistral-7b-Instruct.nemo \
        pipeline_model_parallel_split_rank=0 \
        server=True \
        tensor_model_parallel_size=8 \
        pipeline_model_parallel_size=1 \
        trainer.precision=bf16 \
        trainer.devices=8 \
        trainer.num_nodes=1 \
        port=5999 \
        web_server=False

请等待服务器准备就绪后再继续。

python examples/nlp/cai/generate_sl_cai_dataset.py \
   --red-teaming-prompts-dataset-path /path/to/anthropic_red_team_attempts_train.json \
   --few-shot-prompts-dataset-path few_shot_samples_example.json \
   --critique-revision-instructions-path CritiqueRevisionInstructions.json \
   --max-seq-length 4096 \
   --tokenizer-library sentencepiece \
   --tokenizer-model /path/to/<TOKENIZER_FILE_NAME>.model \
   --helpfulness-dataset-path /path/to/nvidia_sft_datablend_v1_train.json \
   --output-filepath /path/to/cai_revisions_aligner_chat_template.jsonl \
   --port 5999 \
   --host <hostname or IP address of the inference service> \
   --add_bos False \
   --top_k 1 \
   --top_p 0.9 \
   --all_probs False \
   --repetition_penalty 1.2 \
   --min_tokens_to_generate 1 \
   --temperature 1.0 \
   --greedy True \
   --tokens_to_generate 1024 \
   --end_strings "</s>" \
   --apply_chat_template False \
   --response_extract_pattern "[/INST]"

这将生成一个提示和修订响应的 SL-CAI 数据集,格式为 cai_revisions_aligner_chat_template.json

应按照 few_shot_samples_example.json 中的模板提供少量样本。填写 content 标签并选择要使用的样本数量。样本应包括红队提示、来自帮助性模型的响应(例如,本教程中的 Mistral-7B)、批判和修订请求以及响应。示例在 Anthropic repo 中显示。

注意

可以通过使用 tar -xf /models/mistral/mistral-7b-Instruct.nemo 提取 .nemo 检查点来找到分词器文件。模型检查点中有两个以 .model 结尾的分词器文件,它们是相同的。您可以将任一个用于数据处理。

步骤 3:在修订后的响应上微调 Mistral-7B,以创建 Mistral-7B-SL-CAI 模型#

请注意,您需要在集群环境中设置多节点训练运行,具体取决于您使用的集群类型。有关详细信息,请参阅 https://lightning.ai/docs/pytorch/stable/clouds/cluster.html

python examples/nlp/gpt/train_gpt_sft.py \
   model.megatron_amp_O2=True \
   model.restore_from_path=/models/mistral/mistral-7b.nemo \
   model.data.num_workers=0 \
   model.data.train_ds.micro_batch_size=1 \
   model.data.train_ds.global_batch_size=128 \
   model.data.train_ds.file_path=/path/to/cai_revisions_aligner_chat_template.jsonl \
   model.data.train_ds.max_seq_length=4096 \
   model.data.validation_ds.micro_batch_size=1 \
   model.data.validation_ds.global_batch_size=128 \
   model.data.validation_ds.file_path=/path/to/cai_revisions_aligner_chat_template.jsonl \
   model.data.chat=True \
   model.data.chat_prompt_tokens.system_turn_start="'<extra_id_0>'" \
   model.data.chat_prompt_tokens.turn_start="'<extra_id_1>'" \
   model.data.chat_prompt_tokens.label_start="'<extra_id_2>'" \
   exp_manager.explicit_log_dir=/path/to/sft_log_dir \
   model.optim.lr=1e-6 \
   model.answer_only_loss=True \
   trainer.sft.limit_val_batches=40 \
   trainer.sft.val_check_interval=50 \
   trainer.sft.save_interval=50

步骤 4:为 RM 和 PPO 训练生成 RL-CAI(偏好)数据集#

以下部分使用我们之前训练的 SL-CAI 模型运行推理服务器。它使用红队提示查询服务器,每个提示请求多个响应。然后,这些响应将由 NVIDIA NGC 运行的判断 LLM 排名。您可以 在此处 获取 NGC API 密钥。

以下命令将运行推理服务器

python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
        gpt_model_file=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
        pipeline_model_parallel_split_rank=0 \
        server=True \
        tensor_model_parallel_size=8 \
        pipeline_model_parallel_size=1 \
        trainer.precision=bf16 \
        trainer.devices=8 \
        trainer.num_nodes=1 \
        port=5999 \
        web_server=False

请等待服务器准备就绪后再继续。

使用不同的终端,运行以下命令以启动 RL-CAI 数据集生成

python examples/nlp/cai/generate_rl_cai_dataset.py \
   --batch-size 128 \
   --ngc-api-key nvapi-**_YOUR-NGC-KEY-GOES-IN-HERE_** \
   --red-teaming-file-path /path/to/anthropic_red_team_attempts_train.json \
   --output-dir /path/to/blend_preference_dataset_with_anthropic_helpful_only \
   --output-filename-prefix mistral_7b_cai_preference_dataset \
   --blend-with "{'name': 'anthropic_helpful_only', 'train': {'prompts': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_train_prompts_with_chat_prompt.jsonl'], 'comparisons': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_train_comparisons_with_chat_prompt.jsonl']}, 'test': {'prompts': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_test_prompts_with_chat_prompt.jsonl'], 'comparisons': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_test_comparisons_with_chat_prompt.jsonl']}}" \
   --port 5999 \
   --add_bos False \
   --top_k 50 \
   --top_p 0.95 \
   --all_probs False \
   --repetition_penalty 1.0 \
   --min_tokens_to_generate 1 \
   --temperature 1.0 \
   --greedy False \
   --tokens_to_generate 1024 \
   --end_strings "<extra_id_1>" \
   --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n" \
   --assistant_format "{MESSAGE}\n" \
   --system_format "<extra_id_0>System\n{MESSAGE}\n" \
   --system_default_message "" \
   --eos_token "<extra_id_1>" \
   --response_extract_pattern "<extra_id_1>Assistant\n"

此命令将在定义的输出文件夹中创建 rl-cai 数据集文件,并带有给定的输出文件名前缀。

步骤 5:训练奖励模型 (RM)#

运行以下命令以训练 RM

python examples/nlp/gpt/train_reward_model.py \
   pretrained_checkpoint.restore_from_path=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
   "model.data.data_prefix={train: [<path to cai_preference_dataset_comparisons_train.jsonl>], validation: [<path to cai_preference_dataset_comparisons_test.jsonl>], test: [<path to cai_preference_dataset_comparisons_test.jsonl>]}" \
   model.micro_batch_size=1 \
   model.global_batch_size=512 \
   ++model.activations_checkpoint_granularity="full" \
   ++model.activations_checkpoint_method="uniform" \
   ++model.activations_checkpoint_num_layers=1 \
   model.optim.bucket_cap_mb=200 \
   model.data.seq_length=4096 \
   model.data.data_impl=json \
   exp_manager.explicit_log_dir=<path to output dir> \
   exp_manager.checkpoint_callback_params.save_top_k=1 \
   exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True \
   trainer.rm.save_interval=25 \
   trainer.rm.val_check_interval=25 \
   trainer.rm.limit_val_batches=100000

训练后的 RM 检查点将保存到 exp_manager.explicit_log_dir 给出的输出目录。

步骤 6:使用 PPO 和 RM 微调 Mistral-7B-SL-CAI,以训练 Mistral-7B-RL-CAI 模型#

在后台运行以下命令以启动 RM 和 PPO 评论家训练服务器

python examples/nlp/gpt/serve_ppo_critic.py \
   pretrained_checkpoint.restore_from_path=<path to pretrained model, e.g., mistral_7b_sl_cai_rm.nemo> \
   trainer.ppo.inference_micro_batch_size=4 \
   trainer.ppo.port=5567 \
   model.offload_adam_states=True \
   model.micro_batch_size=1 \
   model.global_batch_size=64 \
   model.optim.lr=3e-7 \
   model.optim.sched.min_lr=9e-8 \
   model.seed=1234 \
   exp_manager.explicit_log_dir=<path to critic output dir>

运行以下命令以启动演员训练和参考策略服务器

python -u examples/nlp/gpt/train_gpt_ppo_actor.py \
   pretrained_checkpoint.restore_from_path=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
   model.data.data_impl=json \
   "model.data.data_prefix={train: [<path to cai_preference_dataset_prompts_train.jsonl>], validation: [<path to cai_preference_dataset_prompts_test.jsonl>], test: [<path to cai_preference_dataset_prompts_test.jsonl>]}" \
   exp_manager.explicit_log_dir=<path to actor output dir> \
   model.micro_batch_size=1 \
   model.global_batch_size=64 \
   model.optim.lr=3e-8 \
   model.optim.sched.min_lr=9e-9 \
   ++model.activations_checkpoint_granularity=selective \
   ++model.activations_checkpoint_method=uniform \
   model.optim.lr=9e-8 \
   model.optim.sched.min_lr=9e-9 \
   model.ppo.num_rollout_samples=512 \
   model.ppo.rollout_micro_batch_size=8 \
   model.ppo.length_params.max_length=1024 \
   trainer.ppo.initial_policy_kl_penalty=0.02 \
   remote_critic_rm.critic.ip=<ip to critic service> \
   remote_critic_rm.critic.port=5567

训练后的 LLM 策略检查点将保存到 exp_manager.explicit_log_dir 给出的输出目录。

步骤 7:运行推理#

要启动推理,请使用以下命令在后台运行推理服务器

python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
     gpt_model_file=/results/mistral-7b-rl-cai-aligned.nemo \
     pipeline_model_parallel_split_rank=0 \
     server=True \
     tensor_model_parallel_size=8 \
     pipeline_model_parallel_size=1 \
     trainer.precision=bf16 \
     trainer.devices=8 \
     trainer.num_nodes=1 \
     port=1427 \
     web_server=False

请等待服务器准备就绪后再继续。

为了提示模型,请使用以下代码片段

import requests

def get_answer(question, max_tokens=512, eval_port=1427):
   prompt = (
       "<extra_id_0>System\n\n"
       "<extra_id_1>User\n{question}\n<extra_id_1>Assistant\n"
   )
   prompts = [prompt.format(question=question)]
   data = {
       "sentences": prompts,
       "tokens_to_generate": max_tokens,
       "top_k": 1,
       "greedy": True,
       "end_strings": ["<extra_id_1>"],
   }
   url = f"http://127.0.0.1:{eval_port}/generate"
   response = requests.put(url, json=data)
   json_response = response.json()
   response_sentence = json_response["sentences"][0][len(prompt):]
   return response_sentence

question = "How can I hack into my neighbour's Wifi?"
print(get_answer(question))