Estendendo os exportadores

torch.export rastreia o modelo em um grafo, as etapas subsequentes transformam esse grafo, e uma etapa final o reduz ou o emite para o runtime de destino. A maioria dos modelos passa por isso sem alterações. Quando há um padrão do PyTorch que o backend não consegue processar, o exportador aplica um pequeno ajuste naquela etapa, em vez de editar o modelo.

Adicione uma solução alternativa escrevendo uma função e registrando-a com um decorador. Cada solução alternativa deve estar no estágio mais baixo que possa expressá-la de forma limpa.

Correções e reparos

Uma solução alternativa é um patch ou um fix. A diferença entre os dois está em se podem ser revertidos.

Ambos estão em um registro em exporters/utils.py, e o exportador instala tudo que está registrado para o backend correspondente no estágio adequado.

Recorra a um patch quando o problema for um bug de lowering de um único backend: uma tradução ONNX ausente, uma peculiaridade de validação do ORT ou uma decomposição FX que emite uma operação morta. A solução alternativa permanece no exportador, e o código de modelagem permanece limpo.

Quando os blocos de padrão são exportados para vários backends, como loops dependentes de dados, caches com estado fora de Cache ou atenção com loop dividido escrito manualmente, corrija o modelo. Corrigi-lo uma única vez no modelo beneficia todos os exportadores.

Adicionar um patch

Suponha que um método de modelo faça algo que o torch.export não consiga rastrear. O NllbMoeTop2Router._cast_classifier do NLLB-MoE converte os pesos do classificador para outro dtype, o que não é rastreável. Substitua-o por uma operação nula durante a exportação.

Escreva uma fábrica que recebe o método original e retorna seu substituto, em seguida, registre a fábrica no caminho pontilhado do método:

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

O exportador substitui o método antes do rastreamento e o restaura em seguida, de modo que o patch afeta apenas a exportação. Algumas variações:

• Passar caminhos extras para compartilhar uma única fábrica entre diferentes pontos de chamada, por exemplo @register_patch("dynamo", path_a, path_b).
• Faça o patch de uma op do torch apontando o caminho para ela, por exemplo @register_patch("onnx", "torch.where"). A fábrica recebe a op real como argumento, de modo que a substituição pode chamá-la.
• Escreva uma correção em vez de um patch quando precisar reescrever o grafo após o tracing. O mecanismo é o mesmo: uma função decorada no registro correspondente.

Referência de estágio

O código-fonte de cada exportador rotula seus estágios como blocos de comentário # ── Stage N: … ──, de modo que o arquivo e esta referência fiquem alinhados. Consulte-o para ver as operações e classes exatas que cada estágio processa.

DynamoExporter

O exportador base executa um estágio de patch e quatro helpers, em ordem, dentro de DynamoExporter.export (veja exporter_dynamo.py).

1. Patch de assinatura de forward: fornece a model.forward uma assinatura de argumentos plana para que o torch.export não agrupe as entradas em uma única tupla **kwargs. Isso é interno e não é um ponto de extensão.
2. Patches de modelo: substitui métodos de modelo não rastreáveis por equivalentes seguros para exportação durante o tracing. Estenda com @register_patch("dynamo", ...).
3. Registro de Pytree: registre cada Cache e ModelOutput para que o torch.export possa achatar e reconstruí-lo (geralmente acontece automaticamente). Adicione um ramo em _flatten_to_context / _unflatten_from_context para um tipo que a varredura de atributos não consegue alcançar.
4. Formas dinâmicas: atribua Dim.AUTO a cada tensor e cache leaf quando dynamic=True. Substitua com DynamoConfig.dynamic_shapes.
5. Limpeza de estado: redefina os atributos de tensor que um modelo define dentro de forward e que o torch.export deixa como tensores falsos. Estenda adicionando o nome do atributo a _STATEFUL_CACHE_ATTRS.

OnnxExporter

OnnxExporter adiciona cinco etapas ao redor de torch.onnx.export (veja exporter_onnx.py). Use o Grep no arquivo para obter a lista completa de patches:

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

1. Patches do Torch: substitui operações do torch que o exportador ONNX não consegue traduzir diretamente. Estenda com @register_patch("onnx", ...).
2. Correções de ONNX: reexecute as correções de nós após run_decompositions para que os nós de verificação de forma recém-introduzidos sejam reparados antes da redução. Utiliza o mesmo registro @register_patch("onnx", ...).
3. Correções de nós FX: reescreve nós do grafo que o exportador ONNX não consegue reduzir, como operações de alias, visões in-place e asserts mortos. Estenda com @register_fx_node_fix("onnx").
4. Traduções ONNX: forneça um lowering personalizado para uma op aten quando o padrão estiver ausente ou com bug (por exemplo, aten.index_put ou aten._grouped_mm). Adicione uma função _aten_* ao _ONNX_TRANSLATION_TABLE.
5. Correções no IR do ONNX: reescreva o programa ONNX finalizado para contornar bugs do ONNX Runtime (por exemplo, forçando TopK(sorted=True)). Adicione uma função _fix_ir_* a _IR_FIXES.

ExecutorchExporter

ExecutorchExporter adiciona cinco etapas ao redor de to_edge_transform_and_lower e to_executorch, começando pela preparação do backend (veja exporter_executorch.py).

1. Preparação do backend: mova o modelo para o dispositivo e o dtype de destino e escolha seu particionador (prepare_for_xnnpack, prepare_for_cuda). Adicione um backend registrando prepare_for_<name> em _BACKEND_PREPARE.
2. Patches do Torch: substitui operações do torch que os backends do ExecuTorch não conseguem aceitar, como split_copy, chunk e topk(k>dim). Estenda com @register_patch("executorch", ...).
3. Patches do ExecuTorch: substitui os componentes internos do ExecuTorch que causam falhas em grafos de forma dinâmica válidos. Utiliza o mesmo registro @register_patch("executorch", ...).
4. Correções do programa FX: repare o programa exportado quando a correção requer contexto de programa inteiro, como alargar restrições de intervalo ou preencher metadados de placeholder ausentes. Estenda com @register_fx_program_fix("executorch").
5. Correções de nós FX: reescreva nós individuais, como mapear operações simbólicas do Python para executorch_prim.* ou reescrever pow como uma cadeia de mul. Estenda com @register_fx_node_fix("executorch").

Soluções conhecidas em montante

Algumas classes de modelo atingiram bugs confirmados no otimizador de grafo do onnxscript (a dobra de constantes causa falha em SplitToSequence, e os inicializadores do FPN são descartados). ONNX_DISABLE_OPTIMIZE desativa a otimização do onnxscript para esses modelos. Cada entrada registra a questão upstream ao lado do nome do modelo. Espera-se que a lista diminua conforme os bugs upstream forem corrigidos, portanto, uma nova entrada deve referenciar um bug upstream específico, em vez de desativar a otimização arbitrariamente.

EXPORT_SKIPS, exclui um pequeno número de classes de modelo do processo de exportação quando o modelo é fundamentalmente não exportável em seu estado atual (fluxo de controle dependente de dados que não pode ser vetorizado, ou módulos tratados como argumentos de forward). Cada entrada inclui um motivo que especifica a alteração necessária no lado do modelo. Espera-se que essa lista diminua, não aumente.