Extender los exportadores

torch.export traza el modelo en un grafo, las etapas posteriores transforman ese grafo y una etapa final lo reduce o lo emite para el tiempo de ejecución objetivo. La mayoría de los modelos pasan sin modificaciones. Cuando hay un patrón de PyTorch que el backend no puede manejar, el exportador aplica un pequeño ajuste en esa etapa en lugar de editar el modelo.

Añada una solución temporal escribiendo una función y registrándola con un decorador. Cada solución temporal debe ubicarse en la etapa más baja que pueda expresarla de manera clara.

Parches y correcciones

Una solución temporal es un parche o una corrección. La diferencia entre ambos radica en si pueden revertirse o no.

Ambos se encuentran en un registro en exporters/utils.py, y el exportador instala todo lo registrado para su backend en la etapa adecuada.

Recurre a un parche cuando el problema sea un error de lowering de un solo backend: una traducción a ONNX que falta, una peculiaridad de validación de ORT o una descomposición de FX que emite una operación muerta. El workaround permanece en el exportador y el código de modelado se mantiene limpio.

Cuando los bloques de patrones se exportan a través de varios backends, como bucles dependientes de los datos, cachés con estado fuera de Cache o atención de bucle dividido escrita a mano, corrige el modelo. Corregirlo una vez en el modelo beneficia a todos los exportadores.

Agregar un parche

Supongamos que un método de un modelo realiza algo que torch.export no puede trazar. El NllbMoeTop2Router._cast_classifier de NLLB-MoE convierte los pesos del clasificador a otro dtype, lo cual no es trazable. Reemplázalo por una operación sin efecto durante la exportación.

Escribe una fábrica que reciba el método original y devuelva su reemplazo, y luego registra la fábrica contra la ruta con puntos del 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

El exportador reemplaza el método antes de realizar el trazado y lo restaura después, por lo que el parche solo afecta a la exportación. Algunas variaciones:

• Pase rutas adicionales para compartir una misma fábrica entre varios puntos de llamada, por ejemplo @register_patch("dynamo", path_a, path_b).
• Parchea una operación de torch apuntando la ruta a ella, por ejemplo @register_patch("onnx", "torch.where"). La fábrica recibe la operación real como argumento, por lo que el reemplazo puede llamarla.
• Escribe una corrección en lugar de un parche cuando necesites reescribir el grafo después del trazado. El mecanismo es el mismo: una función decorada en el registro de coincidencias.

Referencia de etapa

Las etiquetas de origen de cada exportador marcan sus etapas como bloques de comentario # ── Stage N: … ──, de modo que el archivo y esta referencia coincidan. Consulta allí las operaciones y clases exactas que maneja cada etapa.

DynamoExporter

El exportador base ejecuta una etapa de parche y cuatro auxiliares, en orden, dentro de DynamoExporter.export (ver exporter_dynamo.py).

1. Parche de firma hacia adelante: proporciona a model.forward una firma de argumentos plana para que torch.export no agrupe las entradas en una sola tupla **kwargs. Esto es interno y no es un punto de extensión.
2. Parches de modelo: sustituye los métodos de modelo no rastreables por equivalentes seguros para la exportación durante el trazado. Extiende con @register_patch("dynamo", ...).
3. Registro de Pytree: registra cada Cache y ModelOutput para que torch.export pueda aplanarlo y reconstruirlo (normalmente ocurre de forma automática). Añade una rama a _flatten_to_context / _unflatten_from_context para un tipo al que la inspección de atributos no puede llegar.
4. Formas dinámicas: asigne Dim.AUTO a cada tensor y caché de hoja cuando dynamic=True. Anule con DynamoConfig.dynamic_shapes.
5. Limpieza de estado: restablece los atributos de tensor que un modelo establece dentro de forward y que torch.export deja como tensores falsos. Amplía esto añadiendo el nombre del atributo a _STATEFUL_CACHE_ATTRS.

OnnxExporter

OnnxExporter añade cinco etapas alrededor de torch.onnx.export (ver exporter_onnx.py). Busca en el archivo la lista completa de parches:

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

1. Parches de Torch: sustituye las operaciones de torch que el exportador de ONNX no puede traducir tal cual. Extiéndelo con @register_patch("onnx", ...).
2. Parches de ONNX: vuelve a ejecutar las correcciones de nodos después de run_decompositions para que los nodos de protección de forma recién introducidos se reparen antes del lowering. Utiliza el mismo registro @register_patch("onnx", ...).
3. Correcciones de nodos FX: reescribe los nodos del grafo que el exportador de ONNX no puede reducir, como operaciones de alias, vistas in-place y asserts muertos. Extiéndelo con @register_fx_node_fix("onnx").
4. Traducciones de ONNX: proporciona un lowering personalizado para un op de aten cuando el predeterminado falta o tiene un error (por ejemplo, aten.index_put o aten._grouped_mm). Añade una función _aten_* a _ONNX_TRANSLATION_TABLE.
5. Correcciones de IR de ONNX: reescribe el programa ONNX terminado para evitar errores de ONNX Runtime (por ejemplo, forzando TopK(sorted=True)). Añade una función _fix_ir_* a _IR_FIXES.

ExecutorchExporter

ExecutorchExporter añade cinco etapas alrededor de to_edge_transform_and_lower y to_executorch, comenzando con la preparación del backend (ver exporter_executorch.py).

1. Preparación del backend: mueve el modelo al dispositivo y dtype de destino y elige su particionador (prepare_for_xnnpack, prepare_for_cuda). Añade un backend registrando prepare_for_<name> en _BACKEND_PREPARE.
2. Parches de Torch: reemplaza las operaciones de torch que los backends de ExecuTorch no pueden aceptar, como split_copy, chunk y topk(k>dim). Extiéndelo con @register_patch("executorch", ...).
3. Parches de ExecuTorch: reemplaza los componentes internos de ExecuTorch que provocan errores en gráficos de forma dinámica válidos. Utiliza el mismo registro @register_patch("executorch", ...).
4. Correcciones del programa FX: repara el programa exportado cuando la corrección requiere el contexto de todo el programa, como ampliar las restricciones de rango o completar metadatos de marcador de posición faltantes. Extiéndelo con @register_fx_program_fix("executorch").
5. Correcciones de nodos FX: reescritura de nodos individuales, como el mapeo de operaciones simbólicas de Python a executorch_prim.* o la reescritura de pow como una cadena de mul. Extiéndase con @register_fx_node_fix("executorch").

Soluciones conocidas aguas arriba

Algunas clases de modelos presentan errores confirmados en el optimizador de grafos de onnxscript (el plegado de constantes se interrumpe en SplitToSequence, y se descartan los inicializadores de FPN). ONNX_DISABLE_OPTIMIZE desactiva la optimización de onnxscript para esos modelos. Cada entrada registra el problema de origen junto al nombre del modelo. Se espera que la lista se reduzca a medida que se resuelvan los errores de origen, por lo que una nueva entrada debe hacer referencia a un error específico de origen en lugar de desactivar la optimización de forma arbitraria.

EXPORT_SKIPS, excluye un pequeño número de clases de modelos del barrido de exportación por completo cuando el modelo es fundamentalmente no exportable tal como está (flujo de control dependiente de datos que no se puede vectorizar, o módulos tratados como argumentos de forward). Cada entrada incluye una razón que nombra el cambio necesario en el lado del modelo. Se espera que esta lista se reduzca, no que crezca.