KEMBAR78
[ONNX] Flip `dynamo` default to True in torch.onnx.export · Issue #151693 · pytorch/pytorch · GitHub
Skip to content

[ONNX] Flip dynamo default to True in torch.onnx.export #151693

@titaiwangms

Description

@titaiwangms

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

  1. The dynamo=True path is the recommended approach for ONNX export, as it uses the modern (PyTorch 2.0 ) torch.export.export logic.
  2. Flipping the default will encourage users to adopt the new export logic without requiring explicit configuration.
  3. This change aligns with the deprecation of older options and the push towards modernizing the ONNX export pipeline.

Impact

  1. dynamic_axes will be replaced by dynamic_shapes, as dynamism is now decided in torch.export.export. The main difference is that while dynamic_axes is defined within flatten structure, dynamic_shapes requires 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)

  2. 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

  3. 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

  4. 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 )

  5. Improve and finalize fallback mechanism ([ONNX] Improve and sort out fallback mechanism #151703 )

  6. 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

Sub-issues

Metadata

Metadata

Assignees

No one assigned

    Labels

    module: onnxRelated to torch.onnxoncall: exportoncall: pt2onnx-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

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions