添加对新架构的支持

注意： ❗ 此部分不适用于通过 transformers-neuronx 集成的自回归采样对解码器模型的推理。如果您想添加对这些模型的支持，请在 Optimum Neuron GitHub 存储库中打开问题，并联系维护者寻求帮助。

您想将新模型导出到 AWS Inferentia 或 Trainium 上运行吗？查看指南，并向 🤗 Optimum Neuron 的 GitHub 存储库提交拉取请求！

要在 Optimum Neuron 库中支持新的模型架构，请按照以下步骤操作

1. 实现自定义 Neuron 配置。
2. 导出并验证模型。
3. 贡献到 GitHub 存储库。

实现自定义 Neuron 配置

为了支持将新的模型导出到 Neuron 兼容格式，首先要定义一个 Neuron 配置，它描述如何通过指定以下内容来导出 PyTorch 模型

1. 输入名称。
2. 输出名称。
3. 用于跟踪模型的虚拟输入：Neuron 编译器通过跟踪记录计算图，并在生成的 TorchScript 模块上运行。
4. 用于控制硬件效率（延迟、吞吐量）和精度之间权衡的编译参数。

根据模型和任务的选择，我们使用配置类来表示上述数据。每个配置类都与特定的模型架构相关联，并遵循命名约定 ArchitectureNameNeuronConfig。例如，指定 BERT 模型的 Neuron 导出的配置是 BertNeuronConfig。

由于许多架构在它们的 Neuron 配置方面具有类似的属性，因此 🤗 Optimum 采用了一个 3 级类层次结构

1. 抽象和通用的基类。这些类处理所有基本功能，但与模态（文本、图像、音频等）无关。
2. 中端类。这些类了解模态。根据它们支持的输入，可以为同一个模态存在多个配置类。它们指定应使用哪些输入生成器来生成虚拟输入，但仍然与模型无关。
3. 模型特定的类，例如上面提到的 BertNeuronConfig。这些类实际上用于导出模型。

示例：添加对 ESM 模型的支持

这里我们以支持 ESM 模型 为例。让我们在 optimum/exporters/neuron/model_configs.py 中创建一个 EsmNeuronConfig 类。

当 Esm 模型被解释为文本编码器时，我们可以从中端类 TextEncoderNeuronConfig 继承。由于 Esm 的建模和配置在被解释为编码器时与 BERT 几乎相同，因此我们可以使用 model_type=bert 的 NormalizedConfigManager 来规范化配置，从而生成虚拟输入来跟踪模型。

最后一步，由于 optimum-neuron 是 optimum 的扩展，我们需要将我们创建的 Neuron 配置使用 register_in_tasks_manager 装饰器注册到 TasksManager 中，并指定模型类型和支持的任务。

@register_in_tasks_manager("esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"])
class EsmNeuronConfig(TextEncoderNeuronConfig):
    NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("bert")
    ATOL_FOR_VALIDATION = 1e-3  # absolute tolerance to compare for comparing model on CPUs

    @property
    def inputs(self) -> List[str]:
        return ["input_ids", "attention_mask"]

导出并验证模型

使用您实现的 Neuron 配置类，现在快速测试它是否按预期工作。

• 导出

optimum-cli export neuron --model facebook/esm2_t33_650M_UR50D --task text-classification --batch_size 1 --sequence_length 16 esm_neuron/

在导出过程中，将调用 validate_model_outputs 以通过将导出的 Neuron 模型的输出与 PyTorch 在 CPU 上的结果进行比较来验证输出。您也可以手动验证模型，方法是：

from optimum.exporters.neuron import validate_model_outputs

validate_model_outputs(
    neuron_config, base_model, neuron_model_path, neuron_named_outputs, neuron_config.ATOL_FOR_VALIDATION
)

• 推理（可选）

from transformers import AutoTokenizer
from optimum.neuron import NeuronModelForSequenceClassification

model = NeuronModelForSequenceClassification.from_pretrained("esm_neuron/")
tokenizer = AutoTokenizer.from_pretrained("esm_neuron/")
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
logits = model(**inputs).logits

贡献到 GitHub 仓库

我们几乎都准备好了。现在提交一个拉取请求，让您的工作对所有社区成员可见！

• 在 Optimum Neuron GitHub 仓库 中打开一个问题，描述新功能，并使其对 Optimum Neuron 的维护者可见。
• 将模型添加到 optimum-neuron/tests/exporters/exporters_utils.py 中的导出器测试和 optimum-neuron/tests/inference/inference_utils.py 中的推理测试。
• 打开一个拉取请求！ (不要忘记将它链接到您打开的问题，这样维护者可以更好地跟踪它并在需要时提供帮助。)

您已使社区在 Neuron 上可以使用新模型！感谢您加入我们，共同努力使机器学习更加民主化 🤗。