From 6d75a598ae5864b1716f9bb4cb9c1d418f305376 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Tue, 20 Oct 2026 09:56:10 +0200 Subject: [PATCH] Add option for eval mode export in QAT recipe-based export. --- export/recipe.py | 20 ++++++++++ export/stages.py | 9 +++-- export/tests/test_export_stages.py | 61 +++++++++++++++++++++++++++++- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/export/recipe.py b/export/recipe.py index 0295be847a6..00322253e5f 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -162,6 +162,13 @@ class QuantizationRecipe: Each callable receives a GraphModule and must return a GraphModule. post_convert_passes: Optional list of callables applied to the GraphModule after convert_pt2e. Each callable receives a GraphModule and must return a GraphModule. + qat_export_in_eval_mode: When True and is_qat=True, the model is put into eval mode before + torch.export.export(). The train_fn is then responsible for calling + model.train() internally after prepare_qat_pt2e. This follows the + torchao-recommended QAT workflow. When False (default) and is_qat=True, + the model is put into train mode before export so that batch_norm and + dropout decompose with training-mode semantics at capture time. + Ignored for PTQ (always eval). """ quantizers: Optional[List[Quantizer]] = None @@ -170,6 +177,7 @@ class QuantizationRecipe: dynamic_batch_size: bool = False calibration_inputs_fn: Optional[Callable[[], Iterable[tuple]]] = None train_fn: Optional[Callable[["torch.fx.GraphModule"], None]] = None + qat_export_in_eval_mode: bool = False pre_prepare_passes: Optional[ List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] ] = None @@ -252,6 +260,7 @@ class _CombineAccumulator: is_qat_values: list = field(default_factory=list) dynamic_batch_size_values: list = field(default_factory=list) train_fn_values: list = field(default_factory=list) + qat_export_in_eval_mode_values: list = field(default_factory=list) calibration_inputs_fn_values: list = field(default_factory=list) strict_values: list = field(default_factory=list) mode_values: list = field(default_factory=list) @@ -411,6 +420,7 @@ def _combine_quantization_recipe( all_post_prepare_passes: list, all_pre_convert_passes: list, all_post_convert_passes: list, + qat_export_in_eval_mode_values: list, ) -> "Optional[QuantizationRecipe]": """ Build the combined QuantizationRecipe from per-recipe collected lists. @@ -420,6 +430,9 @@ def _combine_quantization_recipe( """ # is_qat must agree: the two flows (QAT vs PTQ) are incompatible. cls._assert_scalar_fields_agree("is_qat", is_qat_values) + cls._assert_scalar_fields_agree( + "qat_export_in_eval_mode", qat_export_in_eval_mode_values + ) # At most one recipe may supply a train_fn. non_none_train_fns = [f for f in train_fn_values if f is not None] @@ -468,6 +481,11 @@ def _combined_calib_fn(): post_prepare_passes=all_post_prepare_passes or None, pre_convert_passes=all_pre_convert_passes or None, post_convert_passes=all_post_convert_passes or None, + qat_export_in_eval_mode=( + qat_export_in_eval_mode_values[0] + if qat_export_in_eval_mode_values + else False + ), ) @classmethod @@ -584,6 +602,7 @@ def _collect_quantization_fields( acc.is_qat_values.append(qr.is_qat) acc.dynamic_batch_size_values.append(qr.dynamic_batch_size) acc.train_fn_values.append(qr.train_fn) + acc.qat_export_in_eval_mode_values.append(qr.qat_export_in_eval_mode) acc.calibration_inputs_fn_values.append(qr.calibration_inputs_fn) if qr.pre_prepare_passes: acc.pre_prepare_passes.extend(qr.pre_prepare_passes) @@ -677,6 +696,7 @@ def _combine_recipes( all_post_prepare_passes=acc.post_prepare_passes, all_pre_convert_passes=acc.pre_convert_passes, all_post_convert_passes=acc.post_convert_passes, + qat_export_in_eval_mode_values=acc.qat_export_in_eval_mode_values, ) combined_lowering_recipe = cls._combine_lowering_recipe( diff --git a/export/stages.py b/export/stages.py index a716bede272..dd6f493b51a 100644 --- a/export/stages.py +++ b/export/stages.py @@ -598,10 +598,13 @@ 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: + # For PTQ always export in eval mode. + # For QAT the default is train mode so that batch_norm and dropout decompose with training-mode + # semantics at capture time, or optionally the model can be exported in eval mode. + if recipe.is_qat and not recipe.qat_export_in_eval_mode: model.train() + else: + 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..8005f1e89bc 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -863,12 +863,14 @@ 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 default (qat_export_in_eval_mode=False): model.train() is called before + torch.export.export so that batch_norm and dropout decompose with + training-mode semantics at capture time.""" mock_quantizer = self.create_dummy_quantizer() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] mock_recipe.is_qat = True + mock_recipe.qat_export_in_eval_mode = False mock_recipe.dynamic_batch_size = False mock_recipe.train_fn = Mock() mock_recipe.pre_prepare_passes = None @@ -899,6 +901,61 @@ def capture_training_flag(model, *args, **kwargs): # The model must have been in training mode when export was called. self.assertEqual(training_at_export_time, [True]) + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_model_put_in_eval_mode_before_export( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """QAT with qat_export_in_eval_mode=True: model.eval() is called before + torch.export.export, following the torchao-recommended QAT workflow. + 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] + mock_recipe.is_qat = True + mock_recipe.qat_export_in_eval_mode = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = Mock() + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + # Start model in train mode; the stage must switch it to eval. + self.model.train() + self.assertTrue(self.model.training) + + training_at_export_time = [] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + + def capture_training_flag(model, *args, **kwargs): + training_at_export_time.append(model.training) + return mock_exported_program + + mock_torch_export.side_effect = capture_training_flag + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + # 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") @patch("executorch.export.stages.prepare_pt2e")