[IR] Support float4e2m1 #1908

justinchuby · 2024-10-15T00:48:49Z

Support the float4e2m1 dtype from IRv11 (which is not yet released). This allows our tests to pass in the weekly-onnx CI. We use the ml_dtypes.float4_e2m1fn type for numpy conversion. Since ml_dtypes.float4_e2m1fn is only available in the latest ml_dtypes release which has dropped support for python 3.8, I used a conditional logic to build the numpy dtype mapping table.

codecov · 2024-10-15T00:52:27Z

❌ 14 Tests Failed:

Tests completed	Failed	Passed	Skipped
12379	14	12365	1195

View the full list of 3 ❄️ flaky tests

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime test_function_attribute_by_positional_args

Flake rate in main: 20.18% (Passed 1630 times, Failed 412 times)

Stack Traces | 0.002s run time

..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:91: in run
    res = self._run(x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
..../test_torch_nightly/lib/python3.12.../onnx/reference/custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests/eager_mode_test.py:112: in test_function_attribute_by_positional_args
    self.assertEqual(add_with_alpha(1.0, 2.0, 3.0), 7.0)
onnxscript/values.py:529: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript/evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests/eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
.../onnx_opset/_impl/opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript/values.py:301: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript/evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript/evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
..../test_torch_nightly/lib/python3.12.../onnx/reference/reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types &lt;class 'numpy.ndarray'&gt;, &lt;class 'numpy.ndarray'&gt; (binary operator 'Mul').

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime test_function_some_input_by_kwargs

Flake rate in main: 20.18% (Passed 1630 times, Failed 412 times)

Stack Traces | 0.003s run time

..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:91: in run
    res = self._run(x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
..../test_torch_nightly/lib/python3.12.../onnx/reference/custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests/eager_mode_test.py:106: in test_function_some_input_by_kwargs
    self.assertEqual(add_with_alpha(1.0, other=2.0), 3.0)
onnxscript/values.py:529: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript/evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests/eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
.../onnx_opset/_impl/opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript/values.py:301: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript/evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript/evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
..../test_torch_nightly/lib/python3.12.../onnx/reference/reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types &lt;class 'numpy.ndarray'&gt;, &lt;class 'numpy.ndarray'&gt; (binary operator 'Mul').

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime test_function_all_input_by_kwargs

Flake rate in main: 20.18% (Passed 1630 times, Failed 412 times)

Stack Traces | 0.003s run time

..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:91: in run
    res = self._run(x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
..../test_torch_nightly/lib/python3.12.../onnx/reference/custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests/eager_mode_test.py:109: in test_function_all_input_by_kwargs
    self.assertEqual(add_with_alpha(this=1.0, other=2.0), 3.0)
onnxscript/values.py:529: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript/evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests/eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
.../onnx_opset/_impl/opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript/values.py:301: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript/evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript/evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
..../test_torch_nightly/lib/python3.12.../onnx/reference/reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types &lt;class 'numpy.ndarray'&gt;, &lt;class 'numpy.ndarray'&gt; (binary operator 'Mul').

To view individual test run time comparison to the main branch, go to the Test Analytics Dashboard

onnxscript/ir/_core.py

onnxscript/ir/_type_casting.py

justinchuby · 2024-10-15T16:39:13Z

onnxscript/ir/_type_casting.py

+    Returns:
+        A numpy array of float32 reshaped to dims.
+    """
+    return _unpack_uint4_as_uint8(data, dims).view(ml_dtypes.float4_e2m1fn)


Do we need to handle when ml_dtypes.float4_e2m1fn is not available?

xadupre · 2024-10-15T17:34:57Z

onnxscript/ir/_enums_test.py

@@ -32,6 +32,8 @@ def test_enums_are_the_same_as_spec(self):
        self.assertEqual(_enums.DataType.FLOAT8E5M2FNUZ, onnx.TensorProto.FLOAT8E5M2FNUZ)
        self.assertEqual(_enums.DataType.UINT4, onnx.TensorProto.UINT4)
        self.assertEqual(_enums.DataType.INT4, onnx.TensorProto.INT4)
+        if hasattr(onnx.TensorProto, "FLOAT4E2M1"):


I guess the new functions are already tested somewhere else.

justinchuby added 3 commits October 15, 2024 00:38

[IR] Support float4e2m1

e3a4c9a

update

5184490

serde

5321313

justinchuby added the topic: IR Intermediate representation label Oct 15, 2024

justinchuby requested a review from gramalingam October 15, 2024 00:48

github-advanced-security bot found potential problems Oct 15, 2024

View reviewed changes

onnxscript/ir/_core.py Fixed Show fixed Hide fixed

onnxscript/ir/_type_casting.py Fixed Show fixed Hide fixed

onnxscript/ir/_type_casting.py Fixed Show fixed Hide fixed

justinchuby added 2 commits October 14, 2024 18:34

Update _core.py

d4088a3

Update _type_casting.py

5fb716e

justinchuby commented Oct 15, 2024

View reviewed changes

xadupre reviewed Oct 15, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[IR] Support float4e2m1 #1908

[IR] Support float4e2m1 #1908

justinchuby commented Oct 15, 2024 •

edited

Loading

codecov bot commented Oct 15, 2024 •

edited

Loading

justinchuby Oct 15, 2024

xadupre Oct 15, 2024

[IR] Support float4e2m1 #1908

Are you sure you want to change the base?

[IR] Support float4e2m1 #1908

Conversation

justinchuby commented Oct 15, 2024 • edited Loading

codecov bot commented Oct 15, 2024 • edited Loading

❌ 14 Tests Failed:

justinchuby Oct 15, 2024

Choose a reason for hiding this comment

xadupre Oct 15, 2024

Choose a reason for hiding this comment

justinchuby commented Oct 15, 2024 •

edited

Loading

codecov bot commented Oct 15, 2024 •

edited

Loading