重要提示

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

NeMo 大型语言模型 API#

预训练模型类#

class nemo.collections.nlp.models.language_modeling.megatron_base_model.MegatronBaseModel(*args: Any, **kwargs: Any)#

基类:NLPModel

Megatron 基类。所有 NeMo Megatron 模型都从此类继承。

  • 为 nemo 启用模型并行世界。

  • 开启所有 NVIDIA 优化。

  • 如果 cfg.tokenizer 可用,它将加载分词器并将词汇表填充到张量模型并行的正确大小。

  • 如果使用分布式优化器,则配置为与 O2 级别优化和/或模型并行兼容。

  • 执行梯度裁剪:grad_clip_pl_default 触发 PyTorch Lightning 默认实现,with_distributed_adam 触发分布式优化器的实现,megatron_amp_O2 触发主梯度上的梯度裁剪,否则在模型梯度上执行梯度裁剪。

__init__(
cfg: omegaconf.dictconfig.DictConfig,
trainer: lightning.pytorch.trainer.trainer.Trainer,
no_lm_init=True,
)#

所有 NeMo 模型应继承的基类

参数:
  • cfg (DictConfig) –

    配置对象。cfg 对象应具有(可选)以下子配置

    • train_ds - 实例化训练数据集

    • validation_ds - 实例化验证数据集

    • test_ds - 实例化测试数据集

    • optim - 实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

class nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTModel(*args: Any, **kwargs: Any)#

基类:MegatronBaseModel, TextGeneration

Megatron GPT 预训练

generate(
inputs: List[str] | torch.Tensor | List[dict],
length_params: LengthParam,
sampling_params: SamplingParam | None = None,
*,
strategy: TextGenerationStrategy | None = None,
) OutputType#

生成文本的公共方法。

参数:
  • inputs (Union[List[str], Tensor, List[dict]]) –

    可以是以下 3 种类型之一

    1. 字符串列表。列表的每个元素提供输入提示。模型将对其应用分词器。

      例如 [‘句子’, ‘句子2’ … ]

    2. Pytorch 张量元组(context_tokens, context_lengths)。context_tokens 的形状为 (batch_size, seq_length),它是用作生成提示或用作模型编码器输入的批量序列的标记。

      生成模型将跳过分词和填充步骤。context_lengths 的形状为 (batch_size,),它指示每个输入序列的上下文标记的长度。例如 ( torch.tensor([[23,5234,23,35,…], [223,323,23,23232,232,…] …]), torch.tensor([20, 30, …]))

    3. Python 字典对象列表。用于提示/p-tuning 输入,其中一组键值对被转换为模型的输入标记嵌入。

      例如 [{“prompt-tag”: “sentiment”, “sentence”: “this is a good movie”}, {“prompt-tag”: “qa”, “context”: “some context text”, “question”: “a simple question”} … ] 其中 ‘prompt-tag’ 用于识别要解决的 NLP 任务类型。

  • length_params (LengthParam) –

    控制采样长度的字典类型。

    • max_length:int,要生成的序列的最大长度。

    • min_length:int,要生成的序列的最小长度。

    如果为 None,则 max_length 设置为 30,min_length 设置为 None

  • sampling_params (SamplingParam) –

    包含文本采样参数的字典类型。它具有以下键

    • use_greedy:bool,是否使用采样;否则使用贪婪解码

    • top_k:int,为 top-k 过滤保留的最高概率词汇标记的数量。

    • top_p:float,如果设置为 float < 1,则仅保留概率总和达到 top_p 或更高的最可能的标记以进行生成。

    • repetition_penalty:float,重复惩罚的参数。1.0 表示没有惩罚。

    • add_BOS:bool,是否在提示的开头添加 bos 标记

    • all_probs:bool # 是否返回词汇表中所有标记的对数概率

    • compute_logprob:bool # 用于计算所有输入文本的对数概率的标志,推理的非常特殊的情况,默认为 False

    • end_strings:List[str] # 当生成这些标记之一时,生成将停止

    默认为 None,如果为 None,则 use_greedy 将为 “True”。

返回:

它以字典类型生成输出。它具有以下键,

  • sentences:List[str],输出句子

  • tokens:List[List[str]],输出句子分解为标记

  • logprob:List[List[float]],生成的标记的对数概率

  • full_logprob:List[List[float]],词汇表中所有标记的对数概率

  • token_ids:List[List[int]],输出句子标记 ID

  • offsets:List[List[int]] # 文本中标记起始位置的列表

on_load_checkpoint(checkpoint) None#

LightningModule 钩子:https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-load-checkpoint

on_save_checkpoint(checkpoint) None#

LightningModule 钩子:https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-save-checkpoint

setup(stage=None)#

在 DDP 产生后执行的 PTL 钩子。我们在此处设置数据集,因为 megatron 数据集需要 DDP 才能实例化。有关更多信息,请参阅 https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们将数据加载器迭代器函数传递给微批调度器。每个微批的输入批次使用微批 fwd 函数中的数据加载器函数获取。

validation_step(dataloader_iter, dataloader_idx=0)#

我们的数据加载器生成一个微批,然后我们根据全局批大小和模型并行大小从数据加载器中获取多个微批,以生成微批列表。然后使用 megatron-core fwd/bwd 函数通过管道传输微批列表。

