diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index c22bdac4ed2..98363769e08 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -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,28 @@ def _to_edge_and_lower_llama_xnnpack( ) +def _save_etrecord_if_generated(builder) -> None: + """Write the etrecord the lowering attached, if there is one. + + Called after the model is written, and never allowed to raise, because a debug artifact + must not cost a caller the export itself. A record twice the size of the model makes a + full disk the likely trigger. + """ + try: + etrecord = builder.export_program.get_etrecord() + except RuntimeError: + # Not generated, which is the normal case. + return + + try: + etrecord.save("etrecord.bin") + except Exception as error: + logging.warning("Could not write etrecord.bin: %s", error) + return + + logging.info("Generated etrecord.bin") + + def _to_edge_and_lower_llama_openvino( builder_exported, modelname, @@ -1546,12 +1567,16 @@ def _to_edge_and_lower_llama( # noqa: C901 # 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") + try: + generate_etrecord_func( + et_record="etrecord.bin", + edge_dialect_program=edge_manager_copy, + executorch_program=builder.export_program, + ) + except Exception as error: + logging.warning("Could not write etrecord.bin: %s", error) + else: + logging.info("Generated etrecord.bin") else: builder = builder_exported_to_edge.to_backend(partitioners) if verbose: @@ -1693,6 +1718,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 @@ -1875,6 +1901,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 40ba701f84e..c7aa4f5b5c0 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -5,9 +5,14 @@ # 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 try: from executorch.backends.arm.quantizer.arm_quantizer import ( @@ -28,6 +33,7 @@ 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, Pt2eQuantize, @@ -40,7 +46,89 @@ ] +def _tiny_llm_builder(): + """An exported LLMEdgeManager small enough to lower in a unit test. + + Enough to exercise the lowering helpers without a checkpoint or a real llama model. + """ + + 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),), + ).export() + + class ExportLlamaLibTest(unittest.TestCase): + def _run_tiny_export(self, generate_etrecord): + llm_config = LlmConfig() + llm_config.backend.xnnpack.enabled = True + llm_config.debug.generate_etrecord = generate_etrecord + llm_config.export.output_name = "tiny.pte" + with patch( + "executorch.examples.models.llama.export_llama_lib._prepare_for_llama_export", + return_value=_tiny_llm_builder(), + ): + _export_llama(llm_config) + + def test_export_saves_the_etrecord_it_generates(self): + """A lowering that generates an etrecord must also write it. + + `to_edge_transform_and_lower` builds the record when asked and carries it to the ExecuTorch + program, but nothing saved it, so the flag produced no file and no warning while still + paying for the record. + """ + with tempfile.TemporaryDirectory() as directory: + previous = os.getcwd() + os.chdir(directory) + try: + self._run_tiny_export(generate_etrecord=True) + self.assertEqual(sorted(os.listdir(".")), ["etrecord.bin", "tiny.pte"]) + # An empty file would satisfy the listing, so load it back. + self.assertIsNotNone( + parse_etrecord("etrecord.bin").edge_dialect_program + ) + finally: + os.chdir(previous) + + 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: + previous = os.getcwd() + os.chdir(directory) + try: + self._run_tiny_export(generate_etrecord=False) + self.assertEqual(os.listdir("."), ["tiny.pte"]) + finally: + os.chdir(previous) + + 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 failure to write it costs the record and not + the .pte. A full disk is the likely trigger, since the record is the larger of the two. + """ + with tempfile.TemporaryDirectory() as directory: + previous = os.getcwd() + os.chdir(directory) + try: + # A directory of that name makes the write fail without touching permissions. + os.mkdir("etrecord.bin") + with self.assertLogs(level="WARNING") as logs: + self._run_tiny_export(generate_etrecord=True) + self.assertIn("tiny.pte", os.listdir(".")) + self.assertTrue( + any("Could not write etrecord.bin" in line for line in logs.output) + ) + finally: + os.chdir(previous) + def test_has_expected_ops_and_op_counts(self): """ Checks the presence of unwanted expensive ops.