模型
使用 AutoModelForCausalLMWithValueHead
类,TRL 支持 transformers 中所有解码器模型架构,例如 GPT-2、OPT 和 GPT-Neo。此外,使用 AutoModelForSeq2SeqLMWithValueHead
,您可以使用编码器-解码器架构,例如 T5。TRL 还需要参考模型,这些模型是正在训练的模型的冻结副本。使用 create_reference_model
,您可以轻松创建冻结副本,并共享两个模型之间的层以节省内存。
PreTrainedModelWrapper
class trl.PreTrainedModelWrapper
< 源代码 >( pretrained_model = None score_module = None supports_rm_adapter = False rm_adapter_name = None **kwargs )
一个围绕 (transformers.PreTrainedModel
) 的包装类,使其与 (~transformers.PreTrained
) 类兼容,以便保留 (~transformers.PreTrainedModel
) 类的一些属性和方法。
add_and_load_reward_modeling_adapter
< source >( pretrained_model adapter_model_id adapter_name = 'reward_model_adapter' token = None )
添加并加载奖励建模适配器。此方法仅在模型为 PeftModel
且使用 reward_modeling_adapter_id
参数初始化模型时才可以使用,该参数指向奖励建模适配器的 id。后者还需要包含得分头才能产生奖励。
计算给定输入的奖励分数。该方法首先要启用适配器,然后计算奖励分数。之后,模型将禁用奖励建模适配器并重新启用默认的 ppo 适配器。
from_pretrained
< source >( pretrained_model_name_or_path *model_args **kwargs )
参数
- 预训练模型名称或路径 (
str
或transformers.PreTrainedModel
) — 预训练模型的路径或名称。 - *模型参数 (
list
, 可选)) — 传递给底层模型from_pretrained
方法的额外位置参数。 - **kwargs (
dict
, 可选) — 传递给底层模型from_pretrained
方法的额外关键字参数。我们还会预处理 kwargs 以提取特定于transformers.PreTrainedModel
类的参数和特定于 trl 模型的参数。 kwargs 还支持来自peft
库的prepare_model_for_kbit_training
参数。
从 transformers
中的预训练模型实例化一个新模型。预训练模型使用 transformers.PreTrainedModel
类的 from_pretrained
方法加载。特定于 transformers.PreTrainedModel
类的参数将传递给此方法,并从 kwargs
参数中过滤掉。
模型初始化后调用方法。此方法在模型实例化并从检查点加载后调用。它可用于执行其他操作,例如加载 state_dict。
push_to_hub
< source >( *args **kwargs )
将预训练模型推送到 Hub。此方法是对 transformers.PreTrainedModel.push_to_hub
的包装。有关更多信息,请参阅 transformers.PreTrainedModel.push_to_hub
的文档。
save_pretrained
< source >( *args **kwargs )
将预训练模型保存到目录。此方法是对 transformers.PreTrainedModel.save_pretrained
的包装。有关更多信息,请参阅 transformers.PreTrainedModel.save_pretrained
的文档。
返回预训练模型的 state_dict。
AutoModelForCausalLMWithValueHead
除了语言模型头之外,还具有价值头的自回归模型。此类继承自 ~trl.PreTrainedModelWrapper
并包装了 transformers.PreTrainedModel
类。包装类支持经典功能,例如 from_pretrained
、push_to_hub
和 generate
。要调用包装模型的方法,只需操作此类的 pretrained_model
属性。
类属性
- transformers_parent_class (
transformers.PreTrainedModel
) — 包装模型的父类。对于此类,应将其设置为transformers.AutoModelForCausalLM
。 - lm_head_namings (
tuple
) — 用于识别包装模型的语言模型头的字符串元组。对于此类,将其设置为("lm_head", "embed_out")
,但将来可能会为其他模型更改 - supported_args (
tuple
) — 用于识别ValueHead
类支持的参数的字符串元组。当前,支持的参数是- summary_dropout_prob (
float
,可选
, 默认值为None
) —ValueHead
类的 dropout 概率。 - v_head_initializer_range (
float
,可选
, 默认值为0.2
) — 如果选择了特定的初始化策略,则ValueHead
的初始化范围。 - v_head_init_strategy (
str
,optional
, 默认值None
) —ValueHead
的初始化策略。当前支持的策略有:None
— 使用随机分布初始化ValueHead
的权重。这是默认策略。- “normal” — 使用正态分布初始化
ValueHead
的权重。
- summary_dropout_prob (
__init__
< 源代码 >( pretrained_model **kwargs )
初始化模型。
forward
< 源代码 >( input_ids = None past_key_values = None attention_mask = None return_past_key_values = False **kwargs )
参数
- input_ids (torch.LongTensor 形状为 (batch_size, sequence_length)) — 输入序列标记在词汇表中的索引。
- past_key_values (tuple(tuple(torch.FloatTensor)), optional) — 包含模型计算的预先计算的隐藏状态(注意块中的键和值)(参见 past_key_values 输入)以加快顺序解码速度。
- attention_mask (torch.FloatTensor 形状为 (batch_size, sequence_length), optional) — 避免在填充标记索引上执行注意的掩码。掩码值选择在
[0, 1]
中:- 1 表示未屏蔽的标记,
- 0 表示屏蔽的标记。
- return_past_key_values (bool) — 指示是否应返回计算的隐藏状态的标志。
- kwargs (dict, optional) — 传递给包装模型的其他关键字参数。
对包装模型应用前向传递,并返回值头的 logits。
包装模型 `generate` 方法的简单包装器。有关支持参数的更多信息,请参阅包装模型的 generate
方法。
_init_weights
< source > ( **kwargs )
初始化价值头的权重。默认初始化策略是随机的。用户可以通过在调用 `from_pretrained` 时传递 `v_head_init_strategy` 参数来传递不同的初始化策略。支持的策略包括
normal
: 使用正态分布初始化权重。
AutoModelForSeq2SeqLMWithValueHead
类 trl.AutoModelForSeq2SeqLMWithValueHead
< source >( pretrained_model **kwargs )
除了语言模型头之外,还具有价值头的序列到序列模型。此类继承自 `~trl.PreTrainedModelWrapper` 并包装了 `transformers.PreTrainedModel` 类。包装器类支持经典函数,例如 `from_pretrained` 和 `push_to_hub`,还提供一些附加功能,例如 `generate`。
forward
< source > ( input_ids = None past_key_values = None attention_mask = None return_past_key_values = False **kwargs )
我们在封装的模型上调用 `generate`。
我们初始化价值头的权重。
create_reference_model
trl.create_reference_model
< 源代码 > ( model: PreTrainedModelWrapper num_shared_layers: 可选 = None pattern: 可选 = None ) → PreTrainedModelWrapper
创建模型的静态参考副本。请注意,模型将处于 `eval()` 模式。