重要提示

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

重要提示

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

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

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

SteerLM 2.0：用于属性条件语言模型对齐的迭代训练#

SteerLM 2.0 是一种新颖的方法，用于对齐大型语言模型 (LLM)，以生成具有所需属性值的响应，它建立在原始 SteerLM 方法 [1] 之上。虽然 SteerLM 进行属性条件监督式微调 (SFT) 以引导 LLM 输出，但 SteerLM 2.0 引入了迭代训练程序，以显式强制生成的响应遵循所需的属性分布。

概述#

SteerLM 2.0 的目标是训练一个模型 \(Q_\theta(y|a, x)\)，它可以生成以提示 \(x\) 和所需属性 \(a\) 为条件的响应 \(y\)，同时逼近从属性预测模型 \(P(a|x, y)\) 和无条件响应模型 \(P(y|x)\) 导出的最优条件分布 \(P(y|a, x)\)。SteerLM 2.0 通过最小化 \(P(y|a, x)\) 和 \(Q_\theta(y|a, x)\) 之间的 Kullback-Leibler (KL) 散度来实现此目的

\[\min_\theta \mathbb{E}_{a, x} D_{KL}(P(y|a, x) || Q_\theta(y|a, x))\]

可以使用来自初始 SteerLM 模型 \(Q'(y|a, x)\) 的样本来优化此 KL 散度损失，从而实现高效的梯度估计程序（有关推导，请参见 [2]）。

方法详情#

构建最优条件分布 \(P(y|a, x)\)：使用贝叶斯规则和属性预测模型 \(P(a|x, y)\)，我们可以将最优条件分布导出为

\[P(y|a, x) \propto P(a|x, y) P(y|x)\]

训练 SteerLM 2.0 模型 \(Q_\theta(y|a, x)\)：SteerLM 2.0 模型 \(Q_\theta(y|a, x)\) 被训练为通过最小化使用来自初始 SteerLM 模型 \(Q'(y|a, x)\) 的样本的 KL 散度损失来逼近 \(P(y|a, x)\)。梯度估计为

\[\nabla_\theta L \approx -\sum_{y_i \sim Q'(y|a, x)} (w'_i - b'_i) \nabla_{\theta} \log Q_{\theta}(y_i|a, x)\]

其中 \(w'_i\) 和 \(b'_i\) 是分别以 \(P(y|a, x)\) 和用于稳定优化的基线为目标的归一化重要性权重。（有关详细信息，请参见 [2]）。

迭代训练（可选）：SteerLM 2.0 可以迭代进行（例如，\(n=2\)），在每次迭代后使用优化的策略来采样响应并训练改进的策略。在每次迭代中，从当前模型中采样多个不同的响应，并用于下一轮训练。

通过在此损失上迭代训练，SteerLM 2.0 可以学习生成更符合给定提示 \(x\) 的指定属性值 \(a\) 的响应 \(y\)。

训练 SteerLM 2.0 模型#

准备训练数据集#

SteerLM 2.0 需要特定的数据格式才能有效地训练模型。根据 SteerLM 2.0 方法，需要以下组件

  • 一个 SFT 模型 \(P(y|x)\)，用于生成给定提示 \(x\) 的响应 \(y\)

  • 一个原始 SteerLM 模型 \(Q'(y|a, x)\)，用于生成以属性 \(a\) 和提示 \(x\) 为条件的响应 \(y\)

SteerLM 2.0 模型 \(Q_\theta(y|a, x)\) 使用来自 \(Q'(y|a, x)\) 的权重进行初始化，并进行优化以逼近从属性预测模型 \(P(a|x, y)\) 和无条件响应模型 \(P(y|x)\) 导出的最优条件分布 \(P(y|a, x)\)。

为了方便此训练过程，提出了特定的数据格式

{
"system": "system prompt",
"prompt_turns": [
   {"from": "User", "value": "x_user_turn_1"},
   {"from": "Assistant", "value": "x_assistant_turn_1"},
   {"from": "User", "value": "x_user_turn_2"}
],
"label": "a",
"responses": [
   {
      "from": "Assistant",
      "value": "y_1",
      "log(P(a|x,y))": "v1",
      "log(P(y|x))": "v2",
      "log(Q(y|a,x))": "v3"
   },
   {
      "from": "Assistant",
      "value": "y_2",
      "log(P(a|x,y))": "v1",
      "log(P(y|x))": "v2",
      "log(Q(y|a,x))": "v3"
   },
   ...
   {
      "from": "Assistant",
      "value": "y_n",
      "log(P(a|x,y))": "v1",
      "log(P(y|x))": "v2",
      "log(Q(y|a,x))": "v3"
   }
]
}

对于给定的属性字符串 a 和提示 x（从提示轮次和系统轮次构建），采样 n 个响应 \(y_i\)。为了计算损失，需要以下值

  • \(\log P(a|y_i, x)\)：属性预测模型的输出对数概率，表示给定提示 x 和响应 \(y_i\) 的属性 a

  • \(\log P(y_i|x)\)：无条件响应模型的输出对数概率，表示给定提示 x 的响应 \(y_i\)

  • \(\log Q'(y_i|a, x)\)：原始 SteerLM 模型的输出对数概率，表示给定属性 a 和提示 x 的响应 \(y_i\)

这些值分别作为 log(P(a|x,y))、log(P(y|x)) 和 log(Q(y|a,x)) 提供，用于每个采样的响应 \(y_i\)。

训练示例#

通过以此格式组织数据，可以有效地训练 SteerLM 2.0 模型，以生成符合所需属性值的响应，同时逼近最优条件分布 \(P(y|a, x)\)。以下是启动 SteerLM 2.0 训练的示例

python examples/nlp/gpt/train_steerlm2.py \
     trainer.num_nodes=32 \
     trainer.devices=8 \
     trainer.precision=bf16 \
     trainer.sft.limit_val_batches=40 \
     trainer.sft.max_epochs=1 \
     trainer.sft.max_steps=800 \
     trainer.sft.val_check_interval=800 \
     trainer.sft.save_interval=800 \
     model.megatron_amp_O2=True \
     model.restore_from_path=/path/to/steerlm1/model \
     model.tensor_model_parallel_size=8 \
     model.pipeline_model_parallel_size=2 \
     model.optim.lr=6e-6 \
     model.optim.name=distributed_fused_adam \
     model.optim.weight_decay=0.01 \
     model.optim.sched.constant_steps=200 \
     model.optim.sched.warmup_steps=1 \
     model.optim.sched.min_lr=5e-6 \
     model.answer_only_loss=True \
     model.activations_checkpoint_granularity=selective \
     model.activations_checkpoint_method=uniform \
     model.steerlm2.micro_batch_size=2 \
     model.steerlm2.forward_micro_batch_size=2 \
     model.data.chat=True \
     model.data.num_workers=0 \
     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\>\' \
     model.data.train_ds.max_seq_length=4096 \
     model.data.train_ds.micro_batch_size=1 \
     model.data.train_ds.global_batch_size=128 \
     model.data.train_ds.file_path=data/oasst/train_labeled_2ep.jsonl \
     model.data.train_ds.index_mapping_dir=/indexmap_dir \
     model.data.train_ds.add_eos=False \
     model.data.train_ds.hf_dataset=True \
     model.data.validation_ds.max_seq_length=4096 \
     model.data.validation_ds.file_path=data/oasst/val_labeled.jsonl \
     model.data.validation_ds.micro_batch_size=1 \
     model.data.validation_ds.global_batch_size=128 \
     model.data.validation_ds.index_mapping_dir=/indexmap_dir \
     model.data.validation_ds.add_eos=False \
     model.data.validation_ds.hf_dataset=True \
     exp_manager.create_wandb_logger=True \
     exp_manager.wandb_logger_kwargs.project=steerlm \
     exp_manager.wandb_logger_kwargs.name=acsft_training \
     exp_manager.explicit_log_dir=/results/acsft_70b \
     exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True

/path/to/steerlm1/model 是初始 SteerLM 模型的路径。有关训练初始 SteerLM 模型的详细信息，请参阅 SteerLM 文档。

推理#

由于 SteerLM 2.0 模型是原始 SteerLM 模型的扩展，因此推理过程类似。有关更多详细信息，请参阅 SteerLM 文档。

参考文献#

在本页中