-
Notifications
You must be signed in to change notification settings - Fork 25.7k
Description
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:
pytorch/torch/onnx/_internal/exporter/_core.py
Lines 1263 to 1271 in 1b267a5
| # 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