Make our compile tests actually work (#1522)

pytorch · Sep 8, 2024 · 5d5caca · 5d5caca
1 parent c169bcd
commit 5d5caca
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/tests/recipes/test_full_finetune_single_device.py b/tests/recipes/test_full_finetune_single_device.py
@@ -74,9 +74,6 @@ def test_loss(self, compile, config, model_type, ckpt_type, tmpdir, monkeypatch)
         ckpt_dir = ckpt_path.parent
         log_file = gen_log_file_name(tmpdir)
 
-        # To workaround https://github.com/pytorch/torchtune/issues/676
-        if compile:
-            os.environ["TORCH_COMPILE_BACKEND"] = "aot_eager"
         cmd = f"""
         tune run full_finetune_single_device \
             --config {config} \
@@ -99,8 +96,13 @@ def test_loss(self, compile, config, model_type, ckpt_type, tmpdir, monkeypatch)
         with pytest.raises(SystemExit, match=""):
             runpy.run_path(TUNE_PATH, run_name="__main__")
 
+        # Make sure to clear compile state in between tests
+        if compile:
+            torch._dynamo.reset()
+
         loss_values = get_loss_values_from_metric_logger(log_file)
         expected_loss_values = self._fetch_expected_loss_values(model_type)
+
         torch.testing.assert_close(
             loss_values, expected_loss_values, rtol=1e-4, atol=1e-4
         )

diff --git a/tests/recipes/test_lora_finetune_single_device.py b/tests/recipes/test_lora_finetune_single_device.py
@@ -75,9 +75,6 @@ def test_loss(self, compile, config, model_type, ckpt_type, tmpdir, monkeypatch)
         ckpt_dir = ckpt_path.parent
         log_file = gen_log_file_name(tmpdir)
 
-        # To workaround https://github.com/pytorch/torchtune/issues/676
-        if compile:
-            os.environ["TORCH_COMPILE_BACKEND"] = "aot_eager"
         cmd = f"""
         tune run lora_finetune_single_device \
             --config {config} \
@@ -100,6 +97,10 @@ def test_loss(self, compile, config, model_type, ckpt_type, tmpdir, monkeypatch)
         with pytest.raises(SystemExit, match=""):
             runpy.run_path(TUNE_PATH, run_name="__main__")
 
+        # Make sure to clear compile state in between tests
+        if compile:
+            torch._dynamo.reset()
+
         loss_values = get_loss_values_from_metric_logger(log_file)
         expected_loss_values = self._fetch_expected_loss_values(model_type)
         torch.testing.assert_close(
@@ -119,10 +120,6 @@ def test_loss_qlora(self, compile, dtype, tmpdir, monkeypatch):
         ckpt_dir = ckpt_path.parent
         log_file = gen_log_file_name(tmpdir)
 
-        # To workaround https://github.com/pytorch/torchtune/issues/676
-        if compile:
-            os.environ["TORCH_COMPILE_BACKEND"] = "aot_eager"
-
         cmd = f"""
         tune run lora_finetune_single_device
             --config llama2/7B_qlora_single_device \
@@ -145,6 +142,10 @@ def test_loss_qlora(self, compile, dtype, tmpdir, monkeypatch):
         with pytest.raises(SystemExit):
             runpy.run_path(TUNE_PATH, run_name="__main__")
 
+        # Make sure to clear compile state in between tests
+        if compile:
+            torch._dynamo.reset()
+
         loss_values = get_loss_values_from_metric_logger(log_file)
         expected_loss_values = self._fetch_qlora_expected_loss_values(dtype=dtype)
         torch.testing.assert_close(

diff --git a/torchtune/training/_compile.py b/torchtune/training/_compile.py
@@ -34,6 +34,7 @@ def compile_model(
         verbose (bool): Whether to log compile info. Default: True
     Returns:
         None
+
     """
     backend = os.environ.get("TORCH_COMPILE_BACKEND", "inductor")
     if torch_version_ge("2.5.0"):
@@ -65,7 +66,6 @@ def compile_loss(loss: nn.Module, verbose: bool = True) -> None:
     Returns:
         loss (nn.Module): loss with either entire module compiled or (in the case of
             CEWithChunkedOutputLoss) only the upcast and cross-entropy calculation compiled.
-
     """
     backend = os.environ.get("TORCH_COMPILE_BACKEND", "inductor")
     if verbose: