[pass] Avoid lifting tensors that are too small to initializers

[justinchuby]

Tensors with too few elements are usually not weights and are plenty. Lifting them will make the initializer list very noisy. I added a parameter size_limit to control this and defaulted it to 16.

[Copilot]

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (2)

onnxscript/ir/passes/common/constant_manipulation_test.py:58

• Consider adding test cases that verify the behavior when using the default size_limit (16). It would ensure that tensors with fewer than 16 elements are not lifted.

result = constant_manipulation.LiftConstantsToInitializersPass(lift_all_constants=lift_all_constants, size_limit=0)

onnxscript/ir/passes/common/constant_manipulation.py:121

• [nitpick] Consider using '%d' instead of '%s' in the following logger.debug call to emphasize that size_limit is an integer.

if tensor.size < self.size_limit:

[codecov]

❌ 5 Tests Failed:

Tests completed	Failed	Passed	Skipped
14819	5	14814	2204

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0536_test_layer_normalization_4d_axis1

Stack Traces | 0.006s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_layer_normalization_4d_axis1'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_layer_normalization_4d_axis1' (e=No module named 'tests.onnx_backend_test_code.test_layer_normalization_4d_axis1') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_4d_axis1.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_4d_axis1.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset17
E   
E   @script()
E   def bck_test_layer_normalization_4d_axis1(X: FLOAT[2,3,4,5], W: FLOAT[3,4,5], B: FLOAT[3,4,5]) -> (FLOAT[2,3,4,5], FLOAT[2,1,1,1], FLOAT[2,1,1,1]):
E       Y, Mean, InvStdDev = opset17.LayerNormalization(X, W, B, axis=1)
E       return Y, Mean, InvStdDev

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0397_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.013s run time

onnxscript/converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript/backend/onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript/backend/onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
.../hostedtoolcache/Python/3.11.11.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:95: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:39: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:825: in _translate_call_expr
    attrs = [
onnxscript/converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0125_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.016s run time

onnxscript/converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript/backend/onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript/backend/onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:95: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:39: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:825: in _translate_call_expr
    attrs = [
onnxscript/converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[titaiwangms]

Understand the motivation, but this will create a different behavior to shipped models from GenAI. They lifted all constants with "value". I am not sure whether it's easier to implement for them or they have specific request/process, so they are lifting all tensors.

[justinchuby]

Understand the motivation, but this will create a different behavior to shipped models from GenAI. They lifted all constants with "value". I am not sure whether it's easier to implement for them or they have specific request/process, so they are lifting all tensors.

I think it is possible to users to tweak the value of the size limit to control this behavior. As part of the general optimization passes I think it makes more sense to maintain the proposed default behavior from this PR. If there are special requirements (we need to understand them better than looking at the current proto processing behavior) we can do (1) change the size limit in ort optimization passes (2) update genai so that it is less constrained

[justinchuby]

@kunal-vaishnavi I assume the logic is for pattern matching. Could you share the assumptions the patterns make? What are the patterns? Or is it for different purposes?

[kunal-vaishnavi]

@kunal-vaishnavi I assume the logic is for pattern matching. Could you share the assumptions the patterns make? What are the patterns? Or is it for different purposes?

The ORT transformer optimizer applies its graph fusions on initializers instead of Constant nodes. In order for it to run successfully, all Constant nodes have to be converted to initializers.

[justinchuby]

@kunal-vaishnavi I assume the logic is for pattern matching. Could you share the assumptions the patterns make? What are the patterns? Or is it for different purposes?

The ORT transformer optimizer applies its graph fusions on initializers instead of Constant nodes. In order for it to run successfully, all Constant nodes have to be converted to initializers.

Is this the only fusion? I can see the example is on model weights/biases so this should already be accounted for.

[kunal-vaishnavi]

@kunal-vaishnavi I assume the logic is for pattern matching. Could you share the assumptions the patterns make? What are the patterns? Or is it for different purposes?

The ORT transformer optimizer applies its graph fusions on initializers instead of Constant nodes. In order for it to run successfully, all Constant nodes have to be converted to initializers.

Is this the only fusion? I can see the example is on model weights/biases so this should already be accounted for.

This is just one example. Almost all fusions in the ORT transformer optimizer expect initializers instead of Constant nodes.

[justinchuby]

Thanks. To be on the safe side we can set size_limit to 0 when running the optimize_for_ort api. @titaiwangms

[titaiwangms]

Thanks. To be on the safe side we can set size_limit to 0 when running the optimize_for_ort api. @titaiwangms

We can just apply this pass again with ClearMetadata pass after the export (in builder.py). It's not necessarily to have this change by default if we believe size_limit=16 benefits more in general cases.