From 782f0294cadc7e4c2e2361c385237aaa246fc2e7 Mon Sep 17 00:00:00 2001 From: Guang Yang Date: Wed, 7 Aug 2024 15:53:23 -0700 Subject: [PATCH] Workaround the export issue in torch 2.4 --- tests/utils/test_cache_utils.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index e8d120216a07..5b61999b5625 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -174,6 +174,8 @@ def test_static_cache_exportability(self): """ Tests that static cache works with `torch.export()` """ + import torch + if version.parse(torch.__version__) < version.parse("2.3"): self.skipTest(reason="This test requires torch >= 2.3 to run.") @@ -217,10 +219,15 @@ def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor): set_seed(0) with torch.no_grad(): - from torch.export import ExportedProgram, export + import torch.export._trace + from torch.export import ExportedProgram model = ExportatibleModelWithStaticCache(config, m) - exported_program = export(model, args=(inputs,), kwargs={"input_pos": torch.arange(1)}) + # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal + # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.4.1+ release. + exported_program = torch.export._trace._export( + model, args=(inputs,), kwargs={"input_pos": torch.arange(1)}, pre_dispatch=False, strict=True + ) self.assertTrue(isinstance(exported_program, ExportedProgram))