扩展导出器

torch.export 将模型追踪为图，后续阶段会转换该图，最后一个阶段则将其降低或输出到目标运行时。大多数模型会原样通过。当存在后端无法处理的 PyTorch 模式时，导出器会在该阶段应用一个小变通方案，而不是修改模型。

通过编写一个函数并使用装饰器进行注册来添加变通方案。每个变通方案都应放置在能够清晰表达它的最低阶段。

补丁与修复

解决方法可以是补丁或修复。两者的区别在于是否可以回退。

两者都位于 exporters/utils.py 中的注册表中，exporter 会在合适的阶段安装其后端所注册的所有内容。

当问题源于某个单一后端的降级（lowering）缺陷时，应使用补丁（patch）：例如缺少 ONNX 转换、ORT 验证的怪癖，或 FX 分解生成了无效算子。此时，变通方案保留在导出器中，模型代码保持整洁。

当模式块在多个后端之间导出时，例如数据依赖循环、Cache 之外的有状态缓存，或手工编写的拆分循环注意力，应直接修复模型。在模型中一次性修复有助于所有导出器。

添加补丁

假设某个模型方法执行了 torch.export 无法追踪的操作。NLLB-MoE 的 NllbMoeTop2Router._cast_classifier 会将分类器权重转换为另一种 dtype，而这一操作无法被追踪。在导出期间，请将其替换为空操作（no-op）。

编写一个工厂函数，该函数接收原始方法并返回其替代方法，然后将该工厂函数注册到方法的点分路径上：

from transformers.exporters.utils import register_patch

@register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier")
def _patch_classifier_cast(_original):
    # Replace the untraceable dtype cast with a no-op during export.
    return lambda self, *args, **kwargs: None

导出器在追踪之前替换该方法，并在之后恢复，因此该补丁仅影响导出过程。以下是几种变体：

• 通过传入额外的路径，可以在多个调用点之间共享同一个工厂，例如 @register_patch("dynamo", path_a, path_b)。
• 通过指定路径来修补 torch 算子，例如 @register_patch("onnx", "torch.where")。 工厂函数会接收真实的算子作为参数，因此替换实现可以调用该算子。
• 当需要在追踪后重写计算图时，请编写一个修复（fix）而非补丁（patch）。其机制相同，即在对应的注册表中使用一个带装饰器的函数。

阶段参考

每个导出器的源代码使用 # ── Stage N: … ── 注释块来标记其各个阶段，因此该文件与本参考文档保持一致。请查阅源代码，以了解每个阶段具体处理的操作和类。

DynamoExporter

基础导出器按顺序在 DynamoExporter.export 中运行一个 patch 阶段和四个辅助函数（参见 exporter_dynamo.py）。

1. 前向签名补丁：为 model.forward 提供扁平化的参数签名，以便 torch.export 不会将输入打包成一个 **kwargs 元组。这是内部实现，并非扩展点。
2. 模型补丁：在跟踪期间，将不可跟踪的模型方法替换为可安全导出的等效方法。使用 @register_patch("dynamo", ...) 进行扩展。
3. Pytree 注册：为每个 Cache 和 ModelOutput 进行注册，以便 torch.export 能够对其进行扁平化和重建（通常会自动完成）。对于属性遍历无法触及的类型，在 _flatten_to_context / _unflatten_from_context 中添加一个分支。
4. 动态形状：当 dynamic=True 时，为每个张量和缓存叶节点分配 Dim.AUTO。可通过 DynamoConfig.dynamic_shapes 进行覆盖。
5. 状态清理：重置模型在 forward 中设置、但被 torch.export 保留为伪张量的张量属性。通过向 _STATEFUL_CACHE_ATTRS 添加属性名来扩展。

OnnxExporter

OnnxExporter 在 torch.onnx.export 周围添加了五个阶段（参见 exporter_onnx.py）。 在文件中搜索以获取完整的补丁列表：

grep -nE "^def (_patch_|_fix_|_aten_)" src/transformers/exporters/exporter_onnx.py

1. Torch 补丁：替换 ONNX 导出器无法直接转换的 torch 操作。可通过 @register_patch("onnx", ...) 进行扩展。
2. ONNX 补丁：在 run_decompositions 之后重新运行节点修复，以便在降低之前修复新引入的形状保护节点。使用相同的 @register_patch("onnx", ...) 注册表。
3. FX 节点修复：重写 ONNX 导出器无法降低的图节点，例如别名操作、原地视图和无效断言。可通过 @register_fx_node_fix("onnx") 进行扩展。
4. ONNX 转换：为默认实现缺失或存在缺陷的 aten 算子提供自定义的 lowering（例如 aten.index_put 或 aten._grouped_mm）。向 _ONNX_TRANSLATION_TABLE 中添加一个 _aten_* 函数。
5. ONNX IR 修复：重写已完成的 ONNX 程序以规避 ONNX Runtime 的 bug（例如强制使用 TopK(sorted=True)）。向 _IR_FIXES 中添加一个 _fix_ir_* 函数。

ExecutorchExporter

ExecutorchExporter 在 to_edge_transform_and_lower 和 to_executorch 周围增加了五个阶段，从后端准备开始（参见 exporter_executorch.py）。

1. 后端准备：将模型移动到目标设备和 dtype，并选择其分区器（prepare_for_xnnpack、prepare_for_cuda）。通过在 _BACKEND_PREPARE 中注册 prepare_for_<name> 来添加后端。
2. Torch 补丁：替换 ExecuTorch 后端无法接受的 torch 操作，例如 split_copy、chunk 和 topk(k>dim)。使用 @register_patch("executorch", ...) 进行扩展。
3. ExecuTorch 补丁：替换在有效的动态形状图上会崩溃的 ExecuTorch 内部实现。使用相同的 @register_patch("executorch", ...) 注册表。
4. FX 程序修复：在需要整个程序上下文的情况下修复导出的程序，例如放宽范围约束或填充缺失的占位符元数据。使用 @register_fx_program_fix("executorch") 进行扩展。
5. FX 节点修复：重写单个节点，例如将 Python sym ops 映射到 executorch_prim.*，或将 pow 重写为 mul 链。使用 @register_fx_node_fix("executorch") 进行扩展。

已知的上游解决方案

部分模型类在 onnxscript 图优化器中触发了已确认的 bug（例如 SplitToSequence 上的常量折叠崩溃，以及 FPN 初始化器被丢弃）。ONNX_DISABLE_OPTIMIZE 会为这些模型禁用 onnxscript 优化。每个条目都会在模型名称旁记录对应的上游 issue。随着上游 bug 的修复，该列表预计会逐渐缩短，因此新增条目必须引用具体的上游 bug，而不能随意禁用优化。

EXPORT_SKIPS， 当模型本身无法直接导出时（例如存在无法向量化处理的数据依赖控制流，或某些模块被当作前向参数处理），该选项会将少数模型类完全排除在导出扫描之外。每个条目都附带一个原因，指明模型侧需要做出的相应修改。预计该列表会逐渐缩减，而非增加。