diff --git a/optimum/exporters/executorch/convert.py b/optimum/exporters/executorch/convert.py index 49fb563..fe4bead 100644 --- a/optimum/exporters/executorch/convert.py +++ b/optimum/exporters/executorch/convert.py @@ -55,7 +55,10 @@ def export_to_executorch( output_dir (`Union[str, Path]`): Path to the directory where the resulting ExecuTorch model will be saved. **kwargs: - Additional configuration options passed to the recipe. + Additional configuration options passed to the recipe. Notably `generate_etrecord` + (`bool`, defaults to `False`): if `True`, saves an ETRecord (`{name}_etrecord.bin`) alongside + the `.pte` file for use with the ExecuTorch Inspector. Currently only honored by the + `xnnpack` recipe; other recipes ignore the flag and no ETRecord is produced. Returns: `ExecuTorchProgram`: @@ -67,6 +70,9 @@ def export_to_executorch( - The resulting ExecuTorch program is serialized and saved to the output directory. """ + # Extract generate_etrecord from kwargs (default: False to avoid unwanted files) + generate_etrecord = kwargs.get("generate_etrecord", False) + # Dynamically discover and import registered recipes discover_recipes() @@ -87,4 +93,25 @@ def export_to_executorch( ) prog.write_tensor_data_to_file(output_dir) + # Save ETRecord only if explicitly requested + if generate_etrecord: + try: + etrecord = prog.get_etrecord() + except RuntimeError: + # Raised when the recipe didn't generate an ETRecord (e.g. only the xnnpack + # recipe currently forwards generate_etrecord to the lowering pipeline). + logging.warning( + f"generate_etrecord=True was requested but no ETRecord was produced for {name}; " + "the selected recipe may not support ETRecord generation yet." + ) + else: + etrecord_path = os.path.join(output_dir, f"{name}_etrecord.bin") + try: + etrecord.save(etrecord_path) + except Exception as e: + logging.warning(f"Failed to save ETRecord for {name}: {e}") + else: + etrecord_size_kb = os.path.getsize(etrecord_path) / 1024 + logging.info(f"Saved ETRecord to {etrecord_path} ({etrecord_size_kb:.2f} KB)") + return executorch_progs diff --git a/optimum/exporters/executorch/recipes/xnnpack.py b/optimum/exporters/executorch/recipes/xnnpack.py index ef37cc7..55ad183 100644 --- a/optimum/exporters/executorch/recipes/xnnpack.py +++ b/optimum/exporters/executorch/recipes/xnnpack.py @@ -61,6 +61,8 @@ def export_to_executorch_with_xnnpack( The PyTorch model to be exported to ExecuTorch. **kwargs: Additional keyword arguments for recipe-specific configurations, e.g. export using different example inputs, or different compile/bechend configs. + `generate_etrecord` (`bool`, defaults to `False`): if `True`, generates an ETRecord during + lowering so `convert.py` can save it alongside the `.pte` file. Returns: Dict[str, ExecutorchProgram]: @@ -92,6 +94,7 @@ def _lower_to_executorch( ), constant_methods=metadata, transform_passes=[RemovePaddingIdxEmbeddingPass()], + generate_etrecord=kwargs.get("generate_etrecord", False), ) et_prog = et_prog.to_executorch( config=ExecutorchBackendConfig(**backend_config_dict),