diff --git a/export/stages.py b/export/stages.py index a716bede272..cb9193ae07e 100644 --- a/export/stages.py +++ b/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 []: diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index 8a96063f157..36e44b92c26 100644 --- a/export/tests/test_export_stages.py +++ b/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")