SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fixed test test_static_cache_exportability with torch 2.4.0 (#32516)

Browse Source 
Workaround the export issue in torch 2.4

Co-authored-by: Guang Yang <guangyang@fb.com>
This commit is contained in:
Guang Yang
2024-08-08 10:13:40 -07:00
committed by GitHub
parent 044281605f
commit 0164560353
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/utils/test_cache_utils.py
View File

@@ -174,6 +174,8 @@ class CacheTest(unittest.TestCase):
""" """
Tests that static cache works with `torch.export()` Tests that static cache works with `torch.export()`
""" """
import torch
if version.parse(torch.__version__) < version.parse("2.3"): if version.parse(torch.__version__) < version.parse("2.3"):
self.skipTest(reason="This test requires torch >= 2.3 to run.") self.skipTest(reason="This test requires torch >= 2.3 to run.")
@@ -217,10 +219,15 @@ class CacheTest(unittest.TestCase):
set_seed(0) set_seed(0)
with torch.no_grad(): with torch.no_grad():
from torch.export import ExportedProgram, export import torch.export._trace
from torch.export import ExportedProgram
model = ExportatibleModelWithStaticCache(config, m) 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)) self.assertTrue(isinstance(exported_program, ExportedProgram))