Skip to content

Commit f65eab9

Browse files
xaduprexadupre
committed
fix ut
1 parent efdb81a commit f65eab9

File tree

6 files changed

+69
-37
lines changed

6 files changed

+69
-37
lines changed

‎_unittests/ut_reference/test_backend_onnxruntime_evaluator.py‎

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,25 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
247247
")"
248248
)
249249

250+
if onnx_opset_version() <= 25:
251+
exc = "|".join(
252+
[
253+
"batchnorm_.*_training",
254+
"convinteger_with_padding",
255+
"rms_normalization",
256+
"rotary_embedding_3d",
257+
"rotary_embedding",
258+
# cuda,
259+
"test_Conv3d_dilated.*_cuda",
260+
"test_reduce_.*_empty_set_cuda",
261+
"test_reduce_sum_square_.*_expanded_cuda",
262+
"test_reduce_l1_.*_expanded_cuda",
263+
"test_reduce_l2_.*_expanded_cuda",
264+
"test_reduce_log_sum_.*_expanded_cuda",
265+
]
266+
)
267+
backend_test.exclude(f"({exc})")
268+
250269
if onnx_opset_version() <= 26:
251270
backend_test.exclude(
252271
"(deform_conv"
@@ -261,36 +280,25 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
261280
"|layer_normalization.*expanded"
262281
"|layer_normalization.*expanded"
263282
"|affine_grid.*expanded"
283+
"|test_attention_4d_diff_heads_mask4d_padded_kv.*"
284+
"|test_convinteger_with_padding"
264285
"|test_rnn_seq"
265286
"|test_roialign_aligned_false"
266287
"|test_roialign_aligned_true"
267288
"|test_roialign_mode_max"
289+
"|test_rotary_embedding_no_position_ids_rotary_dim.*"
290+
"|test_rotary_embedding_with_interleaved_rotary_dim.*"
291+
"|test_rotary_embedding_with_rotary_dim*"
268292
"|test_simple_rnn_batchwise"
269293
"|test_simple_rnn_defaults"
270294
"|test_simple_rnn_with_initial_bias"
295+
"|test_swish*"
296+
"|test_tensorscatter*"
297+
"|test_top_k*"
271298
")"
272299
)
273300

274301

275-
if onnx_opset_version() <= 25:
276-
exc = "|".join(
277-
[
278-
"batchnorm_.*_training",
279-
"convinteger_with_padding",
280-
"rms_normalization",
281-
"rotary_embedding_3d",
282-
"rotary_embedding",
283-
# cuda,
284-
"test_Conv3d_dilated.*_cuda",
285-
"test_reduce_.*_empty_set_cuda",
286-
"test_reduce_sum_square_.*_expanded_cuda",
287-
"test_reduce_l1_.*_expanded_cuda",
288-
"test_reduce_l2_.*_expanded_cuda",
289-
"test_reduce_log_sum_.*_expanded_cuda",
290-
]
291-
)
292-
backend_test.exclude(f"({exc})")
293-
294302
if pv.Version(onnxruntime.__version__) <= pv.Version("1.24"):
295303
backend_test.exclude("(test_attention_4d_with|test_attention_4d_gqa)")
296304

‎_unittests/ut_torch_models/test_validate_models.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,17 @@
88
requires_experimental,
99
requires_transformers,
1010
requires_cuda,
11+
has_torch,
12+
has_transformers,
1113
)
1214
from onnx_diagnostic.torch_models.validate import validate_model
1315

1416

17+
torch29_and_tr_main = not has_torch("2.9.9") and has_transformers("4.99999")
18+
19+
1520
class TestValidateModel(ExtTestCase):
21+
@unittest.skipIf(torch29_and_tr_main, "combination not working")
1622
@requires_transformers("4.53")
1723
@requires_torch("2.7.99")
1824
@requires_experimental()
@@ -38,6 +44,7 @@ def test_validate_tiny_llms_bfloat16(self):
3844
self.assertLess(summary["disc_onnx_ort_run_abs"], 2e-2)
3945
self.assertIn("onnx_filename", data)
4046

47+
@unittest.skipIf(torch29_and_tr_main, "combination not working")
4148
@requires_transformers("4.53")
4249
@requires_torch("2.8.99")
4350
@requires_experimental()
@@ -59,6 +66,7 @@ def test_validate_microsoft_phi4_reasoning(self):
5966
self.assertLess(summary["disc_onnx_ort_run_abs"], 2e-5)
6067
self.assertIn("onnx_filename", data)
6168

69+
@unittest.skipIf(torch29_and_tr_main, "combination not working")
6270
@requires_transformers("4.53")
6371
@requires_torch("2.8.99")
6472
@requires_experimental()

