Skip to content

Commit 625b028

Browse files
xaduprexadupre
authored
Add one example on LayerNormalization (#135)
* add one example * doc * improve
1 parent e649e25 commit 625b028

File tree

5 files changed

+247
-26
lines changed

5 files changed

+247
-26
lines changed

‎.github/workflows/documentation.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ jobs:
8282
- name: Generate coverage report
8383
run: |
8484
pip install pytest pytest-cov
85-
PYTHONPATH=. UNITTEST_GOING=1 pytest --cov=./onnx_diagnostic/ --cov-report=xml --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py
85+
PYTHONPATH=. UNITTEST_GOING=1 UNITTEST_DOT=1 pytest --cov=./onnx_diagnostic/ --cov-report=xml --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py
8686
8787
- name: Upload coverage reports to Codecov
8888
uses: codecov/codecov-action@v3

‎_doc/technical/plot_layer_norm_discrepancies.py‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
"""
2+
LayerNormalization implementation cannot be exchanged
3+
=====================================================
4+
5+
This example applies what was illustrated
6+
:ref:`l-plot-parallelized-reduction`, reduction operations
7+
are sensitive to parallelization.
8+
9+
We consider a small model including a layer normalization
10+
followed by a matrix multiplication and we show that replacing
11+
a kernel by another one may significantly impact the output.
12+
13+
The model
14+
+++++++++
15+
"""
16+
17+
import itertools
18+
import pandas
19+
import onnx
20+
import onnx.helper as oh
21+
import onnxruntime
22+
import torch
23+
from onnx_array_api.plotting.graphviz_helper import plot_dot
24+
from onnx_diagnostic.ext_test_case import unit_test_going
25+
from onnx_diagnostic.helpers import max_diff, string_diff, string_type
26+
from onnx_diagnostic.helpers.onnx_helper import onnx_dtype_name, onnx_dtype_to_np_dtype
27+
from onnx_diagnostic.helpers.torch_helper import onnx_dtype_to_torch_dtype
28+
from onnx_diagnostic.helpers.doc_helper import LayerNormalizationOrt, MatMulOrt
29+
from onnx_diagnostic.reference import TorchOnnxEvaluator
30+
31+
TFLOAT = onnx.TensorProto.FLOAT
32+
TFLOAT16 = onnx.TensorProto.FLOAT16
33+
34+
35+
def get_model(itype: int = TFLOAT16):
36+
return oh.make_model(
37+
oh.make_graph(
38+
[
39+
oh.make_node("LayerNormalization", ["X", "scale", "bias"], ["norm"], axis=-1),
40+
oh.make_node("MatMul", ["norm", "weights"], ["mm"]),
41+
oh.make_node("Add", ["mm", "bias2"], ["Z"]),
42+
],
43+
"layer_norm_matmul_add",
44+
[
45+
oh.make_tensor_value_info("X", itype, ["a", "b", "c"]),
46+
oh.make_tensor_value_info("scale", itype, ["c"]),
47+
oh.make_tensor_value_info("bias", itype, ["c"]),
48+
oh.make_tensor_value_info("weights", itype, ["c", "c"]),
49+
oh.make_tensor_value_info("bias2", itype, ["c"]),
50+
],
51+
[oh.make_tensor_value_info("Z", itype, ["a", "b", "c"])],
52+
),
53+
ir_version=9,
54+
opset_imports=[oh.make_opsetid("", 18)],
55+
)
56+
57+
58+
model = get_model()
59+
plot_dot(model)
60+
61+
# %%
62+
# Let's compare two runtimes
63+
# ++++++++++++++++++++++++++
64+
#
65+
# That will be :epkg:`onnxruntime` and
66+
# :class:`onnx_diagnostic.reference.TorchOnnxEvaluator`.
67+
68+
last_dim = 64 if unit_test_going() else 1152
69+
70+
71+
def make_feeds(last_dim: int):
72+
return {
73+
"X": (torch.rand((32, 1024, last_dim), dtype=torch.float16) - 0.5) * 120,
74+
"scale": torch.rand((last_dim,), dtype=torch.float16),
75+
"bias": torch.rand((last_dim,), dtype=torch.float16),
76+
"weights": torch.rand((last_dim, last_dim), dtype=torch.float16),
77+
"bias2": torch.rand((last_dim,), dtype=torch.float16),
78+
}
79+
80+
81+
def cast_feeds(itype, provider, feeds):
82+
np_feeds = {k: v.detach().numpy() for k, v in feeds.items()}
83+
if provider == "CUDA":
84+
if not torch.cuda.is_available():
85+
return None, None
86+
tch_feeds = {k: v.to("cuda") for k, v in feeds.items()}
87+
ort_feeds = np_feeds
88+
else:
89+
tch_feeds = feeds.copy()
90+
tch_feeds["X"] = tch_feeds["X"][:2] # too long otherwise
91+
ort_feeds = np_feeds.copy()
92+
ort_feeds["X"] = ort_feeds["X"][:2]
93+
tch_feeds = {k: v.to(ttype) for k, v in tch_feeds.items()}
94+
ort_feeds = {k: v.astype(np_dtype) for k, v in ort_feeds.items()}
95+
return tch_feeds, ort_feeds
96+
97+
98+
feeds = make_feeds(last_dim)
99+
kws = dict(with_shape=True, with_min_max=True, with_device=True)
100+
data = []
101+
baseline = {}
102+
103+
for provider, itype in itertools.product(["CPU", "CUDA"], [TFLOAT, TFLOAT16]):
104+
ttype = onnx_dtype_to_torch_dtype(itype)
105+
np_dtype = onnx_dtype_to_np_dtype(itype)
106+
tch_feeds, ort_feeds = cast_feeds(itype, provider, feeds)
107+
if tch_feeds is None:
108+
continue
109+
110+
model = get_model(itype)
111+
print()
112+
print(f"-- running on {provider} with {onnx_dtype_name(itype)}")
113+
print("-- running with torch")
114+
torch_sess = TorchOnnxEvaluator(model, providers=[f"{provider}ExecutionProvider"])
115+
expected = torch_sess.run(None, tch_feeds)
116+
baseline[itype, provider, "torch"] = expected
117+
print(f"-- torch: {string_type(expected, **kws)}")
118+
119+
print("-- running with ort")
120+
ort_sess = onnxruntime.InferenceSession(
121+
model.SerializeToString(), providers=[f"{provider}ExecutionProvider"]
122+
)
123+
got = ort_sess.run(None, ort_feeds)
124+
baseline[itype, provider, "ort"] = got
125+
print(f"-- ort: {string_type(got, **kws)}")
126+
diff = max_diff(expected, got, hist=True)
127+
print(f"-- diff {string_diff(diff)}")
128+
129+
# memorize the data
130+
diff["dtype"] = onnx_dtype_name(itype)
131+
diff["provider"] = provider
132+
diff.update(diff["rep"])
133+
del diff["rep"]
134+
del diff["dnan"]
135+
del diff[">100.0"]
136+
del diff[">10.0"]
137+
data.append(diff)
138+
139+
# %%
140+
df = pandas.DataFrame(data).set_index(["provider", "dtype"])
141+
print(df)
142+
143+
# %%
144+
# Visually.
145+
146+
df["abs"].plot(title="Discrepancies ORT / torch for LayerNorm(X) @ W + B")
147+
148+
# %%
149+
# The discrepancies are significant on CUDA, higher for float16.
150+
# Let's see which operator is responsible for them,
151+
# *LayerNormalization* or *MatMul*.
152+
153+
# %%
154+
# The discrepancies come from?
155+
# ++++++++++++++++++++++++++++
156+
#
157+
# We mix torch and onnxruntime to execute the kernels.
158+
159+
data = []
160+
161+
for mod, provider, itype in itertools.product(
162+
["ORT-TORCH", "TORCH-ORT"], ["CPU", "CUDA"], [TFLOAT, TFLOAT16]
163+
):
164+
ttype = onnx_dtype_to_torch_dtype(itype)
165+
np_dtype = onnx_dtype_to_np_dtype(itype)
166+
tch_feeds, _ = cast_feeds(itype, provider, feeds)
167+
if tch_feeds is None:
168+
continue
169+
170+
custom_kernels = (
171+
{("", "LayerNormalization"): LayerNormalizationOrt}
172+
if mod == "ORT-TORCH"
173+
else {("", "MatMul"): MatMulOrt}
174+
)
175+
176+
model = get_model(itype)
177+
print()
178+
print(f"-- {mod} running on {provider} with {onnx_dtype_name(itype)}")
179+
sess = TorchOnnxEvaluator(
180+
model,
181+
custom_kernels=custom_kernels,
182+
providers=[f"{provider}ExecutionProvider"],
183+
)
184+
got = sess.run(None, tch_feeds)
185+
print(f"-- {mod}: {string_type(got, **kws)}")
186+
187+
difft = max_diff(baseline[itype, provider, "torch"], got)
188+
print(f"-- diff with torch {string_diff(difft)}")
189+
diffo = max_diff(baseline[itype, provider, "ort"], got)
190+
print(f"-- diff with ort {string_diff(diffo)}")
191+
192+
data.append(
193+
dict(
194+
model=mod,
195+
dtype=onnx_dtype_name(itype),
196+
provider=provider,
197+
diff_ort=diffo["abs"],
198+
diff_torch=difft["abs"],
199+
)
200+
)
201+
202+
# %%
203+
df = pandas.DataFrame(data).set_index(["model", "provider", "dtype"])
204+
df = df.sort_index()
205+
print(df)
206+
207+
# %%
208+
# Visually.
209+
210+
df[["diff_ort", "diff_torch"]].plot(title="ORT/Torch or Torch/ORT for LayerNorm(X) @ W + B")

‎_doc/technical/plot_parallelized_reduction.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
"""
2+
.. _l-plot-parallelized-reduction:
3+
24
Reproducible Parallelized Reduction is difficult
35
================================================
46

‎_unittests/ut_xrun_doc/test_documentation_technical.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,13 @@ def add_test_methods(cls):
6868
this = os.path.abspath(os.path.dirname(__file__))
6969
fold = os.path.normpath(os.path.join(this, "..", "..", "_doc", "technical"))
7070
found = os.listdir(fold)
71+
has_dot = int(os.environ.get("UNITTEST_DOT", "0"))
7172
for name in found:
7273
if not name.endswith(".py") or not name.startswith("plot_"):
7374
continue
7475
reason = None
76+
if not reason and not has_dot and name in {"plot_layer_norm_discrepancies.py"}:
77+
reason = "dot not installed"
7578

7679
if reason:
7780

‎onnx_diagnostic/helpers/doc_helper.py‎

Lines changed: 31 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@
22
import onnx
33
import onnx.helper as oh
44
import torch
5-
from .torch_helper import onnx_dtype_to_torch_dtype, torch_dtype_to_onnx_dtype
65
from ..reference.torch_ops import OpRunKernel, OpRunTensor
6+
from .torch_helper import onnx_dtype_to_torch_dtype, torch_dtype_to_onnx_dtype
7+
from .ort_session import InferenceSessionForTorch
78

89

910
class LayerNormalizationOrt(OpRunKernel):
@@ -36,50 +37,55 @@ def __init__(
3637
self._cache: Dict[Tuple[int, int], onnx.ModelProto] = {}
3738
self.is_cpu = torch.device("cpu") == self.device
3839

39-
def _make_model(self, itype: int, rank: int) -> onnx.ModelProto:
40+
def _make_model(self, itype: int, rank: int, has_bias: bool) -> onnx.ModelProto:
4041
shape = [*["d{i}" for i in range(rank - 1)], "last"]
4142
layer_model = oh.make_model(
4243
oh.make_graph(
4344
[
4445
oh.make_node(
4546
"LayerNormalization",
46-
["X", "W", "B"],
47+
["X", "W", "B"] if has_bias else ["X", "W"],
4748
["Z"],
4849
axis=self.axis,
4950
epsilon=self.epsilon,
5051
)
5152
],
5253
"dummy",
53-
[
54-
oh.make_tensor_value_info("X", itype, shape),
55-
oh.make_tensor_value_info("W", itype, ["last"]),
56-
oh.make_tensor_value_info("B", itype, ["last"]),
57-
],
54+
(
55+
[
56+
oh.make_tensor_value_info("X", itype, shape),
57+
oh.make_tensor_value_info("W", itype, ["last"]),
58+
oh.make_tensor_value_info("B", itype, ["last"]),
59+
]
60+
if has_bias
61+
else [
62+
oh.make_tensor_value_info("X", itype, shape),
63+
oh.make_tensor_value_info("W", itype, ["last"]),
64+
]
65+
),
5866
[oh.make_tensor_value_info("Z", itype, shape)],
5967
),
6068
ir_version=9,
6169
opset_imports=[oh.make_opsetid("", 18)],
6270
)
63-
import onnxruntime
64-
6571
provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
66-
return onnxruntime.InferenceSession(
67-
layer_model.SerializeToString(), providers=[provider]
68-
)
72+
self._provider = provider
73+
return InferenceSessionForTorch(layer_model, providers=[provider])
6974

7075
def run(self, x, scale, bias=None):
7176
itype = torch_dtype_to_onnx_dtype(x.dtype)
7277
rank = len(x.shape)
7378
key = itype, rank
7479
if key not in self._cache:
75-
self._cache[key] = self._make_model(itype, rank)
80+
self._cache[key] = self._make_model(itype, rank, bias is not None)
7681
sess = self._cache[key]
77-
feeds = dict(X=x, W=scale)
82+
if self.verbose:
83+
print(f"[LayerNormalizationOrt] running on {self._provider!r}")
84+
feeds = dict(X=x.tensor, W=scale.tensor)
7885
if bias is not None:
79-
feeds["B"] = bias
80-
feeds = {k: v.tensor.detach().cpu().numpy() for k, v in feeds.items()}
86+
feeds["B"] = bias.tensor
8187
got = sess.run(None, feeds)[0]
82-
return OpRunTensor(torch.from_numpy(got).to(x.dtype).to(x.device))
88+
return OpRunTensor(got)
8389

8490

8591
class MatMulOrt(OpRunKernel):
@@ -117,12 +123,11 @@ def _make_model(self, itype: int, ranka: int, rankb: int) -> onnx.ModelProto:
117123
[oh.make_tensor_value_info("C", itype, shapec)],
118124
),
119125
ir_version=9,
120-
opset_imports=[oh.make_opsetid("", 17)],
126+
opset_imports=[oh.make_opsetid("", 18)],
121127
)
122-
import onnxruntime
123-
124128
provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
125-
return onnxruntime.InferenceSession(model.SerializeToString(), providers=[provider])
129+
self._provider = provider
130+
return InferenceSessionForTorch(model, providers=[provider])
126131

127132
def run(self, a, b):
128133
itype = torch_dtype_to_onnx_dtype(a.dtype)
@@ -131,7 +136,8 @@ def run(self, a, b):
131136
if key not in self._cache:
132137
self._cache[key] = self._make_model(itype, ranka, rankb)
133138
sess = self._cache[key]
134-
feeds = dict(A=a, B=b)
135-
feeds = {k: v.tensor.detach().cpu().numpy() for k, v in feeds.items()}
139+
if self.verbose:
140+
print(f"[MatMulOrt] running on {self._provider!r}")
141+
feeds = dict(A=a.tensor, B=b.tensor)
136142
got = sess.run(None, feeds)[0]
137-
return OpRunTensor(torch.from_numpy(got).to(a.dtype).to(a.device))
143+
return OpRunTensor(got)

0 commit comments

Comments
 (0)