class nemo.collections.nlp.models.language_modeling.megatron_bert_model.MegatronBertModel(*args: Any, **kwargs: Any)#

基类:MegatronBaseModel

Megatron Bert 预训练。模型返回 [batch, seq, hidden] 形状

on_load_checkpoint(checkpoint) None#

LightningModule 钩子:https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-load-checkpoint

on_save_checkpoint(checkpoint) None#

LightningModule 钩子:https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-save-checkpoint

setup(stage=None)#

在 DDP 产生后执行的 PTL 钩子。我们在此处设置数据集,因为 megatron 数据集需要 DDP 才能实例化。有关更多信息,请参阅 https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

class nemo.collections.nlp.models.language_modeling.megatron_bart_model.MegatronBARTModel(*args: Any, **kwargs: Any)#

基类:MegatronT5Model

Megatron BART 预训练

setup(stage=None)#

在 DDP 产生后执行的 PTL 钩子。我们在此处设置数据集,因为 megatron 数据集需要 DDP 才能实例化。有关更多信息,请参阅 https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们的数据加载器生成一个微批,然后我们根据全局批大小和模型并行大小从数据加载器中获取多个微批,以生成微批列表。批次应为微批列表,这些微批应位于 CPU 上。然后,在管道期间将微批移动到 GPU。然后使用 megatron-core fwd/bwd 函数通过管道传输微批列表。

validation_step(dataloader_iter)#

return_values - 如果给定,则返回带有给定键和相应值的字典

class nemo.collections.nlp.models.language_modeling.megatron_retrieval_model.MegatronRetrievalModel(*args: Any, **kwargs: Any)#

基类:MegatronBaseModel, TextGeneration

Megatron 检索增强语言模型

generate(
inputs: List[str] | torch.Tensor | List[dict],
length_params: LengthParam,
sampling_params: SamplingParam | None = None,
**args,
) OutputType#

生成文本的公共方法。

参数:
  • inputs (Union[List[str], Tensor, List[dict]]) –

    可以是以下 3 种类型之一

    1. 字符串列表。列表的每个元素提供输入提示。模型将对其应用分词器。

      例如 [‘句子’, ‘句子2’ … ]

    2. Pytorch 张量元组(context_tokens, context_lengths)。context_tokens 的形状为 (batch_size, seq_length),它是用作生成提示或用作模型编码器输入的批量序列的标记。

      生成模型将跳过分词和填充步骤。context_lengths 的形状为 (batch_size,),它指示每个输入序列的上下文标记的长度。例如 ( torch.tensor([[23,5234,23,35,…], [223,323,23,23232,232,…] …]), torch.tensor([20, 30, …]))

    3. Python 字典对象列表。用于提示/p-tuning 输入,其中一组键值对被转换为模型的输入标记嵌入。

      例如 [{“prompt-tag”: “sentiment”, “sentence”: “this is a good movie”}, {“prompt-tag”: “qa”, “context”: “some context text”, “question”: “a simple question”} … ] 其中 ‘prompt-tag’ 用于识别要解决的 NLP 任务类型。

  • length_params (LengthParam) –

    控制采样长度的字典类型。

    • max_length:int,要生成的序列的最大长度。

    • min_length:int,要生成的序列的最小长度。

    如果为 None,则 max_length 设置为 30,min_length 设置为 None

  • sampling_params (SamplingParam) –

    包含文本采样参数的字典类型。它具有以下键

    • use_greedy:bool,是否使用采样;否则使用贪婪解码

    • top_k:int,为 top-k 过滤保留的最高概率词汇标记的数量。

    • top_p:float,如果设置为 float < 1,则仅保留概率总和达到 top_p 或更高的最可能的标记以进行生成。

    • repetition_penalty:float,重复惩罚的参数。1.0 表示没有惩罚。

    • add_BOS:bool,是否在提示的开头添加 bos 标记

    • all_probs:bool # 是否返回词汇表中所有标记的对数概率

    • compute_logprob:bool # 用于计算所有输入文本的对数概率的标志,推理的非常特殊的情况,默认为 False

    • end_strings:List[str] # 当生成这些标记之一时,生成将停止

    默认为 None,如果为 None,则 use_greedy 将为 “True”。

返回:

它以字典类型生成输出。它具有以下键,

  • sentences:List[str],输出句子

  • tokens:List[List[str]],输出句子分解为标记

  • logprob:List[List[float]],生成的标记的对数概率

  • full_logprob:List[List[float]],词汇表中所有标记的对数概率

  • token_ids:List[List[int]],输出句子标记 ID

  • offsets:List[List[int]] # 文本中标记起始位置的列表

