Update default QAT recipe path to export in eval mode.

This change aligns the QuantizeStage with the recommended workflow described in https://github.com/pytorch/ao/blob/v0.18.0/torchao/quantization/pt2e/quantize_pt2e.py#L172-L184.

Author: MartinPavella

export/stages.py

@@ -598,10 +598,10 @@ def run(self, artifact: PipelineArtifact) -> None:
                     {0: batch} if isinstance(t, torch.Tensor) else None for t in inputs
                 )
 
-            # QAT requires the model to be in training mode at capture time so
-            # that batch_norm and dropout decompose with training-mode semantics.
-            if recipe.is_qat:
-                model.train()
+            # Export in eval mode for both PTQ and QAT. The torchao-recommended QAT workflow exports the model in eval
+            # mode and calls model.train() only inside the train_fn, after prepare_qat_pt2e.
+            # (see https://github.com/pytorch/ao/blob/v0.18.0/torchao/quantization/pt2e/quantize_pt2e.py#L172-L184)
+            model.eval()
 
             with torch.no_grad():
                 for hook in self._pre_trace_hooks or []:

export/tests/test_export_stages.py

@@ -853,7 +853,7 @@ def test_run_qat_missing_train_fn_raises(self, mock_torch_export: Mock) -> None:
     @patch("executorch.export.stages.prepare_qat_pt2e")
     @patch("executorch.export.stages.ComposableQuantizer")
     @patch("torch.export.export")
-    def test_run_qat_model_put_in_train_mode_before_export(
+    def test_run_qat_model_put_in_eval_mode_before_export(
         self,
         mock_torch_export: Mock,
         mock_composable_quantizer: Mock,
@@ -863,8 +863,9 @@ def test_run_qat_model_put_in_train_mode_before_export(
         mock_move_to_eval: Mock,
         mock_allow_train_eval: Mock,
     ) -> None:
-        """QAT: model.train() must be called before torch.export.export so that
-        batch_norm and dropout decompose with training-mode semantics."""
+        """QAT: model.eval() is called before torch.export.export, following the
+        torchao-recommended QAT workflow (torchao/quantization/pt2e/quantize_pt2e.py).
+        model.train() is called only inside train_fn, after prepare_qat_pt2e."""
         mock_quantizer = self.create_dummy_quantizer()
         mock_recipe = Mock(spec=QuantizationRecipe)
         mock_recipe.quantizers = [mock_quantizer]
@@ -876,9 +877,9 @@ def test_run_qat_model_put_in_train_mode_before_export(
         mock_recipe.pre_convert_passes = None
         mock_recipe.post_convert_passes = None
 
-        # Start model in eval mode; the stage must switch it to train.
-        self.model.eval()
-        self.assertFalse(self.model.training)
+        # Start model in train mode; the stage must switch it to eval.
+        self.model.train()
+        self.assertTrue(self.model.training)
 
         training_at_export_time = []
 
@@ -896,8 +897,8 @@ def capture_training_flag(model, *args, **kwargs):
         stage = QuantizeStage(mock_recipe)
         stage.run(PipelineArtifact(data=self.models_dict, context=self.context))
 
-        # The model must have been in training mode when export was called.
-        self.assertEqual(training_at_export_time, [True])
+        # The model must have been in eval mode when export was called.
+        self.assertEqual(training_at_export_time, [False])
 
     @patch("executorch.export.stages.move_exported_model_to_eval")
     @patch("executorch.export.stages.convert_pt2e")