From e360bd6f8a137deb05d1e538f7ed49a873c5f59e Mon Sep 17 00:00:00 2001 From: Boian Petkantchin Date: Fri, 27 Sep 2024 12:49:24 -0500 Subject: [PATCH] Remove xfail from sharded Llama export to MLIR test https://github.com/iree-org/iree-turbine/pull/175 solves the failing export. --- sharktank/tests/models/llama/sharded_llama_test.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/sharktank/tests/models/llama/sharded_llama_test.py b/sharktank/tests/models/llama/sharded_llama_test.py index bbec7414a..bec2e9775 100644 --- a/sharktank/tests/models/llama/sharded_llama_test.py +++ b/sharktank/tests/models/llama/sharded_llama_test.py @@ -194,7 +194,6 @@ def testCompareToySizedModelToUnsharded(self): actual_decode_cache_state, expected_decode_cache_state, atol=1e-4, rtol=1e-4 ) - @unittest.expectedFailure def testExportToySizedModelToMlir(self): with tempfile.TemporaryDirectory() as temp_dir: sharded_config = deepcopy(self.config) @@ -204,7 +203,7 @@ def testExportToySizedModelToMlir(self): sharded_dataset = Dataset({}, sharded_theta) parameters_path = f"{temp_dir}/parameters.irpa" sharded_dataset.save(f"{temp_dir}/parameters.irpa") - sharded_dataset = Dataset.load(parameters_path) + sharded_dataset = Dataset.load(parameters_path, mmap=False) model = PagedLlamaModelV1(self.theta, self.config) sharded_model = PagedLlamaModelV1( @@ -238,9 +237,5 @@ def _(model, *args, **kwargs) -> torch.Tensor: def _(model, *args, **kwargs) -> torch.Tensor: return model.decode(*args, **kwargs) - # TODO: debug error - # TypeError: Unsupported torch type conversion for - # !torch.vtensor<[3,1,7],complex> - # https://github.com/llvm/torch-mlir/pull/3738 may fix this. output = export(sharded_fxb) output.save_mlir(f"{temp_dir}/program.mlir")