setup(stage=None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时,这会在每个进程上调用。

参数:

stage – fit、validate、test 或 predict

class nemo.collections.nlp.models.language_modeling.megatron_t5_model.MegatronT5Model(*args: Any, **kwargs: Any)#

基类:MegatronLMEncoderDecoderModel

Megatron T5 预训练

complete(request: Dict)#

在推理模式下自动调用语言模型

参数:

request – 具有以下字段的字典 * prompt:模型应完成的文本字符串。 * tokens_to_generate:在执行提示完成时要生成的标记数。

返回:

具有以下字段的 Python 字典
  • prompt:提示的原始文本

  • tokenized_prompt:来自提示的 (str) 标记列表

  • completion:具有以下子字段的 Python 字典
    • tokens:包含完成的三元组(token、token_id、log_prob)列表

    • text:完成文本(作为单个字符串)

返回类型:

response

decode(
tokens_enc,
enc_mask,
num_tokens_to_generate,
encoder_input=None,
tokenizer=None,
enc_output=None,
enc_output_attn_mask=None,
ignore_ids=[],
bos_id=None,
predicted_tokens_dec=None,
batch_data=None,
sampling_method: str = 'greedy-search',
sampling_kwargs: dict = {},
)#
参数:
  • tokens_enc – 形状为 [batch_size, seq_len] 的张量,其中包含输入标记。

  • enc_mask – 形状为 [batch_size, seq_len] 的张量,其中包含输入标记掩码(活动为 1,非活动为 0)。

  • num_tokens_to_generate – 要生成的最大标记数。

  • encoder_input – 形状为 [batch_size, seq_len, hidden_size] 的张量,其中包含编码器隐藏状态(如果给定,则替换 tokens_enc)。

  • tokenizer – 分词器对象。

  • enc_output – 形状为 [batch_size, seq_len, hidden_size] 的张量,其中包含编码器隐藏状态(如果给定,则替换 tokens_enc 和 encoder_input)。

  • enc_output_attn_mask – 形状为 [batch_size, seq_len] 的张量,其中包含编码器注意力掩码(如果给定,则替换 enc_mask)。

  • ignore_ids – 采样时要忽略的标记 ID 列表。

  • bos_id – 句子开头标记的 ID。如果为 None,将使用 tokenizer.bos_id,除非显式设置为其他值。

  • predicted_tokens_dec – 形状为 [batch_size, seq_len] 的张量,其中包含已解码的标记。

  • sampling_method – 解码迭代中使用的采样方法。当前支持的方法是“beam-search”/“greedy-search”/“topkp-sampling”。该参数指定采样函数,该函数接收 logits 张量 [batch_size, vocab_size] 并返回元组(对数概率张量 [batch_size],从 logits [batch_size] 采样的标记 ID 张量)。如果启用了波束搜索,则采样函数返回张量 [batch_size, beam_size]

  • sampling_kwargs – 包含要传递给采样函数的参数的 dict。请参阅方法 get_sampling_token_fn,以查看所选 sampling_method 需要哪些参数。

返回:

预测标记及其对数概率的张量元组 [batch_size, seq_len +1], [batch_size, seq_len]。如果 sampling_method == ‘beam-size’ 且 keep_only_best_tokens 为 False,则张量的形状为 [batch_size, beam_size, seq_len + 1], [batch_size, beam_size, seq_len]

encode(
tokens_enc,
enc_mask,
encoder_input=None,
batch_data=None,
reconfigure_microbatch=True,
)#
参数:
  • tokens_enc – 编码器输入标记

  • enc_mask – 相应的掩码

  • encoder_input – 编码器输入(绕过标记),如果给定,tokens_enc 可以为 None。

  • batch_data – 直接传递给所有隐藏变换和损失。可用于传递其他数据,如类标签。格式未定义,应与所用隐藏模块的预期格式匹配。

setup(stage=None)#

在 DDP 产生后执行的 PTL 钩子。我们在此处设置数据集,因为 megatron 数据集需要 DDP 才能实例化。有关更多信息,请参阅 https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们的数据加载器生成一个微批,然后我们根据全局批大小和模型并行大小从数据加载器中获取多个微批,以生成微批列表。批次应为微批列表,这些微批应位于 CPU 上。然后,在管道期间将微批移动到 GPU。然后使用 megatron-core fwd/bwd 函数通过管道传输微批列表。

validation_step(dataloader_iter)#

return_values - 如果给定,则返回带有给定键和相应值的字典

自定义模型类#

class nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model.MegatronGPTSFTModel(*args: Any, **kwargs: Any)#

基类: NLPAdapterModelMixin, MegatronGPTModel

Megatron GPT 监督式微调

generate(
inputs: List[str] | torch.Tensor | List[dict],
length_params: LengthParam,
sampling_params: SamplingParam | None = None,
*,
strategy: TextGenerationStrategy | None = None,
) OutputType#

生成文本的公共方法。

参数:
  • inputs (Union[List[str], Tensor, List[dict]]) –

    可以是以下 3 种类型之一

    1. 字符串列表。列表的每个元素提供输入提示。模型将对其应用分词器。

      例如 [‘句子’, ‘句子2’ … ]

    2. Pytorch 张量元组(context_tokens, context_lengths)。context_tokens 的形状为 (batch_size, seq_length),它是用作生成提示或用作模型编码器输入的批量序列的标记。

      生成模型将跳过分词和填充步骤。context_lengths 的形状为 (batch_size,),它指示每个输入序列的上下文标记的长度。例如 ( torch.tensor([[23,5234,23,35,…], [223,323,23,23232,232,…] …]), torch.tensor([20, 30, …]))

    3. Python 字典对象列表。用于提示/p-tuning 输入,其中一组键值对被转换为模型的输入标记嵌入。

      例如 [{“prompt-tag”: “sentiment”, “sentence”: “this is a good movie”}, {“prompt-tag”: “qa”, “context”: “some context text”, “question”: “a simple question”} … ] 其中 ‘prompt-tag’ 用于识别要解决的 NLP 任务类型。

  • length_params (LengthParam) –

    控制采样长度的字典类型。

    • max_length:int,要生成的序列的最大长度。

    • min_length:int,要生成的序列的最小长度。

    如果为 None,则 max_length 设置为 30,min_length 设置为 None

  • sampling_params (SamplingParam) –

    包含文本采样参数的字典类型。它具有以下键

    • use_greedy:bool,是否使用采样;否则使用贪婪解码

    • top_k:int,为 top-k 过滤保留的最高概率词汇标记的数量。

    • top_p:float,如果设置为 float < 1,则仅保留概率总和达到 top_p 或更高的最可能的标记以进行生成。

    • repetition_penalty:float,重复惩罚的参数。1.0 表示没有惩罚。

    • add_BOS:bool,是否在提示的开头添加 bos 标记

    • all_probs:bool # 是否返回词汇表中所有标记的对数概率

    • compute_logprob:bool # 用于计算所有输入文本的对数概率的标志,推理的非常特殊的情况,默认为 False

    • end_strings:List[str] # 当生成这些标记之一时,生成将停止

    默认为 None,如果为 None,则 use_greedy 将为 “True”。

返回:

它以字典类型生成输出。它具有以下键,

  • sentences:List[str],输出句子

  • tokens:List[List[str]],输出句子分解为标记

  • logprob:List[List[float]],生成的标记的对数概率

  • full_logprob:List[List[float]],词汇表中所有标记的对数概率

  • token_ids:List[List[int]],输出句子标记 ID

  • offsets:List[List[int]] # 文本中标记起始位置的列表

setup(stage=None)#

在 DDP 产生后执行的 PTL 钩子。我们在此处设置数据集,因为 megatron 数据集需要 DDP 才能实例化。有关更多信息,请参阅 https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们将数据加载器迭代器函数传递给微批调度器。每个微批的输入批次使用微批 fwd 函数中的数据加载器函数获取。

validation_step(dataloader_iter)#

我们的数据加载器生成一个微批,然后我们根据全局批大小和模型并行大小从数据加载器中获取多个微批,以生成微批列表。然后使用 megatron-core fwd/bwd 函数通过管道传输微批列表。

class nemo.collections.nlp.models.language_modeling.megatron_gpt_adapter_model.MegatronGPTAdapterLearningModel(*args: Any, **kwargs: Any)#

基类: MegatronGPTBaseAdapterModel

MegatronGPTAdapterLearningModel 是一个将基础模型 (GPTModel) 与适配器结合的模型。此类仅支持 Houlsby 等人描述的规范适配器训练 (https://arxiv.org/pdf/1902.00751.pdf)

两个适配器被插入到基础 GPT 模型中的每个 Transformer 层。

假设这组适配器随后将针对特定任务进行训练。训练完成后,适配器权重将被保存,并可以重新加载并注入到相同的 GPT 模型中以进行推理。

__init__(
cfg: omegaconf.dictconfig.DictConfig,
trainer: lightning.pytorch.trainer.trainer.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:
  • cfg (DictConfig) –

    配置对象。cfg 对象应具有(可选)以下子配置

    • train_ds - 实例化训练数据集

    • validation_ds - 实例化验证数据集

    • test_ds - 实例化测试数据集

    • optim - 实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

generate(
inputs: List[str] | torch.Tensor | List[dict],
length_params: LengthParam,
sampling_params: SamplingParam | None = None,
batch_size: int | None = 1,
)#

生成文本的公共方法。

参数:
  • inputs (Union[List[str], Tensor, List[dict]]) –

    可以是以下 3 种类型之一

    1. 字符串列表。列表的每个元素提供输入提示。模型将对其应用分词器。

      例如 [‘句子’, ‘句子2’ … ]

    2. Pytorch 张量元组(context_tokens, context_lengths)。context_tokens 的形状为 (batch_size, seq_length),它是用作生成提示或用作模型编码器输入的批量序列的标记。

      生成模型将跳过分词和填充步骤。context_lengths 的形状为 (batch_size,),它指示每个输入序列的上下文标记的长度。例如 ( torch.tensor([[23,5234,23,35,…], [223,323,23,23232,232,…] …]), torch.tensor([20, 30, …]))

    3. Python 字典对象列表。用于提示/p-tuning 输入,其中一组键值对被转换为模型的输入标记嵌入。

      例如 [{“prompt-tag”: “sentiment”, “sentence”: “this is a good movie”}, {“prompt-tag”: “qa”, “context”: “some context text”, “question”: “a simple question”} … ] 其中 ‘prompt-tag’ 用于识别要解决的 NLP 任务类型。

  • length_params (LengthParam) –

    控制采样长度的字典类型。

    • max_length:int,要生成的序列的最大长度。

    • min_length:int,要生成的序列的最小长度。

    如果为 None,则 max_length 设置为 30,min_length 设置为 None

  • sampling_params (SamplingParam) –

    包含文本采样参数的字典类型。它具有以下键

    • use_greedy:bool,是否使用采样;否则使用贪婪解码

    • top_k:int,为 top-k 过滤保留的最高概率词汇标记的数量。

    • top_p:float,如果设置为 float < 1,则仅保留概率总和达到 top_p 或更高的最可能的标记以进行生成。

    • repetition_penalty:float,重复惩罚的参数。1.0 表示没有惩罚。

    • add_BOS:bool,是否在提示的开头添加 bos 标记

    • all_probs:bool # 是否返回词汇表中所有标记的对数概率

    • compute_logprob:bool # 用于计算所有输入文本的对数概率的标志,推理的非常特殊的情况,默认为 False

    • end_strings:List[str] # 当生成这些标记之一时,生成将停止

    默认为 None,如果为 None,则 use_greedy 将为 “True”。

返回:

它以字典类型生成输出。它具有以下键,

  • sentences:List[str],输出句子

  • tokens:List[List[str]],输出句子分解为标记

  • logprob:List[List[float]],生成的标记的对数概率

  • full_logprob:List[List[float]],词汇表中所有标记的对数概率

  • token_ids:List[List[int]],输出句子标记 ID

  • offsets:List[List[int]] # 文本中标记起始位置的列表

setup(stage=None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时,这会在每个进程上调用。

参数:

stage – fit、validate、test 或 predict

state_dict(
destination=None,
prefix=None,
keep_vars=False,
)#

创建一个仅使用适配器参数的 state_dict。这确保了此包装器类将仅检查点适配器权重,而不是基础 GPT 模型的其余部分。

class nemo.collections.nlp.models.language_modeling.megatron_gpt_adapter_model.MegatronGPTInfusedAdapterModel(*args: Any, **kwargs: Any)#

基类: MegatronGPTBaseAdapterModel

MegatronGPTInfusedAdapterModel 是一个将基础模型 (GPTModel) 与“可以抑制和放大内部激活的注入适配器”(称为 IA3)结合的模型。此类支持将 IA3 添加到基于 Transformer 的 LM 中,如 Liu 等人所述 (https://arxiv.org/pdf/2205.05638.pdf)

三个适配器被插入到基础 GPT 模型中的每个 Transformer 层。每个适配器基本上是一个向量,它仅缩放键、值或 ffn 隐藏表示。

假设这组适配器随后将针对特定任务进行训练。训练完成后,适配器权重将被保存,并可以重新加载并注入到相同的 GPT 模型中以进行推理。

__init__(
cfg: omegaconf.dictconfig.DictConfig,
trainer: lightning.pytorch.trainer.trainer.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:
  • cfg (DictConfig) –

    配置对象。cfg 对象应具有(可选)以下子配置

    • train_ds - 实例化训练数据集

    • validation_ds - 实例化验证数据集

    • test_ds - 实例化测试数据集

    • optim - 实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

generate(
inputs: List[str] | torch.Tensor | List[dict],
length_params: LengthParam,
sampling_params: SamplingParam | None = None,
batch_size: int | None = 1,
)#

生成文本的公共方法。

参数:
  • inputs (Union[List[str], Tensor, List[dict]]) –

    可以是以下 3 种类型之一

    1. 字符串列表。列表的每个元素提供输入提示。模型将对其应用分词器。

      例如 [‘句子’, ‘句子2’ … ]

    2. Pytorch 张量元组(context_tokens, context_lengths)。context_tokens 的形状为 (batch_size, seq_length),它是用作生成提示或用作模型编码器输入的批量序列的标记。

      生成模型将跳过分词和填充步骤。context_lengths 的形状为 (batch_size,),它指示每个输入序列的上下文标记的长度。例如 ( torch.tensor([[23,5234,23,35,…], [223,323,23,23232,232,…] …]), torch.tensor([20, 30, …]))

    3. Python 字典对象列表。用于提示/p-tuning 输入,其中一组键值对被转换为模型的输入标记嵌入。

      例如 [{“prompt-tag”: “sentiment”, “sentence”: “this is a good movie”}, {“prompt-tag”: “qa”, “context”: “some context text”, “question”: “a simple question”} … ] 其中 ‘prompt-tag’ 用于识别要解决的 NLP 任务类型。

  • length_params (LengthParam) –

    控制采样长度的字典类型。

    • max_length:int,要生成的序列的最大长度。

    • min_length:int,要生成的序列的最小长度。

    如果为 None,则 max_length 设置为 30,min_length 设置为 None

  • sampling_params (SamplingParam) –

    包含文本采样参数的字典类型。它具有以下键

    • use_greedy:bool,是否使用采样;否则使用贪婪解码

    • top_k:int,为 top-k 过滤保留的最高概率词汇标记的数量。

    • top_p:float,如果设置为 float < 1,则仅保留概率总和达到 top_p 或更高的最可能的标记以进行生成。

    • repetition_penalty:float,重复惩罚的参数。1.0 表示没有惩罚。

    • add_BOS:bool,是否在提示的开头添加 bos 标记

    • all_probs:bool # 是否返回词汇表中所有标记的对数概率

    • compute_logprob:bool # 用于计算所有输入文本的对数概率的标志,推理的非常特殊的情况,默认为 False

    • end_strings:List[str] # 当生成这些标记之一时,生成将停止

    默认为 None,如果为 None,则 use_greedy 将为 “True”。

返回:

它以字典类型生成输出。它具有以下键,

  • sentences:List[str],输出句子

  • tokens:List[List[str]],输出句子分解为标记

  • logprob:List[List[float]],生成的标记的对数概率

  • full_logprob:List[List[float]],词汇表中所有标记的对数概率

  • token_ids:List[List[int]],输出句子标记 ID

  • offsets:List[List[int]] # 文本中标记起始位置的列表

setup(stage=None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时,这会在每个进程上调用。

参数:

stage – fit、validate、test 或 predict

state_dict(
destination=None,
prefix=None,
keep_vars=False,
)#

创建一个仅使用适配器参数的 state_dict。这确保了此包装器类将仅检查点适配器权重,而不是基础 GPT 模型的其余部分。

class nemo.collections.nlp.models.language_modeling.megatron_gpt_prompt_learning_model.MegatronGPTPromptLearningModel(*args: Any, **kwargs: Any)#

基类: MegatronBasePromptLearningModel

用于 prompt-tuning 或 p-tuning 预训练 Megatron GPT 模型的模型类。

Prompt Tuning 直接从预训练的 GPT 模型词汇表中某些 token embeddings 的副本初始化 virtual prompt embeddings,并直接调整这些 embedding 权重。初始化中使用的 token embeddings 由用户在 config 文件中指定。该模型可以一次 prompt-tuned 用于多个任务。virtual prompts 存储在 prompt table 中,可以添加或删除,而不会破坏其他任务的 virtual prompts。

P-tuning 初始化一个 LSTM encoder model,该模型为每个任务生成 virtual prompt embeddings。每个任务共享同一个 encoder。在 ptuning 完成后,学习到的 virtual prompts 可以使用 add_ptuned_prompts_to_prompt_table() 保存到 prompt table 中。因此,如果用户想通过 p-tuning 添加新的 virtual prompt,他们不需要在所有以前的任务上重新训练。这使 p-tuning 具有与 prompt-tuning 相同的任务灵活性。

generate(
inputs: List[str] | torch.Tensor | List[dict],
length_params: LengthParam,
sampling_params: SamplingParam | None = None,
batch_size: int | None = 1,
)#

生成文本的公共方法。

参数:
  • inputs (Union[List[str], Tensor, List[dict]]) –

    可以是以下 3 种类型之一

    1. 字符串列表。列表的每个元素提供输入提示。模型将对其应用分词器。

      例如 [‘句子’, ‘句子2’ … ]

    2. Pytorch 张量元组(context_tokens, context_lengths)。context_tokens 的形状为 (batch_size, seq_length),它是用作生成提示或用作模型编码器输入的批量序列的标记。

      生成模型将跳过分词和填充步骤。context_lengths 的形状为 (batch_size,),它指示每个输入序列的上下文标记的长度。例如 ( torch.tensor([[23,5234,23,35,…], [223,323,23,23232,232,…] …]), torch.tensor([20, 30, …]))

    3. Python 字典对象列表。用于提示/p-tuning 输入,其中一组键值对被转换为模型的输入标记嵌入。

      例如 [{“prompt-tag”: “sentiment”, “sentence”: “this is a good movie”}, {“prompt-tag”: “qa”, “context”: “some context text”, “question”: “a simple question”} … ] 其中 ‘prompt-tag’ 用于识别要解决的 NLP 任务类型。

  • length_params (LengthParam) –

    控制采样长度的字典类型。

    • max_length:int,要生成的序列的最大长度。

    • min_length:int,要生成的序列的最小长度。

    如果为 None,则 max_length 设置为 30,min_length 设置为 None

  • sampling_params (SamplingParam) –

    包含文本采样参数的字典类型。它具有以下键

    • use_greedy:bool,是否使用采样;否则使用贪婪解码

    • top_k:int,为 top-k 过滤保留的最高概率词汇标记的数量。

    • top_p:float,如果设置为 float < 1,则仅保留概率总和达到 top_p 或更高的最可能的标记以进行生成。

    • repetition_penalty:float,重复惩罚的参数。1.0 表示没有惩罚。

    • add_BOS:bool,是否在提示的开头添加 bos 标记

    • all_probs:bool # 是否返回词汇表中所有标记的对数概率

    • compute_logprob:bool # 用于计算所有输入文本的对数概率的标志,推理的非常特殊的情况,默认为 False

    • end_strings:List[str] # 当生成这些标记之一时,生成将停止

    默认为 None,如果为 None,则 use_greedy 将为 “True”。

返回:

它以字典类型生成输出。它具有以下键,

  • sentences:List[str],输出句子

  • tokens:List[List[str]],输出句子分解为标记

  • logprob:List[List[float]],生成的标记的对数概率

  • full_logprob:List[List[float]],词汇表中所有标记的对数概率

  • token_ids:List[List[int]],输出句子标记 ID

  • offsets:List[List[int]] # 文本中标记起始位置的列表

setup(stage=None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时,这会在每个进程上调用。

参数:

stage – fit、validate、test 或 predict

class nemo.collections.nlp.models.language_modeling.megatron_t5_adapter_model.MegatronT5AdapterLearningModel(*args: Any, **kwargs: Any)#

基类: MegatronT5BaseAdapterModel

TODO (@adithyare)

__init__(
cfg: omegaconf.dictconfig.DictConfig,
trainer: lightning.pytorch.trainer.trainer.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:
  • cfg (DictConfig) –

    配置对象。cfg 对象应具有(可选)以下子配置

    • train_ds - 实例化训练数据集

    • validation_ds - 实例化验证数据集

    • test_ds - 实例化测试数据集

    • optim - 实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

setup(stage=None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时,这会在每个进程上调用。

参数:

stage – fit、validate、test 或 predict

state_dict(
destination=None,
prefix=None,
keep_vars=False,
)#

创建一个仅使用适配器参数的 state_dict。这确保了此包装器类将仅检查点适配器权重,而不是基础 GPT 模型的其余部分。

class nemo.collections.nlp.models.language_modeling.megatron_t5_adapter_model.MegatronT5InfusedAdapterModel(*args: Any, **kwargs: Any)#

基类: MegatronT5BaseAdapterModel

MegatronGPTInfusedAdapterModel 是一个将基础模型 (GPTModel) 与“可以抑制和放大内部激活的注入适配器”(称为 IA3)结合的模型。此类支持将 IA3 添加到基于 Transformer 的 LM 中,如 Liu 等人所述 (https://arxiv.org/pdf/2205.05638.pdf)

三个适配器被插入到基础 GPT 模型中的每个 Transformer 层。每个适配器基本上是一个向量,它仅缩放键、值或 ffn 隐藏表示。

假设这组适配器随后将针对特定任务进行训练。训练完成后,适配器权重将被保存,并可以重新加载并注入到相同的 GPT 模型中以进行推理。

__init__(
cfg: omegaconf.dictconfig.DictConfig,
trainer: lightning.pytorch.trainer.trainer.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:
  • cfg (DictConfig) –

    配置对象。cfg 对象应具有(可选)以下子配置

    • train_ds - 实例化训练数据集

    • validation_ds - 实例化验证数据集

    • test_ds - 实例化测试数据集

    • optim - 实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

setup(stage=None)#

在 fit、validate、test 或 predict 的开始时调用。当使用 DDP 时,这会在每个进程上调用。

参数:

stage – fit、validate、test 或 predict

state_dict(
destination=None,
prefix=None,
keep_vars=False,
)#

创建一个仅使用适配器参数的 state_dict。这确保了此包装器类将仅检查点适配器权重,而不是基础 GPT 模型的其余部分。

模块#

class nemo.collections.nlp.modules.common.megatron.module.MegatronModule(*args: Any, **kwargs: Any)#

基类: Module

Megatron 特定的 torch Module 扩展,支持流水线。

class nemo.collections.nlp.modules.common.megatron.module.Float16Module(*args: Any, **kwargs: Any)#

基类: MegatronModule

class nemo.collections.nlp.models.language_modeling.megatron.gpt_model.GPTModel(*args: Any, **kwargs: Any)#

基类: MegatronModule

GPT-2 语言模型。

class nemo.collections.nlp.models.language_modeling.megatron.bert.bert_model.NeMoBertModel(*args: Any, **kwargs: Any)#

基类: MegatronModule

Bert 语言模型。模型返回 [seq, batch, hidden] 形状

class nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder.MegatronTokenLevelEncoderDecoderModule(*args: Any, **kwargs: Any)#

基类: MegatronModule, AdapterModuleMixin

基于 Token 的(输入/输出是 token)编码器-解码器模型(例如 T5 语言模型)。

forward(
enc_input_ids=None,
enc_attn_mask=None,
dec_input_ids=None,
dec_attn_mask=None,
token_type_ids=None,
labels=None,
batch_data=None,
enc_output=None,
enc_output_attn_mask=None,
enc_input=None,
output_enc_hidden_only=False,
)#

返回值是每个 token / 每个维度(即,非折叠损失值)

class nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder.MegatronRetrievalTokenLevelEncoderDecoderModule(
*args: Any,
**kwargs: Any,
)#

基类: MegatronModule

基于 Token 的(输入/输出是 token)检索编码器-解码器模型

forward(
input_ids,
input_attn_mask,
retrieved_ids,
retrieved_attn_mask,
token_type_ids=None,
labels=None,
input_emb=None,
set_inference_key_value_memory=False,
inference_max_sequence_len=None,
neighbors=None,
position_ids=None,
)#

返回值是每个 token / 每个维度(即,非折叠损失值)

数据集#

class nemo.collections.nlp.data.language_modeling.megatron.blendable_dataset.BlendableDataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.GPTDataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.MockGPTDataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.bert_dataset.BertDataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.base_prompt_learning_dataset.BasePromptLearningDataset(*args: Any, **kwargs: Any)#

基类: Dataset

prompt-tuning 或 p-tuning 的基础数据集类。TODO: (@adithyare) 应该合并到 GPTPromptLearningDataset 中

class nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_dataset.GPTSFTDataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_chat_dataset.GPTSFTChatDataset(*args: Any, **kwargs: Any)#

基类: GPTSFTDataset

class nemo.collections.nlp.data.language_modeling.megatron.retro_dataset.RETRODataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.t5_dataset.T5Dataset(*args: Any, **kwargs: Any)#

基类: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.t5_prompt_learning_dataset.T5PromptLearningDataset(*args: Any, **kwargs: Any)#

基类: BasePromptLearningDataset

用于 prompt-tuning 或 p-tuning 预训练 T5 模型的数据集类。

class nemo.collections.nlp.data.language_modeling.megatron.ul2_dataset.UL2Dataset(*args: Any, **kwargs: Any)#

基类: T5Dataset

来自 https://arxiv.org/abs/2205.05131 的 UL2 数据集。包含三个不同的目标

  1. 短跨度掩码,概率较小(例如:T5)。通常最大 ngram 大小为 5,掩码概率为 0.15。

  2. 极端跨度掩码,具有较大概率或较大 ngram 大小或两者兼有。

  3. Prefx-LM,如 T5 或 LM 适配的 T5(prompt-tuning 论文)中所示。

适配器 Mixin 类#

class nemo.collections.nlp.parts.mixins.nlp_adapter_mixins.NLPAdapterModelMixin(*args, **kwargs)#

基类: object

NLP 适配器 Mixin,可以增强任何具有适配器模块支持的基于 Transformer 的模型。此 mixin 类应仅与顶层 ModelPT 子类一起使用,该子类包括 model 或 enc_dec_model 子模块。此 mixin 类添加了几个实用方法来添加、加载和保存适配器。

适配器模块是任何具有一些属性的 Pytorch nn.Module

  • 它的输入和输出维度相同,而隐藏维度不必相同。

  • 适配器模块的最后一层是零初始化的,以便与适配器的残差连接产生原始输出。

此 mixin 类旨在与 PEFT 集成,PEFT 是一个或多个适配器模块。PEFT 的两个功能,层选择和权重绑定,也在此 mixin 类中得到支持。

add_adapter(
peft_cfgs: PEFTConfig | List[PEFTConfig],
)#

用于向模型添加一个或多个适配器模块并冻结基础权重的高级 API。此方法支持从 PEFTConfig 或 PEFTConfig 列表添加适配器模块。它将添加相应的适配器模块。如果 PEFTConfig 中存在层选择和权重绑定,则将应用它们

参数:

peft_cfgs – 指定 PEFT 方法配置的一个或多个 PEFTConfig 对象

load_adapters(
filepath: str,
peft_cfgs: PEFTConfig | List[PEFTConfig] | None = None,
map_location: str | None = None,
)#

仅恢复适配器模块,而不是整个模型本身的实用方法。这允许共享适配器,适配器通常只是完整模型大小的一小部分,从而更容易交付。

注意

在恢复期间,假设模型当前尚未有一个或多个适配器模块。

参数:
  • filepath – .ckpt 或 .nemo 文件的文件路径。

  • peft_cfgs – 指定 PEFT 方法配置的一个或多个 PEFTConfig 对象。如果没有,将从 .nemo 检查点推断

  • map_location – Pytorch 标志,用于放置适配器状态字典。

classmethod merge_cfg_with(
path: str,
cfg: omegaconf.DictConfig,
) omegaconf.DictConfig#

将给定的配置字典 cfg 与从恢复指定路径的 MegatronGPTSFTModel 或 MegatronT5SFTModel 获得的配置字典合并。

参数:
  • path (str) – 要恢复的 SFT 模型检查点的路径。

  • cfg (DictConfig) – 要合并的配置字典。

返回:

合并后的配置字典。

返回类型:

DictConfig

示例

>>> path = "/path/to/model/checkpoint"
>>> cfg = DictConfig({"model": {"key": "value"}, "trainer": {"precision": 16}})
>>> merged_cfg = merge_cfg_with(path, cfg)

注释

  • 该函数使用 OmegaConf.resolve 解析 cfg 字典中的变量。

  • cfg.model 中的键将覆盖输出字典中的相应键。

  • 如果 “train_ds” 存在于 cfg.model.data 中,它会更新 micro_batch_size 和 global_batch_size。

  • 如果 cfg.trainer 包含 “precision” 键,它会更新 output.precision

classmethod merge_inference_cfg(
path: str,
cfg: omegaconf.DictConfig,
) omegaconf.DictConfig#

通过给定的配置字典 cfg,以及从在指定 path 恢复 MegatronGPTSFTModel 或 MegatronT5SFTModel 获取的配置字典,生成一个新的配置字典,并为了推理修改 cfg

参数:
  • path (str) – 要恢复的 SFT 模型检查点的路径。

  • cfg (DictConfig) – 为了推理而修改的配置字典。

返回:

用于推理的配置字典。

返回类型:

DictConfig

示例

>>> path = "/path/to/model/checkpoint"
>>> cfg = DictConfig({"model": {"key": "value"}, "trainer": {"precision": 16}})
>>> merged_cfg = merge_inference_cfg(path, cfg)

注释

  • 来自 cfg 的 “precision” 和 “test_ds” 将覆盖输出字典中对应的键

  • “activations_checkpoint” 将在输出字典中被覆盖为 None

  • 如果其中一个配置字典中为 True,则 “use_flash_attention” 将为 True

  • 如果检查点中不是 None,则 “seq_len_interpolation_factor” 将从 cfg 中被覆盖

可导出模型类#

class nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTExportableModel(*args: Any, **kwargs: Any)#

基类: Module, Exportable

用于 ONNX 导出的 Megatron GPT 包装器