添加对新架构的支持

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

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

要在 Optimum Neuron 库中支持新的模型架构，需要遵循以下步骤：

1. 实现自定义的 Neuron 配置。
2. 导出并验证模型。
3. 向 GitHub 仓库贡献代码。

实现自定义的 Neuron 配置

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

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

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

由于许多架构的 Neuron 配置具有相似的属性，🤗 Optimum 采用了一个三级类层次结构：

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

示例：为 ESM 模型添加支持

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

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

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

@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` 的推理测试中。
• 提交一个拉取请求！（别忘了将它与你创建的 issue 关联起来，这样维护者就能更好地跟踪并及时提供帮助。）

你已经为社区在 Neuron 上实现了一个新模型的可用性！感谢你加入我们，共同推动优秀机器学习的民主化进程 🤗。