KEMBAR78
[ONNX] Improve and sort out fallback mechanism · Issue #151703 · pytorch/pytorch · GitHub
Skip to content

[ONNX] Improve and sort out fallback mechanism #151703

@titaiwangms

Description

@titaiwangms

In torch.onnx.export(..., dynamo=True), when fallback=True, the current behavior involves attempting four different strategies within the FX-based exporter (exported program-based exporter) before falling back to the TorchScript-based exporter. These strategies include TorchExportNonStrictStrategy, TorchExportStrategy, TorchExportDraftExportStrategy, and JitTraceConvertStrategy:

# Record the status
if strategy_class is _capture_strategies.TorchExportNonStrictStrategy:
export_status.torch_export_non_strict = result.success
elif strategy_class is _capture_strategies.TorchExportStrategy:
export_status.torch_export = result.success
elif strategy_class is _capture_strategies.TorchExportDraftExportStrategy:
export_status.torch_export_draft_export = result.success
elif strategy_class is _capture_strategies.JitTraceConvertStrategy:
export_status.torch_jit = result.success

The original intent was to maximize export coverage by trying all known strategies. However, the additional strategies (TorchExportStrategy, TorchExportDraftExportStrategy, and JitTraceConvertStrategy) do not significantly improve coverage beyond what TorchExportNonStrictStrategy already handles. Furthermore, the complex try-catch error messages make it harder for users to understand the root cause of failures.

To simplify the process and reduce confusion, especially in light of upcoming changes #151693, torch.onnx.export(..., dynamo=True) should directly use torch.export.export(..., strict=False) (TorchExportNonStrictStrategy), as the primary strategy. If this fails, it should immediately fall back to torch.onnx.export(..., dynamo=False).

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

Metadata

Metadata

Assignees

Labels

module: onnxRelated to torch.onnxonnx-triagedtriaged by ONNX teamtriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions