Skip to content

fix: guard skip-rmsnorm fusion against mismatched dtypes (#3155) #3158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aanelluri wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: guard skip-rmsnorm fusion against mismatched dtypes (#3155) #3158

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions onnxscript/rewriter/ort_fusions/skip_normalization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,14 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
if stash_type != 1:
return check_result.fail("Stash type is not supported.")

if (
input.dtype is not None
and skip.dtype is not None
and gamma.dtype is not None
and (input.dtype != skip.dtype or input.dtype != gamma.dtype)
):
return check_result.fail("Dtype mismatch", [input, skip, gamma])

return check_result

def rewrite(
Expand Down
43 changes: 43 additions & 0 deletions onnxscript/rewriter/ort_fusions/skip_normalization_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

import unittest

import onnx
import onnx_ir
import onnxscript.optimizer
from onnxscript.rewriter.models._bart_encoder import bart_encoder_test
from onnxscript.rewriter.models._smollm_1 import smollm_test_1
Expand All @@ -18,6 +20,47 @@


class TestSkipNormalization(unittest.TestCase):
def test_skip_rms_norm_dtype_mismatch(self):
def make_model(gamma_dtype):
graph = onnx.helper.make_graph(
[
onnx.helper.make_node("Add", ["input", "skip"], ["skip_sum"]),
onnx.helper.make_node(
"SimplifiedLayerNormalization",
["skip_sum", "gamma"],
["output"],
axis=-1,
stash_type=1,
),
],
"skip_rms_norm",
[
onnx.helper.make_tensor_value_info(
"input", onnx.TensorProto.FLOAT16, [2, 8, 64]
),
onnx.helper.make_tensor_value_info(
"skip", onnx.TensorProto.FLOAT16, [2, 8, 64]
),
onnx.helper.make_tensor_value_info("gamma", gamma_dtype, [64]),
],
[
onnx.helper.make_tensor_value_info(
"output", onnx.TensorProto.FLOAT16, [2, 8, 64]
)
],
)
return onnx_ir.from_proto(onnx.helper.make_model(graph))

mixed_dtype_model = make_model(onnx.TensorProto.FLOAT)
fuse_skip_rms_normalization(mixed_dtype_model)
mixed_dtype_op_types = [node.op_type for node in mixed_dtype_model.graph]
self.assertNotIn("SkipSimplifiedLayerNormalization", mixed_dtype_op_types)

matching_dtype_model = make_model(onnx.TensorProto.FLOAT16)
fuse_skip_rms_normalization(matching_dtype_model)
matching_dtype_op_types = [node.op_type for node in matching_dtype_model.graph]
self.assertIn("SkipSimplifiedLayerNormalization", matching_dtype_op_types)

def test_smollm(self):
smollm_test = smollm_test_1()
model = smollm_test.get_onnx_model()
Expand Down