-
Notifications
You must be signed in to change notification settings - Fork 25.7k
Description
NOTE
The task will not commence until the sub-issues outlined below are addressed or resolved.
Description
The torch.onnx.export function currently defaults the dynamo parameter to False. However, the dynamo=True path is the recommended approach for ONNX export, as it leverages the modern torch.export.export logic introduced in PyTorch 2.0. This task proposes flipping the default value of dynamo to True to align with the latest export pipeline and encourage users to adopt the updated functionality.
Before implementing this change, several sub-issues need to be addressed to ensure a smooth transition for users.
Proposed Change
Update torch.onnx.export to set dynamo=True as default value:
def export(
model: torch.nn.Module
| torch.export.ExportedProgram
| torch.jit.ScriptModule
| torch.jit.ScriptFunction,
args: tuple[Any, ...] = (),
f: str | os.PathLike | None = None,
*,
kwargs: dict[str, Any] | None = None,
export_params: bool = True,
verbose: bool | None = None,
input_names: Sequence[str] | None = None,
output_names: Sequence[str] | None = None,
opset_version: int | None = None,
dynamic_axes: Mapping[str, Mapping[int, str]]
| Mapping[str, Sequence[int]]
| None = None,
keep_initializers_as_inputs: bool = False,
dynamo: bool = True, # Flip default to True
...
) -> ONNXProgram | None:Rationale
- The dynamo=True path is the recommended approach for ONNX export, as it uses the modern (PyTorch 2.0 ) torch.export.export logic.
- Flipping the default will encourage users to adopt the new export logic without requiring explicit configuration.
- This change aligns with the deprecation of older options and the push towards modernizing the ONNX export pipeline.
Impact
-
dynamic_axeswill be replaced bydynamic_shapes, as dynamism is now decided intorch.export.export. The main difference is that whiledynamic_axesis defined within flatten structure,dynamic_shapesrequires the nested tree structure as input argument. Although the auto-conversion will be provided, there are corner cases. More detail can be found: [ONNX] Improve dynamic_axes to dynamic_shapes conversion in exporter #150940 (TODO: Summarize the pifalls) -
Loops in the models will need to be re-written with torch.scan. [ONNX] Implement scan #151327 and run_decompositions fails with torch.scan #151564
-
Control flows (if/else) in the models will need to be re-written with torch.cond. For example, torch.cond + torch.non_zero does not work with torch.export.export #144691, https://pytorch.org/tutorials/beginner/onnx/export_control_flow_model_to_onnx_tutorial.html
-
torch.jit.ScriptModule and torch.jit.ScriptFunction will automatically fallback to depreacted torchscript-based exporter ([ONNX] Auto-fallback in torch.onnx.export(..., dynamo=True) with deprecation warning #151714 )
-
Improve and finalize fallback mechanism ([ONNX] Improve and sort out fallback mechanism #151703 )
-
Migrate torchscript-based ONNX tests to the new exporter ([ONNX] Create unit tests for the new export path by adapting all existing tests #129279)
cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @xadupre @justinchuby @shubhambhokare1 @gramalingam