Improper model conversion from PyTorch to ONNX with torch.onnx.OperatorExportTypes.ONNX_ATEN flag

[oviazlo]

🐛 Describe the bug

I have spotted improper model conversion to ONNX when using torch.onnx.OperatorExportTypes.ONNX_ATEN flag with PyTorch >=1.12.0

My understanding from the documentation

OperatorExportTypes.ONNX_ATEN: All ATen ops (in the TorchScript namespace “aten”) are exported as ATen ops (in opset domain “org.pytorch.aten”). ATen is PyTorch’s built-in tensor library, so this instructs the runtime to use PyTorch’s implementation of these ops.

is that when one uses torch.onnx.OperatorExportTypes.ONNX_ATEN flag, all ATen operators has to be exported to ONNX graph. When I run the code snippet below:

import torch


class ModelWithAtenFmod(torch.nn.Module):
    def forward(self, x, y):
        return torch.fmod(x, y)


x = torch.randn(3, 4, dtype=torch.float32)
y = torch.randn(3, 4, dtype=torch.float32)
OUT = torch.onnx.export_to_pretty_string(
    ModelWithAtenFmod(), (x, y),
    add_node_names=False,
    do_constant_folding=False,
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN)
print(OUT)

with PyTorch 1.11.0 I obtain ONNX graph with ATen operator that is correct:

### PyTorch 1.11.0

ModelProto {
  producer_name: "pytorch"
  domain: ""
  doc_string: ""
  graph:
    GraphProto {
      name: "torch-jit-export"
      inputs: [{name: "aten::ATen_0", type:Tensor dtype: 1, Tensor dims: 3 4},{name: "aten::ATen_1", type:Tensor dtype: 1, Tensor dims: 3 4}]
      outputs: [{name: "2", type:Tensor dtype: 1, Tensor dims: ? ?}]
      value_infos: []
      initializers: []
      nodes: [
        Node {type: "ATen", inputs: [aten::ATen_0,aten::ATen_1], outputs: [2], attributes: [{ name: 'operator', type: string, value: 'fmod'}]}
      ]
    }
  opset_import: [OperatorSetIdProto { domain: , version: 9}OperatorSetIdProto { domain: org.pytorch.aten, version: 1}],
}

However, when I run the same code with PyTorch 1.12.1, ATen operator got substituted with ONNX native operator and that is not expected behaviour:

### PyTorch 1.12.1

ModelProto {
  producer_name: "pytorch"
  domain: ""
  doc_string: ""
  graph:
    GraphProto {
      name: "torch_jit"
      inputs: [{name: "onnx::Mod_0", type:Tensor dtype: 1, Tensor dims: 3 4},{name: "onnx::Mod_1", type:Tensor dtype: 1, Tensor dims: 3 4}]
      outputs: [{name: "2", type:Tensor dtype: 1, Tensor dims: 3 4}]
      value_infos: []
      initializers: []
      nodes: [
        Node {type: "Mod", inputs: [onnx::Mod_0,onnx::Mod_1], outputs: [2], attributes: [{ name: 'fmod', type: int, value: 1}]}
      ]
    }
  opset_import: [OperatorSetIdProto { domain: , version: 13}],
}

Versions

PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.6 (arm64)
GCC version: Could not collect
Clang version: 14.0.6
CMake version: version 3.24.1
Libc version: N/A

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:14) [Clang 12.0.1 ] (64-bit runtime)
Python platform: macOS-12.6-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.4
[pip3] torch==1.12.1
[pip3] torchvision==0.13.1
[conda] numpy 1.23.4 py38h09ac2d9_0 conda-forge
[conda] pytorch 1.12.1 py3.8_0 pytorch
[conda] torchvision 0.13.1 py38_cpu pytorch

[oviazlo]

I got the example of the code snippet from one existing unit test:

pytorch/test/jit/test_export_modes.py

Lines 107 to 121 in 664058f

	# torch.fmod is using to test ONNX_ATEN.
	# If you plan to remove fmod from aten, or found this test failed.
	# please contact @Rui.
	def test_onnx_aten(self):
	class ModelWithAtenFmod(nn.Module):
	def forward(self, x, y):
	return torch.fmod(x, y)
	x = torch.randn(3, 4, dtype=torch.float32)
	y = torch.randn(3, 4, dtype=torch.float32)
	torch.onnx.export_to_pretty_string(
	ModelWithAtenFmod(), (x, y),
	add_node_names=False,
	do_constant_folding=False,
	operator_export_type=OperatorExportTypes.ONNX_ATEN)

[oviazlo]

Operator substitution is happening inside the _optimize_graph function in the code below:

pytorch/torch/onnx/utils.py

Lines 726 to 743 in 664058f

	model = _pre_trace_quant_model(model, args)
	graph, params, torch_out, module = _create_jit_graph(model, args)
	params_dict = _get_named_param_dict(graph, params)
	try:
	graph = _optimize_graph(
	graph,
	operator_export_type,
	_disable_torch_constant_prop=_disable_torch_constant_prop,
	fixed_batch_size=fixed_batch_size,
	params_dict=params_dict,
	dynamic_axes=dynamic_axes,
	input_names=input_names,
	module=module,
	)
	except Exception as e:
	torch.onnx.log("Torch IR graph at exception: ", graph)
	raise

The graph before _optimize_graph function call looks like:

graph(%0 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %1 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %6 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = aten::fmod(%0, %1)
  return (%6)

And after the call:

graph(%0 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %1 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %2 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = onnx::Mod[fmod=1](%0, %1)
  return (%2)

despite operator_export_type being equal to OperatorExportTypes.ONNX_ATEN

[thiagocrepaldi]

This is happening inside graph = _C._jit_pass_onnx(graph, operator_export_type)

[thiagocrepaldi]

ca37477 - [ONNX] update default opset_version to 13 (#73898) introduced this regression. Looking into it further

[thiagocrepaldi]

Fix: #87735