avoid compiling compute_loss

Summary:
# Context
[Describe motivations and existing situation that led to creating this diff. Don't be cheap with context, it is the basis for a good code review.]

# This diff
[List all the changes that this diff introduces and explain the ones that are not trivial. Give directions for the reviewer if needed.]

# What’s next
[If this diff is part of a stack or if it has direct continuation in a future diff, share these plans with your reviewer.]

Differential Revision: D48581289

fbshipit-source-id: 4abd5932a1a18f2983605cacc886abe6433a0711

Author: galrotem

tests/framework/test_auto_unit.py

@@ -284,94 +284,6 @@ def test_compile_state_dict(self) -> None:
                 torch.allclose(my_module_state_dict[k], compiled_state_dict[k])
             )
 
-    @unittest.skipUnless(
-        condition=COMPILE_AVAIL,
-        reason="This test needs PyTorch 1.13 or greater to run.",
-    )
-    def test_compile_eager(self) -> None:
-        """
-        e2e torch compile test
-        """
-
-        my_module = torch.nn.Linear(2, 2)
-
-        input_dim = 2
-        dataset_len = 16
-        batch_size = 2
-        max_epochs = 1
-
-        auto_unit = DummyAutoUnit(
-            module=my_module,
-            torch_compile_params=TorchCompileParams(backend="eager"),
-        )
-
-        train_dl = generate_random_dataloader(dataset_len, input_dim, batch_size)
-        self.assertFalse(auto_unit._compile_used)
-        train(auto_unit, train_dl, max_epochs=max_epochs)
-        self.assertTrue(auto_unit._compile_used)
-
-    @unittest.skipUnless(
-        condition=COMPILE_AVAIL,
-        reason="This test needs PyTorch 1.13 or greater to run.",
-    )
-    @unittest.skipUnless(
-        condition=cuda_available, reason="This test needs a GPU host to run."
-    )
-    def test_compile_train(self) -> None:
-        """
-        e2e torch compile on train
-        """
-
-        my_module = torch.nn.Linear(2, 2)
-
-        input_dim = 2
-        dataset_len = 16
-        batch_size = 2
-        max_epochs = 1
-
-        auto_unit = DummyAutoUnit(
-            module=my_module,
-            torch_compile_params=TorchCompileParams(backend="inductor"),
-        )
-
-        train_dl = generate_random_dataloader(dataset_len, input_dim, batch_size)
-
-        self.assertFalse(auto_unit._compile_used)
-        train(auto_unit, train_dl, max_epochs=max_epochs)
-        self.assertTrue(auto_unit._compile_used)
-
-    @unittest.skipUnless(
-        condition=COMPILE_AVAIL,
-        reason="This test needs PyTorch 1.13 or greater to run.",
-    )
-    @unittest.skipUnless(
-        condition=cuda_available, reason="This test needs a GPU host to run."
-    )
-    def test_compile_eval(self) -> None:
-        """
-        e2e torch compile on eval
-        """
-
-        my_module = torch.nn.Linear(2, 2)
-
-        input_dim = 2
-        dataset_len = 16
-        batch_size = 2
-
-        auto_unit = DummyAutoUnit(
-            module=my_module,
-            torch_compile_params=TorchCompileParams(backend="inductor"),
-        )
-
-        input_dim = 2
-        dataset_len = 8
-        batch_size = 2
-
-        eval_dl = generate_random_dataloader(dataset_len, input_dim, batch_size)
-        self.assertFalse(auto_unit._compile_used)
-        evaluate(auto_unit, eval_dl)
-        self.assertTrue(auto_unit._compile_used)
-
     @unittest.skipUnless(
         condition=COMPILE_AVAIL,
         reason="This test needs PyTorch 1.13 or greater to run.",
@@ -983,16 +895,9 @@ def test_predict_detect_anomaly(self, mock_detect_anomaly) -> None:
 
 # pyre-fixme[11]: Annotation `Batch` is not defined as a type.
 class DummyAutoUnit(AutoUnit[Batch]):
-    # pyre-fixme[3]: Return type must be annotated.
-    # pyre-fixme[2]: Parameter must be annotated.
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._compile_used = False
 
     # pyre-fixme[3]: Return annotation cannot contain `Any`.
     def compute_loss(self, state: State, data: Batch) -> Tuple[torch.Tensor, Any]:
-        if COMPILE_AVAIL:
-            self._compile_used = torch._dynamo.is_compiling()
         inputs, targets = data
         outputs = self.module(inputs)
         loss = torch.nn.functional.cross_entropy(outputs, targets)

torchtnt/framework/auto_unit.py

@@ -470,11 +470,6 @@ def __init__(
             )
 
         if torch_compile_params:
-            # pyre-ignore
-            self.compute_loss = torch.compile(
-                self.compute_loss,
-                **asdict(torch_compile_params),
-            )
             try:
                 # use in-place compile to avoid altering the state_dict keys
                 module.compile(**asdict(torch_compile_params))