[ONNX] Float8 support
Labels: module: onnx, low priority, triaged

### 🚀 The feature, motivation and pitch

Support exporting float8 models to ONNX.

### Alternatives

_No response_

### Additional context

Exporting the model

```python
import torch


class Float8Module(torch.nn.Module):
    def forward(self, input: torch.Tensor):
        input = input.to(torch.float8_e5m2)
        return input + torch.tensor(1.0, dtype=torch.float8_e5m2)

torch.onnx.dynamo_export(Float8Module(), torch.randn(1, 3, 224, 224))
```

Raises

```
/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py:112: UserWarning: Torchlib only supports opset version 18 for now. If you need to use a             different opset version, please register them with register_custom_op.
  warnings.warn(
Traceback (most recent call last):
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py", line 1006, in dynamo_export
    ).export()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py", line 810, in export
    graph_module = self.options.fx_tracer.generate_fx(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py", line 198, in generate_fx
    graph_module, graph_guard = torch._dynamo.export(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1014, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 321, in _fn
    return fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py", line 153, in wrapped
    return output_adapter.apply(model_func(*args, **kwargs))
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 481, in catch_errors
    return callback(frame, cache_size, hooks, frame_state)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 129, in _fn
    return fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 361, in _convert_frame_assert
    return _compile(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 194, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 431, in _compile
    out_code = transform_code_object(code, transform)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 416, in transform
    tracer.run()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2071, in run
    super().run()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 168, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 561, in call_function
    return wrap_fx_proxy(tx, proxy, **options)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1156, in wrap_fx_proxy
    return wrap_fx_proxy_cls(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1230, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1360, in get_fake_value
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1328, in get_fake_value
    return wrap_fake_exception(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 907, in wrap_fake_exception
    return fn()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1329, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1394, in run_node
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1381, in run_node
    return node.target(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1235, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1525, in dispatch
    r = func(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_ops.py", line 435, in __call__
    return self._op(*args, **kwargs or {})
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 229, in _fn
    result = fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 120, in _fn
    compute_dtype, result_dtype = utils.elementwise_dtypes(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/__init__.py", line 1373, in elementwise_dtypes
    highest_type = get_higher_type(highest_type, dtype_to_type(x.dtype))
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/__init__.py", line 915, in dtype_to_type
    raise ValueError("Invalid dtype!")
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in function add>(*(FakeTensor(..., size=(1, 3, 224, 224), dtype=torch.float8_e5m2), FakeTensor(..., size=(), dtype=torch.float8_e5m2)), **{}):
Invalid dtype!

from user code:
   File "/home/justinchu/dev/onnx-script/float8export.py", line 7, in forward
    return input + torch.tensor(1.0, dtype=torch.float8_e5m2)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/justinchu/dev/onnx-script/float8export.py", line 9, in <module>
    torch.onnx.dynamo_export(Float8Module(), torch.randn(1, 3, 224, 224))
  File "<@beartype(torch.onnx._internal.exporter.dynamo_export) at 0x7f5cd348da20>", line 53, in dynamo_export
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py", line 1016, in dynamo_export
    raise OnnxExporterError(
torch.onnx.OnnxExporterError: Failed to export the model to ONNX. Generating SARIF report at report_dynamo_export.sarif. Please report a bug on PyTorch Github: https://github.com/pytorch/pytorch/issues
```

After adding `torch.float8_e5m2` to `_float_dtypes` in `torch/_prims_common/__init__.py`, the new error is

