From 8a71747474c2458181b0614149b604ccb1025e9a Mon Sep 17 00:00:00 2001 From: Kevin Wu <158515169+kevinwuTT@users.noreply.github.com> Date: Wed, 22 Jan 2025 16:42:01 -0500 Subject: [PATCH] Fix and add workaround for falcon-7b model test (#719) * Add fallback to arange and argmax variations for falcon-7b * Append argmax blocklist instead since it already exists * Add arange.start_step to constantfolding pass instead * Fix arange test * Correct error message --- tests/lowering/creation/test_arange.py | 2 +- torch_ttnn/passes/constant_folding_pass.py | 1 + torch_ttnn/passes/lowering/to_tt_guard.py | 7 +++++++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/lowering/creation/test_arange.py b/tests/lowering/creation/test_arange.py index be863f311..fa81cb1dd 100644 --- a/tests/lowering/creation/test_arange.py +++ b/tests/lowering/creation/test_arange.py @@ -90,6 +90,6 @@ def test_arange_start_step(device, input_shapes): # Check the graph has be rewritten and contain ttnn ops nodes = list(option._out_fx_graphs[0].nodes) - assert [node.target for node in nodes].count(ttnn.arange) == 1 + assert [node.target for node in nodes].count(ttnn.arange) == 1 or [node.op for node in nodes].count("get_attr") # Check inference result assert torch.allclose(result_before, result_after) diff --git a/torch_ttnn/passes/constant_folding_pass.py b/torch_ttnn/passes/constant_folding_pass.py index dbd8de010..0b62d404f 100644 --- a/torch_ttnn/passes/constant_folding_pass.py +++ b/torch_ttnn/passes/constant_folding_pass.py @@ -11,6 +11,7 @@ def __init__(self): torch.ops.aten.lift_fresh_copy.default, torch.ops.aten.pow.Tensor_Tensor, torch.ops.aten.arange.start, + torch.ops.aten.arange.start_step, torch.ops.aten.unsqueeze.default, torch.ops.aten.arange.default, torch.ops.aten.view.default, diff --git a/torch_ttnn/passes/lowering/to_tt_guard.py b/torch_ttnn/passes/lowering/to_tt_guard.py index bfe128809..a8ac05350 100644 --- a/torch_ttnn/passes/lowering/to_tt_guard.py +++ b/torch_ttnn/passes/lowering/to_tt_guard.py @@ -83,6 +83,13 @@ ["List[Tensor] tensors = [<[13600]>, <[13600]>, <[13600]>, <[13600]>]", "int dim = 1"], ] +############################################################ +# EXTRA BLOCKLIST OF falcon-7b-instruct +############################################################ +# RuntimeError: TT_THROW @ /tmp/build-via-sdist-d26xvola/ttnn-0.54.0rc18+wormhole.b0/ttnn/cpp/ttnn/device_operation.hpp:487: tt::exception +# Unsupported storage type +aten_argmax_default_blocklist += [["Tensor<[1, 7]> self = ?", "Optional[int] dim = 1", "bool keepdim = True"]] + ############################################################ # EXTRA BLOCKLIST OF retinanet_resnet50_fpn_v2 ############################################################