添加对新架构的支持

注意： ❗本节不适用于通过 transformers-neuronx 集成的自回归采样解码器模型的推理。如果您想添加对这些模型的支持，请在 Optimum Neuron GitHub 仓库上开启一个 issue，并 ping 维护人员寻求帮助。

您想在 AWS Inferentia 或 Trainium 上导出并运行新模型吗？请查看指南，并向 🤗 Optimum Neuron 的 GitHub 仓库提交 pull request！

为了在 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 作为编码器解释时几乎相同，我们可以使用 NormalizedConfigManager 和 model_type=bert 来标准化配置，以生成用于追踪模型的虚拟输入。

最后一步，由于 optimum-neuron 是 optimum 的扩展，我们需要使用 register_in_tasks_manager 装饰器将我们创建的 Neuron 配置注册到 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 模型的输出，方法是将它们与 CPU 上 PyTorch 的结果进行比较。您也可以使用以下命令手动验证模型：

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 仓库

我们几乎都准备好了。现在提交一个 pull request，让您的工作可以被所有社区成员访问！

• 在 Optimum Neuron GitHub 仓库 中开启一个 issue，描述新功能并使其对 Optimum Neuron 的维护人员可见。
• 在 optimum-neuron/tests/exporters/exporters_utils.py 中将模型添加到导出器测试，并在 optimum-neuron/tests/inference/inference_utils.py 中添加到推理测试。
• 打开一个 pull request! （不要忘记将其链接到您打开的 issue，以便维护人员可以更好地跟踪它并在需要时提供帮助。）

您已经为社区提供了一个可以在 Neuron 上使用的新模型！感谢您加入我们 democratizing good machine learning 🤗 的努力。