```
[2023-08-09 17:40:21,961] torch._dynamo.eval_frame: [DEBUG] skipping __init__ /home/justinchu/anaconda3/envs/onnx/lib/python3.10/contextlib.py
[2023-08-09 17:40:21,961] torch._dynamo.eval_frame: [DEBUG] skipping __enter__ /home/justinchu/anaconda3/envs/onnx/lib/python3.10/contextlib.py
[2023-08-09 17:40:21,962] torch._dynamo.eval_frame: [DEBUG] skipping helper /home/justinchu/anaconda3/envs/onnx/lib/python3.10/contextlib.py
[2023-08-09 17:40:21,962] torch._dynamo.eval_frame: [DEBUG] skipping __init__ /home/justinchu/anaconda3/envs/onnx/lib/python3.10/contextlib.py
[2023-08-09 17:40:21,962] torch._dynamo.eval_frame: [DEBUG] skipping __enter__ /home/justinchu/anaconda3/envs/onnx/lib/python3.10/contextlib.py
[2023-08-09 17:40:21,962] torch._dynamo.eval_frame: [DEBUG] skipping enable_dynamic /home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py
[2023-08-09 17:40:21,962] torch._dynamo.eval_frame: [DEBUG] skipping wrapped /home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py
[2023-08-09 17:40:21,962] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing forward /home/justinchu/dev/onnx-script/float8export.py:5
[2023-08-09 17:40:21,965] torch._dynamo.symbolic_convert.__trace_source: [DEBUG] TRACE starts_line forward /home/justinchu/dev/onnx-script/float8export.py:5
[2023-08-09 17:40:21,965] torch._dynamo.symbolic_convert.__trace_source: [DEBUG]         def forward(self, input: torch.Tensor):
[2023-08-09 17:40:22,012] torch._dynamo.variables.builder: [DEBUG] wrap_to_fake L['input'] (1, 3, 224, 224) [<DimDynamic.STATIC: 2>, <DimDynamic.STATIC: 2>, <DimDynamic.STATIC: 2>, <DimDynamic.STATIC: 2>] [None, None, None, None]
[2023-08-09 17:40:22,013] torch._dynamo.symbolic_convert.__trace_source: [DEBUG] TRACE starts_line forward /home/justinchu/dev/onnx-script/float8export.py:6
[2023-08-09 17:40:22,013] torch._dynamo.symbolic_convert.__trace_source: [DEBUG]             input = input.to(torch.float8_e5m2)
[2023-08-09 17:40:22,013] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_FAST input []
[2023-08-09 17:40:22,013] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_ATTR to [TensorVariable()]
[2023-08-09 17:40:22,014] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_GLOBAL torch [GetAttrVariable(TensorVariable(), to)]
[2023-08-09 17:40:22,014] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_ATTR float8_e5m2 [GetAttrVariable(TensorVariable(), to), TorchVariable(<module 'torch' from '/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/__init__.py'>)]
[2023-08-09 17:40:22,015] torch._dynamo.symbolic_convert: [DEBUG] TRACE CALL_FUNCTION 1 [GetAttrVariable(TensorVariable(), to), ConstantVariable(dtype)]
[2023-08-09 17:40:22,017] torch._dynamo.symbolic_convert: [DEBUG] TRACE STORE_FAST input [TensorVariable()]
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert.__trace_source: [DEBUG] TRACE starts_line forward /home/justinchu/dev/onnx-script/float8export.py:7
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert.__trace_source: [DEBUG]             return input + torch.tensor(1.0, dtype=torch.float8_e5m2)
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_FAST input []
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_GLOBAL torch [TensorVariable()]
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_ATTR tensor [TensorVariable(), TorchVariable(<module 'torch' from '/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/__init__.py'>)]
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_CONST 1.0 [TensorVariable(), TorchVariable(<built-in method tensor of type object at 0x7f9d7359c000>)]
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_GLOBAL torch [TensorVariable(), TorchVariable(<built-in method tensor of type object at 0x7f9d7359c000>), ConstantVariable(float)]
[2023-08-09 17:40:22,018] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_ATTR float8_e5m2 [TensorVariable(), TorchVariable(<built-in method tensor of type object at 0x7f9d7359c000>), ConstantVariable(float), TorchVariable(<module 'torch' from '/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/__init__.py'>)]
[2023-08-09 17:40:22,019] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_CONST ('dtype',) [TensorVariable(), TorchVariable(<built-in method tensor of type object at 0x7f9d7359c000>), ConstantVariable(float), ConstantVariable(dtype)]
[2023-08-09 17:40:22,019] torch._dynamo.symbolic_convert: [DEBUG] TRACE CALL_FUNCTION_KW 2 [TensorVariable(), TorchVariable(<built-in method tensor of type object at 0x7f9d7359c000>), ConstantVariable(float), ConstantVariable(dtype), ConstantVariable(tuple)]
[2023-08-09 17:40:22,033] torch._dynamo.symbolic_convert: [DEBUG] TRACE BINARY_ADD None [TensorVariable(), TensorVariable()]
Traceback (most recent call last):
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py", line 1006, in dynamo_export
    ).export()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py", line 810, in export
    graph_module = self.options.fx_tracer.generate_fx(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py", line 198, in generate_fx
    graph_module, graph_guard = torch._dynamo.export(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1014, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 321, in _fn
    return fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py", line 153, in wrapped
    return output_adapter.apply(model_func(*args, **kwargs))
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 481, in catch_errors
    return callback(frame, cache_size, hooks, frame_state)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 129, in _fn
    return fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 361, in _convert_frame_assert
    return _compile(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 194, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 431, in _compile
    out_code = transform_code_object(code, transform)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 416, in transform
    tracer.run()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2071, in run
    super().run()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 168, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 561, in call_function
    return wrap_fx_proxy(tx, proxy, **options)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1156, in wrap_fx_proxy
    return wrap_fx_proxy_cls(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1230, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1360, in get_fake_value
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1328, in get_fake_value
    return wrap_fake_exception(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 907, in wrap_fake_exception
    return fn()
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1329, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1394, in run_node
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1381, in run_node
    return node.target(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1235, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1525, in dispatch
    r = func(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_ops.py", line 435, in __call__
    return self._op(*args, **kwargs or {})
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 229, in _fn
    result = fn(*args, **kwargs)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 120, in _fn
    compute_dtype, result_dtype = utils.elementwise_dtypes(
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/__init__.py", line 1373, in elementwise_dtypes
    highest_type = get_higher_type(highest_type, dtype_to_type(x.dtype))
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/_prims_common/__init__.py", line 915, in dtype_to_type
    raise ValueError("Invalid dtype!")
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in function add>(*(FakeTensor(..., size=(1, 3, 224, 224), dtype=torch.float8_e5m2), FakeTensor(..., size=(), dtype=torch.float8_e5m2)), **{}):
Invalid dtype!

from user code:
   File "/home/justinchu/dev/onnx-script/float8export.py", line 7, in forward
    return input + torch.tensor(1.0, dtype=torch.float8_e5m2)


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/justinchu/dev/onnx-script/float8export.py", line 9, in <module>
    torch.onnx.dynamo_export(Float8Module(), torch.randn(1, 3, 224, 224))
  File "<@beartype(torch.onnx._internal.exporter.dynamo_export) at 0x7f9cb6e0dbd0>", line 53, in dynamo_export
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/torch/onnx/_internal/exporter.py", line 1016, in dynamo_export
    raise OnnxExporterError(
torch.onnx.OnnxExporterError: Failed to export the model to ONNX. Generating SARIF report at report_dynamo_export.sarif. Please report a bug on PyTorch Github: https://github.com/pytorch/pytorch/issues
[2023-08-09 17:40:22,046] torch._dynamo.utils: [INFO] TorchDynamo compilation metrics:
[2023-08-09 17:40:22,046] torch._dynamo.utils: [INFO] Function      Runtimes (s)
[2023-08-09 17:40:22,046] torch._dynamo.utils: [INFO] ----------  --------------
[2023-08-09 17:40:22,046] torch._dynamo.utils: [INFO] _compile                 0
```

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @anijain2305 @ipiszy