‎_unittests/ut_torch_models/test_validate_whole_models1.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
requires_experimental,
1212
requires_onnxscript,
1313
requires_transformers,
14+
has_torch,
15+
has_transformers,
1416
)
1517
from onnx_diagnostic.torch_models.validate import (
1618
get_inputs_for_task,
@@ -22,6 +24,9 @@
2224
from onnx_diagnostic.tasks import supported_tasks
2325

2426

27+
torch29_and_tr_main = not has_torch("2.9.9") and has_transformers("4.99999")
28+
29+
2530
class TestValidateWholeModels1(ExtTestCase):
2631
def test_a_get_inputs_for_task(self):
2732
fcts = supported_tasks()
@@ -193,6 +198,7 @@ def test_k_filter_inputs(self):
193198
ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["a"], model=["a", "b"])
194199
self.assertEqual((ni, nd), (((None,), {"b": 4}), {"b": 30}))
195200

201+
@unittest.skipIf(torch29_and_tr_main, "combination not working")
196202
@requires_torch("2.9.99")
197203
@hide_stdout()
198204
@ignore_warnings(FutureWarning)

‎_unittests/ut_torch_models/test_validate_whole_models2.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,16 @@
77
ignore_warnings,
88
requires_torch,
99
requires_transformers,
10+
has_torch,
11+
has_transformers,
1012
)
1113
from onnx_diagnostic.torch_models.validate import validate_model
1214

15+
torch29_and_tr_main = not has_torch("2.9.9") and has_transformers("4.99999")
16+
1317

1418
class TestValidateWholeModels2(ExtTestCase):
19+
@unittest.skipIf(torch29_and_tr_main, "combination not working")
1520
@requires_torch("2.9")
1621
@hide_stdout()
1722
@ignore_warnings(FutureWarning)

‎_unittests/ut_torch_models/test_validate_whole_models3.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,16 @@
55
ignore_warnings,
66
requires_torch,
77
requires_transformers,
8+
has_torch,
9+
has_transformers,
810
)
911
from onnx_diagnostic.torch_models.validate import validate_model
1012

13+
torch29_and_tr_main = not has_torch("2.9.9") and has_transformers("4.99999")
14+
1115

1216
class TestValidateWholeModels3(ExtTestCase):
17+
@unittest.skipIf(torch29_and_tr_main, "combination not working")
1318
@requires_torch("2.7")
1419
@hide_stdout()
1520
@ignore_warnings(FutureWarning)

‎onnx_diagnostic/torch_export_patches/patches/patch_transformers.py‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2343,30 +2343,32 @@ def forward(
23432343
)
23442344
elif _is_torchdynamo_exporting():
23452345

2346-
def _iteration(a, b, query_states, key_states, value_states):
2346+
def _iteration(start_end, query_states, key_states, value_states):
2347+
a, b = start_end
23472348
q = query_states[:, :, a:b, :]
23482349
k = key_states[:, :, a:b, :]
23492350
v = value_states[:, :, a:b, :]
2350-
return attention_interface(
2351-
self,
2352-
q,
2353-
k,
2354-
v,
2355-
attention_mask=None,
2356-
scaling=self.scaling,
2357-
dropout=0.0 if not self.training else self.attention_dropout,
2358-
is_causal=False,
2359-
**kwargs,
2360-
)[0]
2351+
return (
2352+
attention_interface(
2353+
self,
2354+
q,
2355+
k,
2356+
v,
2357+
attention_mask=None,
2358+
scaling=self.scaling,
2359+
dropout=0.0 if not self.training else self.attention_dropout,
2360+
is_causal=False,
2361+
**kwargs,
2362+
)[0],
2363+
)
23612364

23622365
starts = cu_seqlens[:-1]
23632366
ends = cu_seqlens[1:]
2367+
starts_ends = torch.cat([starts.unsqueeze(1), ends.unsqueeze(1)], dim=1)
23642368
attn_outputs = [
2365-
_iteration(a, b, query_states, key_states, value_states)
2366-
for a, b in zip(starts, ends)
2369+
_iteration(start_end, query_states, key_states, value_states)
2370+
for start_end in starts_ends
23672371
]
2368-
for att in attn_outputs:
2369-
print("B", _is_torchdynamo_exporting(), att.shape)
23702372
attn_output = torch.cat(attn_outputs, dim=1)
23712373
else:
23722374
# Other implementations: Process each chunk separately
@@ -2390,8 +2392,6 @@ def _iteration(a, b, query_states, key_states, value_states):
23902392
)[0]
23912393
for q, k, v in zip(*splits)
23922394
]
2393-
for att in attn_outputs:
2394-
print("A", _is_torchdynamo_exporting(), att.shape)
23952395
attn_output = torch.cat(attn_outputs, dim=1)
23962396

23972397
attn_output = attn_output.reshape(seq_length, -1).contiguous()

0 commit comments

Comments
 (0)