diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 5cdc9a40940..ffe9a89a570 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -14,8 +14,10 @@ import json import logging import math +import os import re import shlex +import tempfile from functools import partial from importlib import resources as _resources from json import JSONDecodeError @@ -24,7 +26,6 @@ import torch from executorch.devtools.backend_debug import print_delegation_info -from executorch.devtools.etrecord import generate_etrecord as generate_etrecord_func from executorch.examples.models.llama.hf_download import ( download_and_convert_hf_checkpoint, ) @@ -762,7 +763,6 @@ def _ensure_static_kvq_out_variants(llm_config: LlmConfig) -> None: libraries = [llm_config.export.so_library] else: import glob - import os import executorch from executorch.extension.pybindings import portable_lib # noqa # usort: skip @@ -1186,7 +1186,6 @@ def _to_edge_and_lower_llama_xnnpack( for partitioner in partitioners: logging.info(f"--> {partitioner.__class__.__name__}") - # TODO: Enable generating ETRecord with XNNPack and to_edge_transform_and_lower(). if generate_etrecord: builder_exported.generate_etrecord = True @@ -1214,6 +1213,53 @@ def _to_edge_and_lower_llama_xnnpack( ) +def _etrecord_path_beside_model(builder) -> str: + """Where to write the etrecord so it lands beside the model. + + `save_pte_program` uses an output name ending in `.pte` as the path verbatim, ignoring + `output_dir`, so the model's own directory is the only reliable anchor. Falls back to + `output_dir` before the model is saved. + """ + saved = builder.get_saved_pte_filename() + directory = os.path.dirname(saved) if saved else builder.output_dir + return os.path.join(directory, "etrecord.bin") + + +def _save_etrecord_if_generated(builder) -> None: + """Write the etrecord the lowering produced, if there is one. + + The record goes beside the model rather than under `output_dir`, because an output name + ending in `.pte` is used as the path verbatim and so bypasses `output_dir`. + + A failed write costs the record and not the export, because a debug artifact must not + cost a caller the model. The record is staged under a temporary name and moved into + place, so a failure part way cannot leave a half-written record under the real name. + """ + try: + etrecord = builder.export_program.get_etrecord() + except RuntimeError: + # Not generated, which is the normal case. + return + + path = _etrecord_path_beside_model(builder) + try: + with tempfile.NamedTemporaryFile( + dir=os.path.dirname(path) or ".", suffix=".bin", delete=False + ) as temporary: + staged = temporary.name + try: + etrecord.save(staged) + os.replace(staged, path) + except BaseException: + os.unlink(staged) + raise + except Exception as error: + logging.warning("Could not write %s: %s", path, error) + return + + logging.info("Generated %s", path) + + def _to_edge_and_lower_llama_openvino( builder_exported, modelname, @@ -1420,6 +1466,8 @@ def _to_edge_and_lower_llama( # noqa: C901 generate_etrecord: bool = False, verbose: bool = False, ): + # Set before the edge export, because that is where the record is created. + builder_exported.generate_etrecord = generate_etrecord builder_exported_to_edge = builder_exported.pt2e_quantize( quantizers ).export_to_edge() @@ -1528,9 +1576,6 @@ def _to_edge_and_lower_llama( # noqa: C901 if not builder_exported_to_edge.edge_manager: raise ValueError("Unable to generate etrecord due to missing edge manager.") - logging.info("Generating etrecord") - # Copy the edge manager which will be serialized into etrecord. This is memory-wise expensive. - edge_manager_copy = copy.deepcopy(builder_exported_to_edge.edge_manager) builder = builder_exported_to_edge.to_backend(partitioners) if verbose: print_delegation_info(builder.edge_manager.exported_program().graph_module) @@ -1543,15 +1588,8 @@ def _to_edge_and_lower_llama( # noqa: C901 builder = builder.to_executorch( passes=additional_passes, ) - - # Generate ETRecord - if edge_manager_copy: - generate_etrecord_func( - et_record="etrecord.bin", - edge_dialect_program=edge_manager_copy, - executorch_program=builder.export_program, - ) - logging.info("Generated etrecord.bin") + # The record rides along with the program, so _save_etrecord_if_generated writes it + # after the model like every other path. else: builder = builder_exported_to_edge.to_backend(partitioners) if verbose: @@ -1693,6 +1731,7 @@ def _export_llama_multimethod(llm_config: LlmConfig) -> LLMEdgeManager: first_builder.dtype, ) first_builder.save_to_pte(output_file) + _save_etrecord_if_generated(first_builder) return first_builder @@ -1873,6 +1912,7 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 builder.dtype, ) builder.save_to_pte(output_file) + _save_etrecord_if_generated(builder) return builder diff --git a/examples/models/llama/tests/test_export_llama_lib.py b/examples/models/llama/tests/test_export_llama_lib.py index f33f4489c91..96640045e33 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -5,10 +5,15 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import os +import tempfile import unittest from unittest.mock import patch +import torch from executorch.devtools.backend_debug import get_delegation_info +from executorch.devtools.etrecord import parse_etrecord +from executorch.devtools.etrecord._etrecord import ETRecord try: from executorch.backends.arm.quantizer.arm_quantizer import ( @@ -30,8 +35,10 @@ build_args_parser, get_quantizer_and_quant_params, ) +from executorch.extension.llm.export.builder import LLMEdgeManager from executorch.extension.llm.export.config.llm_config import ( LlmConfig, + MethodConfig, Pt2eQuantize, VgfQuantizeScope, ) @@ -42,6 +49,25 @@ ] +def _tiny_llm_builder(): + """An LLMEdgeManager small enough to lower in a unit test. + + `_export_llama` exports it, so this returns it unexported. + """ + + class Tiny(torch.nn.Module): + def forward(self, tokens): + return tokens.to(torch.float32) * 2.0 + 1.0 + + return LLMEdgeManager( + model=Tiny(), + modelname="tiny", + max_seq_len=4, + use_kv_cache=False, + example_inputs=(torch.ones(1, 4, dtype=torch.long),), + ) + + class ExportLlamaLibTest(unittest.TestCase): def _assert_routes_to(self, lowering, coreml=False, vulkan=False, qnn=False): """Assert which lowering an export routes to, without running one.""" @@ -96,6 +122,148 @@ def test_core_ml_with_vulkan_routes_to_the_combined_lowering(self): self.assertTrue(target.call_args.kwargs["coreml"]) self.assertTrue(target.call_args.kwargs["vulkan"]) + def _run_tiny_export( + self, + generate_etrecord, + directory, + output_name="tiny.pte", + output_dir=None, + xnnpack=True, + ): + """Export the tiny model, writing the .pte into `directory`. + + The builder is pointed at a directory rather than chdir'ing the process, because these + tests run under pytest-xdist where the cwd is shared between tests in a worker. + + `output_dir` defaults to `directory`, and is set separately only by the test that + checks a `.pte` output name does not separate the model from its record. + + `xnnpack=False` selects the combined lowering, which reaches the edge export by a + different route than the backend-specific helpers. + """ + llm_config = LlmConfig() + llm_config.backend.xnnpack.enabled = xnnpack + llm_config.debug.generate_etrecord = generate_etrecord + llm_config.model.enable_dynamic_shape = False + llm_config.export.output_dir = output_dir or directory + llm_config.export.output_name = os.path.join(directory, output_name) + builder = _tiny_llm_builder().set_output_dir(output_dir or directory) + with patch( + "executorch.examples.models.llama.export_llama_lib._prepare_for_llama_export", + return_value=builder, + ): + _export_llama(llm_config) + return sorted(os.listdir(directory)) + + def test_combined_lowering_saves_the_etrecord_beside_the_model(self): + """The combined lowering must write a record too. + + It reaches the edge export by a different route than the backend helpers, so the flag + has to be set on the builder before that export rather than at the lowering. + """ + with tempfile.TemporaryDirectory() as directory: + self.assertEqual( + self._run_tiny_export( + generate_etrecord=True, directory=directory, xnnpack=False + ), + ["etrecord.bin", "tiny.pte"], + ) + self.assertIsNotNone( + parse_etrecord( + os.path.join(directory, "etrecord.bin") + ).edge_dialect_program + ) + + def test_export_saves_the_etrecord_beside_the_model(self): + """The record must be written, and land beside the model. + + A `.pte` output name is used as the path verbatim, so `output_dir` points somewhere + else here: a record keyed off `output_dir` would separate it from the model. + """ + with tempfile.TemporaryDirectory() as directory, tempfile.TemporaryDirectory() as elsewhere: + self.assertEqual( + self._run_tiny_export( + generate_etrecord=True, directory=directory, output_dir=elsewhere + ), + ["etrecord.bin", "tiny.pte"], + ) + # An empty file would satisfy the listing, so load it back. + self.assertIsNotNone( + parse_etrecord( + os.path.join(directory, "etrecord.bin") + ).edge_dialect_program + ) + + def test_export_writes_no_etrecord_when_not_asked(self): + """The common case must stay silent rather than raise or leave a stray file.""" + with tempfile.TemporaryDirectory() as directory: + with self.assertNoLogs(level="WARNING"): + landed = self._run_tiny_export( + generate_etrecord=False, directory=directory + ) + self.assertEqual(landed, ["tiny.pte"]) + + def test_export_keeps_the_model_when_the_etrecord_cannot_be_written(self): + """A debug artifact must not cost the caller the export. + + The record is written after the model, so a failed record write costs the record and + not the .pte. A full disk is the likely trigger, since the record is the larger. + """ + with tempfile.TemporaryDirectory() as directory: + # A directory of that name makes the rename fail without touching permissions. + os.mkdir(os.path.join(directory, "etrecord.bin")) + with self.assertLogs(level="WARNING") as logs: + landed = self._run_tiny_export( + generate_etrecord=True, directory=directory + ) + self.assertIn("tiny.pte", landed) + self.assertTrue(any("Could not write" in line for line in logs.output)) + # The staged record must not be left behind next to the model. + self.assertEqual(landed, ["etrecord.bin", "tiny.pte"]) + + def test_export_keeps_the_previous_etrecord_when_a_rewrite_fails(self): + """A failed rewrite must not destroy the record that was already there. + + The record format truncates its target on open, so writing in place would leave a + short file under the real name and report only a warning. + """ + with tempfile.TemporaryDirectory() as directory: + self._run_tiny_export(generate_etrecord=True, directory=directory) + record = os.path.join(directory, "etrecord.bin") + good = os.path.getsize(record) + + with patch.object( + ETRecord, "_save_graph_map", side_effect=RuntimeError("disk full") + ): + with self.assertLogs(level="WARNING"): + self._run_tiny_export(generate_etrecord=True, directory=directory) + + self.assertEqual(os.path.getsize(record), good) + self.assertIsNotNone(parse_etrecord(record).edge_dialect_program) + + def test_multimethod_export_saves_the_etrecord_beside_the_model(self): + """The multimethod path must write the record too. + + It builds a record the same way the single-method paths do, so leaving out the save + would reproduce the silent-flag bug on that path alone. + """ + with tempfile.TemporaryDirectory() as directory: + llm_config = LlmConfig() + llm_config.backend.xnnpack.enabled = True + llm_config.debug.generate_etrecord = True + llm_config.multimethod.methods = [MethodConfig(method_name="forward")] + llm_config.export.output_dir = directory + llm_config.export.output_name = os.path.join(directory, "tiny.pte") + with patch.object( + export_llama_lib, + "_prepare_for_llama_export", + side_effect=lambda _: _tiny_llm_builder().set_output_dir(directory), + ): + _export_llama(llm_config) + self.assertEqual( + sorted(os.listdir(directory)), ["etrecord.bin", "tiny.pte"] + ) + def test_has_expected_ops_and_op_counts(self): """ Checks the presence of unwanted expensive ops. diff --git a/extension/export_util/utils.py b/extension/export_util/utils.py index 782b2b0ae63..e7534b22448 100644 --- a/extension/export_util/utils.py +++ b/extension/export_util/utils.py @@ -56,6 +56,7 @@ def _core_aten_to_edge( edge_constant_methods: Optional[Dict[str, Any]] = None, edge_compile_config=None, verbose=True, + generate_etrecord: bool = False, ) -> EdgeProgramManager: if not edge_compile_config: edge_compile_config = exir.EdgeCompileConfig( @@ -65,6 +66,7 @@ def _core_aten_to_edge( core_aten_exir_ep, constant_methods=edge_constant_methods, compile_config=edge_compile_config, + generate_etrecord=generate_etrecord, ) if verbose: logging.info(f"Exported graph:\n{edge_manager.exported_program()}") @@ -81,6 +83,7 @@ def export_to_edge( edge_compile_config=_EDGE_COMPILE_CONFIG, strict=True, verbose=True, + generate_etrecord: bool = False, ) -> EdgeProgramManager: core_aten_ep = _to_core_aten( model, @@ -91,7 +94,11 @@ def export_to_edge( verbose=verbose, ) return _core_aten_to_edge( - core_aten_ep, edge_constant_methods, edge_compile_config, verbose=verbose + core_aten_ep, + edge_constant_methods, + edge_compile_config, + verbose=verbose, + generate_etrecord=generate_etrecord, ) diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py index 825768ede3d..6ec9e320eb8 100644 --- a/extension/llm/export/builder.py +++ b/extension/llm/export/builder.py @@ -471,6 +471,7 @@ def export_to_edge(self) -> "LLMEdgeManager": edge_constant_methods=self.metadata, edge_compile_config=edge_config, verbose=self.verbose, + generate_etrecord=self.generate_etrecord, ) return self